Currently only tool calling supports MCP auth. This refactors the auth logic into a auth_utils file and uses it for tool listing as well. Fixes https://github.com/google/adk-python/issues/2168.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 856798589
Merge https://github.com/google/adk-python/pull/4025
**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.**
### Link to Issue or Description of Change
**1. Link to an existing issue (if applicable):**
- Closes:
- #3950
- #3731
- #3708
**2. Or, if no issue exists, describe the change:**
**Problem:**
- `ClientSession` of https://github.com/modelcontextprotocol/python-sdk uses AnyIO for async task management.
- AnyIO TaskGroup requires its start and close must happen in a same task.
- Since `McpSessionManager` does not create task per client, the client might be closed by different task, cause the error: `Attempted to exit cancel scope in a different task than it was entered in`.
**Solution:**
I Suggest 2 changes:
Handling the `ClientSession` in a single task
- To start and close `ClientSession` by the same task, we need to wrap the whole lifecycle of `ClientSession` to a single task.
- `SessionContext` wraps the initialization and disposal of `ClientSession` to a single task, ensures that the `ClientSession` will be handled only in a dedicated task.
Add timeout for `ClientSession`
- Since now we are using task per `ClientSession`, task should never be leaked.
- But `McpSessionManager` does not deliver timeout directly to `ClientSession` when the type is not STDIO.
- There is only timeout for `httpx` client when MCP type is SSE or StreamableHTTP.
- But the timeout applys only to `httpx` client, so if there is an issue in MCP client itself(e.g. https://github.com/modelcontextprotocol/python-sdk/issues/262), a tool call waits the result **FOREVER**!
- To overcome this issue, I propagated the `sse_read_timeout` to `ClientSession`.
- `timeout` is too short for timeout for tool call, since its default value is only 5s.
- `sse_read_timeout` is originally made for read timeout of SSE(default value of 5m or 300s), but actually most of SSE implementations from server (e.g. FastAPI, etc.) sends ping periodically(about 15s I assume), so in a normal circumstances this timeout is quite useless.
- If the server does not send ping, the timeout is equal to tool call timeout. Therefore, it would be appropriate to use `sse_read_timeout` as tool call timeout.
- Most of tool calls should finish within 5 minutes, and sse timeout is adjustable if not.
- If this change is not acceptable, we could make a dedicate parameter for tool call timeout(e.g. `tool_call_timeout`).
### Testing Plan
- Although this does not change the interface itself, it changes its own session management logics, some existing tests are no longer valid.
- I made changes to those tests, especially those of which validate session states(e.g. checking whether `initialize()` called).
- Since now session is encapsulated with `SessionContext`, we cannot validate the initialized state of the session in `TestMcpSessionManager`, should validate it at `TestSessionContext`.
- Added a simple test for reproducing the issue(`test_create_and_close_session_in_different_tasks`).
- Also made a test for the new component: `SessionContext`.
**Unit Tests:**
- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.
```plaintext
=================================================================================== 3689 passed, 1 skipped, 2205 warnings in 63.39s (0:01:03) ===================================================================================
```
**Manual End-to-End (E2E) Tests:**
_Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix._
### Checklist
- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective or that my feature works.
- [x] New and existing unit tests pass locally with my changes.
- [x] I have manually tested my changes end-to-end.
- [ ] ~~Any dependent changes have been merged and published in downstream modules.~~ `no deps has been changed`
### Additional context
This PR is related to https://github.com/modelcontextprotocol/python-sdk/pull/1817 since it also fixes endless tool call awaiting.
Co-authored-by: Kathy Wu <wukathy@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4025 from challenger71498:feat/task-based-mcp-session-manager f7f7cd0c9c96840361c30499d08c33a189f57d86
PiperOrigin-RevId: 856438147
There was an extra test_mcp_toolset in the tools/ directory with only one test; I moved it into the main file.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 856419611
gcloud auth team requested that we audit ADK's codebase for places where ADC (google.auth.default) is used, and make sure that the quota project id header is being populated.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 855322964
API registry uses a nextPageToken for pagination, so we should loop through all the pages by sending the received nextPageToken to subsequent requests until no token is returned. The first page contains 1P MCP servers, and later pages contain 3P customer servers.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 853880449
Fix discovery engine search tool, bigquery agent analytics plugin, and application integration tool to correctly handle the ADC quota project override -- the x-goog-user-project should be set based on the ADC quota project, per gcloud auth team's requirements.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 853841124
There isn't a consistent format for MCP server urls in the registry-- some customers add the https:// but others don't. To standardize, only prepend if it isn't there already.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 846451152
This gcloud authentication header is required as a project override when using google ADC credentials. It is required for all OneMCP servers. Originally I manually put it in the big query sample, but since it is used for all OneMCP, it makes sense to just add it to the core API Registry logic. I also refactored the auth header logic a bit to deduplicate.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 845967178
This calls the cloudapiregistry.googleapis.com API to get MCP tools from the project's registry, and adds them to ADK.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 837166909
Retrying only on closed_resource error is not enough to be reliable for production environments due to the other network errors that may occur -- remote protocol error, read timeout, etc. We will update this to retry on all errors. Since it is only a one-time retry, it should not affect latency significantly. Fixes https://github.com/google/adk-python/issues/2561.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 834514264
Currently logic to check for a disconnected session only checks for certain headers but doesn't detect all cases, leading to situations where it tries to connect to a session that is down. This adds logic so that we ping the server to check if it is disconnected. Fixes https://github.com/google/adk-python/issues/3321.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 833487825
Currently logic to check for a disconnected session only checks for certain headers but doesn't detect all cases, leading to situations where it tries to connect to a session that is down. This adds logic so that we ping the server to check if it is disconnected. Fixes https://github.com/google/adk-python/issues/3321.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 832460068
Retrying only on closed_resource error is not enough to be reliable for production environments due to the other network errors that may occur -- remote protocol error, read timeout, etc. We will update this to retry on all errors. Since it is only a one-time retry, it should not affect latency significantly. Fixes https://github.com/google/adk-python/issues/2561.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 831123085
Currently ADK web hangs and logs "AGSI callable returned without completing response" when the server is unreachable. To fix, set timeouts for connecting to server.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 828648928
Add support for MCP prompts via the McpInstructionProvider class, which can be specified as an agent's instruction.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 828166051
Users were getting spammed with this log even though their tools didn't require authentication. To fix, reduce the log level to DEBUG so that it doesn't show up by default.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 828161281
Populate the usage_metadata field for live events with the metadata provided by the Gemini live API.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 828124232
This fixes MCP authentication for gcloud service accounts. Previously it was failing to authenticate tool calls.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 826639044
While testing the bidi streaming sample agent, I noticed that the import was erroring and I think it's a typo -- the other bidi sample agents all import `from google.adk.agents.llm_agent`.
PiperOrigin-RevId: 823662586
Since it appears in the same bubble as the rest of the LLM's response text, make it more human readable so it doesn't look out of place.
PiperOrigin-RevId: 822729061
Previously BuiltInCodeExecutor was missing the logic to save output files from executed code as artifacts, so images/visualizations wouldn't show up in the UI. This fix will iterate through all parts of the LlmResponse, and if any of them are images, it will save the image data using artifact_service (similar to what is done in VertexAICodeExecutor).
This fixes the backend, but there are still UI bugs that should be fixed -- events without content are currently ignored, so the image doesn't appear even though it is saved. We will add the UI fix in a separate change.
PiperOrigin-RevId: 822245140
Add a header_provider param which is a callable[ReadonlyContext, Dict[str, Any]] for users to build headers in MCPToolset
fix: https://github.com/google/adk-python/issues/3156
PiperOrigin-RevId: 820412372
