Skip to content
27 changes: 19 additions & 8 deletions src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -369,14 +369,6 @@ def get_author_for_event(llm_response):
while True:
async with Aclosing(llm_connection.receive()) as agen:
async for llm_response in agen:
if llm_response.live_session_resumption_update:
logger.info(
'Update session resumption handle:'
f' {llm_response.live_session_resumption_update}.'
)
invocation_context.live_session_resumption_handle = (
llm_response.live_session_resumption_update.new_handle
)
model_response_event = Event(
id=Event.new_id(),
invocation_id=invocation_context.invocation_id,
Expand Down Expand Up @@ -739,6 +731,25 @@ async def _postprocess_live(
async for event in agen:
yield event

# Handle session resumption updates for cross-connection resumption.
# Must be before skip condition - resumption updates have no content.
if llm_response.live_session_resumption_update:
# Update internal handle for auto-resumption within run_live()
logger.info(
'Update session resumption handle: %s',
llm_response.live_session_resumption_update,
)
invocation_context.live_session_resumption_handle = (
llm_response.live_session_resumption_update.new_handle
)

# Expose update in event for application-level cross-connection resumption
model_response_event.live_session_resumption_update = (
llm_response.live_session_resumption_update
)
yield model_response_event
return
Comment thread
AshleyMaloney marked this conversation as resolved.
Comment thread
AshleyMaloney marked this conversation as resolved.

# Skip the model response event if there is no content and no error code.
# This is needed for the code executor to trigger another loop.
# But don't skip control events like turn_complete or transcription events.
Expand Down
Loading