Skip to content

Commit 51e66e4

Browse files
fix(sessions): fix httpx client closure during event pagination
Move event iteration inside the api_client context manager in VertexAiSessionService.get_session() to prevent client closure during multi-page event fetching.
1 parent 8da61be commit 51e66e4

File tree

1 file changed

+17
-18
lines changed

1 file changed

+17
-18
lines changed

src/google/adk/sessions/vertex_ai_session_service.py

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -162,26 +162,25 @@ async def get_session(
162162
**list_events_kwargs,
163163
),
164164
)
165+
if get_session_response.user_id != user_id:
166+
raise ValueError(
167+
f'Session {session_id} does not belong to user {user_id}.'
168+
)
165169

166-
if get_session_response.user_id != user_id:
167-
raise ValueError(
168-
f'Session {session_id} does not belong to user {user_id}.'
170+
update_timestamp = get_session_response.update_time.timestamp()
171+
session = Session(
172+
app_name=app_name,
173+
user_id=user_id,
174+
id=session_id,
175+
state=getattr(get_session_response, 'session_state', None) or {},
176+
last_update_time=update_timestamp,
169177
)
170-
171-
update_timestamp = get_session_response.update_time.timestamp()
172-
session = Session(
173-
app_name=app_name,
174-
user_id=user_id,
175-
id=session_id,
176-
state=getattr(get_session_response, 'session_state', None) or {},
177-
last_update_time=update_timestamp,
178-
)
179-
# Preserve the entire event stream that Vertex returns rather than trying
180-
# to discard events written milliseconds after the session resource was
181-
# updated. Clock skew between those writes can otherwise drop tool_result
182-
# events and permanently break the replayed conversation.
183-
async for event in events_iterator:
184-
session.events.append(_from_api_event(event))
178+
# Preserve the entire event stream that Vertex returns rather than trying
179+
# to discard events written milliseconds after the session resource was
180+
# updated. Clock skew between those writes can otherwise drop tool_result
181+
# events and permanently break the replayed conversation.
182+
async for event in events_iterator:
183+
session.events.append(_from_api_event(event))
185184

186185
if config:
187186
# Filter events based on num_recent_events.

0 commit comments

Comments
 (0)