Commit Graph

1956 Commits

Author SHA1 Message Date
Kathy Wu cce430da79 feat: start and close ClientSession in a single task in McpSessionManager
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
2026-01-14 18:10:03 -08:00
Google Team Member 1133ce219c feat: convert A2UI messages between A2A DataPart metadata and ADK events
1. Convert A2A responses containing a DataPart to ADK events. By default, this is done by serializing the DataPart to JSON and embedding it within the inline_data field of a GenAI Part, wrapped with custom tags (<a2a_datapart_json> and </a2a_datapart_json>).
2. Convert ADK events back to A2A requests. Specifically, messages stored in inline_data with the text/plain mime type and content wrapped within the custom tags (<a2a_datapart_json> and </a2a_datapart_json>) are deserialized from JSON back into an A2A DataPart

PiperOrigin-RevId: 856426615
2026-01-14 17:07:54 -08:00
Xiang (Sean) Zhou 712b5a393d fix: Only filter out audio content when sending history
audio is transcribed thus no need to be sent, but other blob(e.g. image) should still be sent.

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 856422986
2026-01-14 16:51:19 -08:00
George Weale 89bed43f5e fix: Add finish reason mapping and remove custom file URI handling in LiteLLM
Introduces a function to map LiteLLM finish reason strings to the internal types.FinishReason enum and populates the finish_reason field in LlmResponse. Removes custom logic for handling file URIs, including special casing for different providers, and updates tests accordingly

Close #4125

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 856421317
2026-01-14 16:44:44 -08:00
Kathy Wu 8264211f98 chore: Consolidate test_mcp_toolset.py into one file
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
2026-01-14 16:37:12 -08:00
Xuan Yang 8e7cc16f12 docs: Refactor ADK release analyzer with workflow agents
Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 856412858
2026-01-14 16:19:02 -08:00
George Weale fdc98d5c92 fix: Convert unsupported inline artifact MIME types to text in LoadArtifactsTool
The LoadArtifactsTool now checks if an artifact's inline data MIME type is supported by Gemini. If not, it attempts to convert the artifact content into a text Part

Close #4028

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 856404510
2026-01-14 16:01:06 -08:00
Xiang (Sean) Zhou 7b035aa9fc chore: Always log api backend when connecting to live model
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 856404282
2026-01-14 16:00:15 -08:00
Google Team Member 672b57f1b7 chore: add a sample BigQuery agent using BigQuery MCP tools
PiperOrigin-RevId: 856400285
2026-01-14 15:49:23 -08:00
Kotaro Saito 38d52b2476 fix(cli): pass log_level to uvicorn in web and api_server commands (#4144)
Explicitly pass the log_level parameter to uvicorn.Config in both
adk web and adk api_server commands. This ensures that Uvicorn's
internal logging respects the configured log level.

Closes #4139
2026-01-14 09:48:32 -08:00
Xuan Yang 79fcddb39f feat: Add --enable_features CLI option to ADK CLI
This flag can be used to override default feature enable state.

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 856067979
2026-01-13 23:58:26 -08:00
Xuan Yang 8973618b0b chore: Add a DebugLoggingPlugin to record human readable debugging logs
Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 856067925
2026-01-13 23:57:36 -08:00
lwangverizon 8e69a58df4 feat: Add support to automatically create a session if one does not exist
feature/auto-create-new-session

Merge https://github.com/google/adk-python/pull/4072

**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

**2. Or, if no issue exists, describe the change:**

**Problem:**

When building frontend applications with ADK, there's a limitation where frontends cannot always guarantee that `create_session` is called before initiating a conversation. This creates friction in the user experience because:

- Users may refresh the page or navigate directly to a conversation URL with a specific session_id
- Frontend state management may lose track of whether a session was already created
- Mobile apps or single-page applications have complex lifecycle management where ensuring `create_session` is called first adds unnecessary complexity
- This forces developers to implement additional logic to check session existence before every conversation

Currently, if `get_session` is called with a non-existent session_id, it returns `None`, requiring the frontend to explicitly handle this case and call `create_session` separately.

**Solution:**

Modified the `get_session` method in `DatabaseSessionService` to automatically create a session if it doesn't exist in the database. This "get or create" pattern is common in many frameworks and provides a more developer-friendly API.

The implementation:
1. Attempts to fetch the session from the database
2. If the session doesn't exist (returns `None`), automatically calls `create_session` with the provided parameters
3. Retrieves and returns the newly created session
4. Maintains backward compatibility - existing code continues to work without changes

This allows frontends to simply call `get_session` with a session_id and be confident that the session will be available, regardless of whether it was previously created.

**Benefits:**
- Simplifies frontend integration by removing the need to track session creation state
- Reduces API calls (no need to check existence before calling get_session)
- Follows the principle of least surprise - getting a session with an ID should work reliably
- No breaking changes to existing code that checks for `None` return values

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

**pytest results:**

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4072 from lwangverizon:feature/auto-create-new-session 5475c6ae91d12332598d521302736eb1db79a8be
PiperOrigin-RevId: 856019482
2026-01-13 21:23:49 -08:00
Xiang (Sean) Zhou 1bedffe457 chore: Remove dead codes for flushing model audio when generation completes
LlmResponse so far doesn't expose generation_complete signal, removing the dead codes.

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 855835802
2026-01-13 12:28:24 -08:00
Google Team Member 277084e313 refactor(tools): Update ToolboxToolset to wrap toolbox-adk
### Description of Change

**Problem:**
The `ToolboxToolset` was implemented directly within `adk-python`, leading to code duplication and potential drift from the core `toolbox-adk` implementation.

**Solution:**
Refactored `ToolboxToolset` to act as a rigorous wrapper around the `toolbox-adk` package, which delegates all functionality to `toolbox_adk.ToolboxToolset`.

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

Summary:
-   Verified initialization flows through to `toolbox-adk`.
-   Verified `auth_token_getters` are correctly propagated.
-   Verified type hints are static-analysis friendly.

**Manual End-to-End (E2E) Tests:**

Manually verified standard toolbox loading and execution with the new wrapper:
```python
from google.adk.tools import ToolboxToolset
from toolbox_adk import CredentialStrategy

# Loading with toolset_name
ts = ToolboxToolset(
    server_url='http://localhost:8080',
    toolset_name='calculator',
    credentials=CredentialStrategy.toolbox_identity()
)
tools = await ts.get_tools()
print(f'Loaded {len(tools)} tools')
```

### 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.
- [x] Any dependent changes have been merged and published in downstream modules.

PiperOrigin-RevId: 855798474
2026-01-13 10:59:51 -08:00
Xiang (Sean) Zhou ab62b1bffd fix: Use the agent name as the author of the audio event
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 855789317
2026-01-13 10:38:53 -08:00
Xiang (Sean) Zhou f668a5de44 chore: Update comments about why we can return upon flushing audio caches
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 855788674
2026-01-13 10:37:28 -08:00
Xuan Yang a5f0d333d7 feat: Use json schema for RestApiTool declaration when feature enabled
Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 855767527
2026-01-13 09:50:59 -08:00
Google Team Member fd2c0f556b chore: Upgrade the sample BQ agent model version to gemini-2.5-flash
This change should give a more intelligent out-of-the-box experience to the sample BigQuery agent users.

PiperOrigin-RevId: 855552871
2026-01-12 23:30:16 -08:00
Joseph Pagadora 6d2f33a59c feat: Update EvalConfig and EvalMetric data models to support custom metrics
Co-authored-by: Joseph Pagadora <jcpagadora@google.com>
PiperOrigin-RevId: 855517478
2026-01-12 21:37:21 -08:00
George Weale 905604faac refactor: Import migration_runner lazily within the migrate command
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 855496468
2026-01-12 20:21:57 -08:00
Xiang (Sean) Zhou a4732a8e50 chore: Updatedate the doc string of LiveRequest
activity_start, activity_end , and blob is sent via send_realtime_input, according to the docstring : https://github.com/googleapis/python-genai/blob/d98c757c7d9a88026ac0f9eb2b1b578e2e7f3bfe/google/genai/live.py#L251, they should be send separately, so makes sense they are mutally exclusive.

And we already make them exclusive in our current logic:
https://github.com/google/adk-python/blob/6c0bf85042c38c7bafe1c183f1bba8bee1ba3570/src/google/adk/flows/llm_flows/base_llm_flow.py#L264-L272, make it clear in docstring

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 855492864
2026-01-12 20:06:08 -08:00
Liang Wu 75231a30f1 fix: Handle NOT_FOUND error when fetching Vertex AI sessions
Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 855464349
2026-01-12 18:17:41 -08:00
Alexis Marasigan b725045e5a fix: fix httpx client closure during event pagination
Merge https://github.com/google/adk-python/pull/3756

move event iteration inside api_client context in get_session

Move event iteration inside the api_client context manager in VertexAiSessionService.get_session() to prevent client closure during multi-page event fetching.

**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: #3757

**2. Or, if no issue exists, describe the change:**

**Problem:**

When a session contains more than 100 events (requiring pagination), `VertexAiSessionService.get_session()` fails with:

```
RuntimeError: Cannot send a request, as the client has been closed.
```

The root cause is that the `events_iterator` is consumed **outside** the `async with self._get_api_client() as api_client:` context block. When the iterator needs to fetch page 2, 3, etc., the API client has already been closed because the `async with` block has exited.

```python
# Current buggy flow:
async with self._get_api_client() as api_client:
    get_session_response, events_iterator = await asyncio.gather(...)
# ← Client closed here

async for event in events_iterator:  # ← Fails on page 2+ (client closed)
    session.events.append(...)
```

**Solution:**

Move the session creation, user validation, and event iteration **inside** the `async with` block so the API client remains open during the entire pagination process:

```python
async with self._get_api_client() as api_client:
    get_session_response, events_iterator = await asyncio.gather(...)
    # Validation and session creation...
    async for event in events_iterator:  # ← Now works for all pages
        session.events.append(...)
# Client closed after all events are fetched
```

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

```bash
pytest tests/unittests/sessions/test_vertex_ai_session_service.py -v
```

**Added regression test:** `test_get_session_pagination_keeps_client_open`
- Creates a `MockAsyncClientWithPagination` that tracks whether it's inside the `async with` context
- Raises `RuntimeError` if iteration happens outside the context (matching real httpx behavior)
- Simulates 3 pages of events (100 + 100 + 50 = 250 events)
- Verifies all 250 events are successfully retrieved

**Manual End-to-End (E2E) Tests:**

1. Deploy an ADK agent to Vertex AI Agent Engine
2. Create a session and send 100+ messages to accumulate >100 events
3. Verify `get_session()` successfully retrieves all events without error

**Before fix:**
```
RuntimeError: Cannot send a request, as the client has been closed.
```

**After fix:**
- Session with 201 events (3 pages) loads successfully
- All events are retrieved and appended to the session

### 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.
- [x] Any dependent changes have been merged and published in downstream modules.

### Additional context

This bug affects any production deployment where users have extended conversations. Sessions accumulating >100 events (which triggers pagination) become completely unusable as the agent cannot load the session to process new messages.

The fix is minimal and maintains backward compatibility - it only changes the scope of the `async with` block without altering any logic or return values.

**Affected versions:** Tested on google-adk 1.19.0, but the bug exists in earlier versions as well.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3756 from AlexisMarasigan:fix/vertex-ai-session-service-paginatio 01fbafa6524312f24f7c9feaffb07bff0ad49b77
PiperOrigin-RevId: 855451813
2026-01-12 17:27:29 -08:00
Google Team Member a4116a6cbf feat: Enhance TraceManager async safety, enrich BigQuery plugin logging, and fix serialization
*   **Async Safety:** Improved TraceManager context variable handling to ensure correct context isolation in concurrent asynchronous operations. This was achieved by using immutable tuples for the span stack and making copies of context dictionaries before modification.
*   **Enhanced Logging:** The BigQueryAgentAnalyticsPlugin now captures richer metadata, including:
    *   Root agent name (via a new context variable).
    *   LLM model name and version.
    *   Usage metadata from LLM requests and responses.
*   **Serialization Fix:** Updated BigQueryAgentAnalyticsPlugin to prevent JSON serialization errors when logging custom objects (e.g., Dataclasses). These are now automatically converted to dictionaries or string representations to ensure successful insertion into BigQuery.

PiperOrigin-RevId: 855415320
2026-01-12 15:46:12 -08:00