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
This commit is contained in:
Alexis Marasigan
2026-01-12 17:27:29 -08:00
committed by Copybara-Service
parent a4116a6cbf
commit b725045e5a
2 changed files with 114 additions and 18 deletions
@@ -162,26 +162,25 @@ class VertexAiSessionService(BaseSessionService):
**list_events_kwargs,
),
)
if get_session_response.user_id != user_id:
raise ValueError(
f'Session {session_id} does not belong to user {user_id}.'
)
if get_session_response.user_id != user_id:
raise ValueError(
f'Session {session_id} does not belong to user {user_id}.'
update_timestamp = get_session_response.update_time.timestamp()
session = Session(
app_name=app_name,
user_id=user_id,
id=session_id,
state=getattr(get_session_response, 'session_state', None) or {},
last_update_time=update_timestamp,
)
update_timestamp = get_session_response.update_time.timestamp()
session = Session(
app_name=app_name,
user_id=user_id,
id=session_id,
state=getattr(get_session_response, 'session_state', None) or {},
last_update_time=update_timestamp,
)
# Preserve the entire event stream that Vertex returns rather than trying
# to discard events written milliseconds after the session resource was
# updated. Clock skew between those writes can otherwise drop tool_result
# events and permanently break the replayed conversation.
async for event in events_iterator:
session.events.append(_from_api_event(event))
# Preserve the entire event stream that Vertex returns rather than trying
# to discard events written milliseconds after the session resource was
# updated. Clock skew between those writes can otherwise drop tool_result
# events and permanently break the replayed conversation.
async for event in events_iterator:
session.events.append(_from_api_event(event))
if config:
# Filter events based on num_recent_events.
@@ -397,6 +397,103 @@ class MockAsyncClient:
self.event_dict[session_id] = ([event_json], None)
class MockAsyncClientWithPagination:
"""Mock client that simulates pagination requiring an open client connection.
This mock tracks whether the client context is active and raises RuntimeError
if iteration occurs outside the context, simulating the real httpx behavior.
"""
def __init__(self, session_data: dict, events_pages: list[list[dict]]):
self._session_data = session_data
self._events_pages = events_pages
self._context_active = False
self.agent_engines = mock.AsyncMock()
self.agent_engines.sessions.get.side_effect = self._get_session
self.agent_engines.sessions.events.list.side_effect = self._list_events
async def __aenter__(self):
self._context_active = True
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self._context_active = False
async def _get_session(self, name: str):
return _convert_to_object(self._session_data)
async def _list_events(self, name: str, **kwargs):
return self._paginated_events_iterator()
async def _paginated_events_iterator(self):
for page in self._events_pages:
for event in page:
if not self._context_active:
raise RuntimeError(
'Cannot send a request, as the client has been closed.'
)
yield _convert_to_object(event)
def _generate_events_for_page(session_id: str, start_idx: int, count: int):
events = []
start_time = isoparse('2024-12-12T12:12:12.123456Z')
for i in range(count):
idx = start_idx + i
event_time = start_time + datetime.timedelta(microseconds=idx * 1000)
events.append({
'name': (
'projects/test-project/locations/test-location/'
f'reasoningEngines/123/sessions/{session_id}/events/{idx}'
),
'invocation_id': f'invocation_{idx}',
'author': 'pagination_user',
'timestamp': event_time.isoformat().replace('+00:00', 'Z'),
})
return events
@pytest.mark.asyncio
async def test_get_session_pagination_keeps_client_open():
"""Regression test: event iteration must occur inside the api_client context.
This test verifies that get_session() keeps the API client open while
iterating through paginated events. Before the fix, the events_iterator
was consumed outside the async with block, causing RuntimeError when
fetching subsequent pages.
"""
session_data = {
'name': (
'projects/test-project/locations/test-location/'
'reasoningEngines/123/sessions/pagination_test'
),
'update_time': '2024-12-12T12:12:12.123456Z',
'user_id': 'pagination_user',
}
page1_events = _generate_events_for_page('pagination_test', 0, 100)
page2_events = _generate_events_for_page('pagination_test', 100, 100)
page3_events = _generate_events_for_page('pagination_test', 200, 50)
mock_client = MockAsyncClientWithPagination(
session_data=session_data,
events_pages=[page1_events, page2_events, page3_events],
)
session_service = mock_vertex_ai_session_service()
with mock.patch.object(
session_service, '_get_api_client', return_value=mock_client
):
session = await session_service.get_session(
app_name='123', user_id='pagination_user', session_id='pagination_test'
)
assert session is not None
assert len(session.events) == 250
assert session.events[0].invocation_id == 'invocation_0'
assert session.events[249].invocation_id == 'invocation_249'
def mock_vertex_ai_session_service(
project: Optional[str] = 'test-project',
location: Optional[str] = 'test-location',