Fix bidi session restoration 3573 #4194
Conversation
…3940) Repeated tool execution can happen when the model emits the same function call multiple times (same tool name + same arguments), which is especially painful for slow or expensive tools. This change introduces an invocation-scoped single-flight cache: - Adds RunConfig.dedupe_tool_calls (opt-in) to reuse results for identical calls - Always dedupes BaseTool.is_long_running tools (safer default for long ops) - Ensures only the first execution applies tool side effects (state/artifacts) - Emits cached responses with per-call function_response.id and debug metadata Tests: - Added unit tests covering duplicate calls across steps and within one step.
…#3573) When a client reconnects to a live (BIDI) session using SQLite or database persistence, user messages were stored correctly in the database but not replayed to the reconnecting client. This caused the UI to show only agent responses after reconnection, making conversations appear incomplete. Root Cause: The /run_live websocket endpoint was loading the session (including all events) but only forwarding new events generated by the runner. There was no logic to replay existing events back to the reconnecting client. Changes: - Modified adk_web_server.py run_agent_live() to replay all existing session events to the client before starting the live runner - Added comprehensive error handling to continue replay even if individual events fail to serialize - Added logging to track event replay for debugging Testing: - Added test_live_session_restoration.py with 3 unit tests covering: * Successful replay of all user and agent events on reconnection * Graceful handling of sessions with no events * Continuation of replay even if one event fails Fixes google#3573
Summary of ChangesHello @jayy-77, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces significant improvements to the system's robustness and efficiency. It resolves a critical bug in bidirectional session management by ensuring that clients reconnecting to a live session are provided with a complete and accurate history of all prior events, including user inputs. Concurrently, it implements a new optimization for tool execution, preventing redundant calls to the same tool with identical arguments, which can lead to performance gains and more consistent behavior, particularly in scenarios involving repeated or slow tool operations. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Response from ADK Triaging Agent Hello @jayy-77, thank you for creating this PR! To help us review your contribution more effectively, could you please fill out the pull request description template? Specifically, we're looking for:
This information is really helpful for our reviewers to understand the context and validate the fix. You can find more details in our contribution guide. Thanks! |
There was a problem hiding this comment.
Code Review
This pull request introduces two main features: fixing bidirectional session restoration by replaying events upon reconnection, and adding a tool call deduplication mechanism to prevent redundant executions. The implementation for both features appears solid. The session restoration now correctly replays the entire conversation history, and the tool call caching is robustly implemented with proper asynchronous locking and error handling.
My review includes two main points. First, I've identified a significant feature inconsistency between the live and async tool execution paths regarding plugin support, which this PR's refactoring makes more apparent. Second, the new tests for session restoration are tightly coupled to the implementation, which is a fragile pattern. I've suggested a refactoring to improve the testability and maintainability of this part of the code. The new tests for tool call deduplication, on the other hand, are well-written and effectively test the feature's behavior.
| @pytest.mark.asyncio | ||
| async def test_live_session_replays_all_events_on_reconnection(): | ||
| """Test that reconnecting to a live session replays all events including user messages. | ||
| This tests the fix for Issue #3573 where user messages were stored in the | ||
| database but not sent back to the client on reconnection. | ||
| """ | ||
| # Create a mock session with both user and agent events | ||
| user_event = Event( | ||
| id="event-user-1", | ||
| author="user", | ||
| content=types.Content(parts=[types.Part(text="Hello, assistant!")]), | ||
| invocation_id="inv-1", | ||
| ) | ||
| agent_event = Event( | ||
| id="event-agent-1", | ||
| author="test_agent", | ||
| content=types.Content( | ||
| parts=[types.Part(text="Hello! How can I help you?")] | ||
| ), | ||
| invocation_id="inv-1", | ||
| ) | ||
| user_event2 = Event( | ||
| id="event-user-2", | ||
| author="user", | ||
| content=types.Content( | ||
| parts=[types.Part(text="What's the weather today?")] | ||
| ), | ||
| invocation_id="inv-2", | ||
| ) | ||
| agent_event2 = Event( | ||
| id="event-agent-2", | ||
| author="test_agent", | ||
| content=types.Content( | ||
| parts=[types.Part(text="I can help you check the weather.")] | ||
| ), | ||
| invocation_id="inv-2", | ||
| ) | ||
|
|
||
| mock_session = Session( | ||
| app_name="test_app", | ||
| user_id="test_user", | ||
| id="test_session", | ||
| state={}, | ||
| events=[user_event, agent_event, user_event2, agent_event2], | ||
| last_update_time=1234567890.0, | ||
| ) | ||
|
|
||
| # Mock WebSocket to capture replayed events | ||
| mock_websocket = AsyncMock() | ||
| replayed_events = [] | ||
|
|
||
| async def capture_send_text(data): | ||
| replayed_events.append(data) | ||
|
|
||
| mock_websocket.send_text = capture_send_text | ||
|
|
||
| # Test the core event replay logic that should be in run_agent_live | ||
| # This simulates what happens when a client reconnects: | ||
| # 1. Session is loaded (with all events) | ||
| session = mock_session | ||
|
|
||
| # 2. All existing events should be replayed to the client | ||
| if session and session.events: | ||
| for event in session.events: | ||
| await mock_websocket.send_text( | ||
| event.model_dump_json(exclude_none=True, by_alias=True) | ||
| ) | ||
|
|
||
| # Verify that all 4 events were replayed (2 user + 2 agent) | ||
| assert len(replayed_events) == 4 | ||
|
|
||
| # Verify that events were sent in order | ||
| import json | ||
|
|
||
| event_data = [json.loads(data) for data in replayed_events] | ||
|
|
||
| # Check that user messages are included | ||
| assert event_data[0]["author"] == "user" | ||
| assert "Hello, assistant!" in event_data[0]["content"]["parts"][0]["text"] | ||
|
|
||
| assert event_data[1]["author"] == "test_agent" | ||
|
|
||
| assert event_data[2]["author"] == "user" | ||
| assert "weather" in event_data[2]["content"]["parts"][0]["text"] | ||
|
|
||
| assert event_data[3]["author"] == "test_agent" |
There was a problem hiding this comment.
The tests in this file are reimplementing the event replay logic from adk_web_server.py instead of testing the actual implementation. This is a fragile testing pattern because if the production code changes, these tests might not fail, giving a false sense of security. For example, test_live_session_replays_all_events_on_reconnection doesn't include the try-except block that the production code has.
To make these tests more robust and maintainable, I recommend extracting the event replay logic into a separate, testable helper function within adk_web_server.py. The tests can then directly call and verify this helper function.
Here is a suggested refactoring for src/google/adk/cli/adk_web_server.py:
async def _replay_session_events(
session: Session, websocket: WebSocket, logger: logging.Logger
):
"""Replays all existing events in a session to a websocket client."""
if not session.events:
return
logger.info(
"Replaying %d existing events for session %s",
len(session.events),
session.id,
)
for event in session.events:
try:
await websocket.send_text(
event.model_dump_json(exclude_none=True, by_alias=True)
)
except Exception as e:
logger.error(
"Failed to replay event %s during session restoration: %s",
event.id,
e,
)
# Continue replaying other events even if one fails
continue
@app.websocket("/run_live")
async def run_agent_live(
# ... function signature ...
) -> None:
# ... existing setup code ...
if not session:
await websocket.close(code=1002, reason="Session not found")
return
# Replay existing session events to the reconnecting client
await _replay_session_events(session, websocket, logger)
live_request_queue = LiveRequestQueue()
# ... rest of the function ...With this change, your tests can import and test _replay_session_events directly, which would be a much more robust approach.
418070b to
b86c3d5
Compare
) This fix addresses duplicate agent responses that occur during agent transfers and session resumption when using Gemini Live API. Problem: When the Gemini Live API performs transparent session resumption (replaying the model's last response automatically), our manual event replay logic in the /run_live endpoint conflicts with this, causing the same response to be sent twice to the client. Root Cause: Issue google#3573's fix added event replay for BIDI session reconnections, which works perfectly for text-only sessions. However, for live/audio sessions, the Gemini Live API has built-in session resumption that already replays the model's response. When both mechanisms are active, duplicates occur. Solution: Added intelligent session detection logic that distinguishes between: - Text-only sessions: Replay all events (maintains google#3573 fix) - Live/audio sessions: Skip manual replay, rely on Gemini Live API's transparent session resumption Detection is based on the presence of: - Input/output transcription data - Audio or video content in recent events Changes: - Modified src/google/adk/cli/adk_web_server.py: * Added is_live_session() helper to detect live/audio sessions * Made event replay conditional based on session type * Added detailed logging for debugging session restoration - Enhanced tests/unittests/cli/test_live_session_restoration.py: * Added test_live_session_skips_replay_for_audio_sessions * Added test_text_session_replays_events_normally * Ensures both google#3573 and google#3395 fixes work correctly Dependencies: This fix builds on PR google#4194 (Issue google#3573) which adds the event replay mechanism. The google#3395 fix makes that mechanism smart enough to avoid conflicts with Gemini Live API. Testing: ✅ All 5 unit tests pass: * test_live_session_replays_all_events_on_reconnection (google#3573) * test_live_session_handles_empty_events_gracefully (google#3573) * test_live_session_continues_after_replay_failure (google#3573) * test_live_session_skips_replay_for_audio_sessions (google#3395) * test_text_session_replays_events_normally (google#3395) Fixes google#3395
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):