From 1de09c85c84445cf948af3c40bfb7bf4f906f9bc Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Tue, 18 Aug 2026 03:17:57 +0800 Subject: [PATCH] fix(live): stop a live run from writing session state onto the RunConfig The basic request processor forwards `RunConfig.session_resumption` and `RunConfig.history_config` to `LiveConnectConfig` by reference, and the connect loop in `BaseLlmFlow.run_live` then writes onto whatever object it finds there: the newest server-issued resumption handle on every reconnect, `transparent` on the Vertex AI backend, and `initial_history_in_client_content` when it replays history. All of that lands on the caller's own RunConfig, so a caller that reuses it for the next run resumes the session that just ended instead of starting a new one. Forward a copy of both, which is what `http_options` already does for the same reason. Agent transfer is unaffected: it clears the handle on a deep-copied run config, and the child's request assembly copies from there. --- src/google/adk/flows/llm_flows/basic.py | 15 +- .../flows/llm_flows/test_base_llm_flow.py | 139 ++++++++++++++++++ .../flows/llm_flows/test_basic_processor.py | 28 ++++ 3 files changed, 180 insertions(+), 2 deletions(-) diff --git a/src/google/adk/flows/llm_flows/basic.py b/src/google/adk/flows/llm_flows/basic.py index bfaf1343ea..a26346f453 100644 --- a/src/google/adk/flows/llm_flows/basic.py +++ b/src/google/adk/flows/llm_flows/basic.py @@ -183,10 +183,21 @@ def _build_basic_request( llm_request.live_connect_config.proactivity = ( None if is_gemini_3_x else run_config.proactivity ) + # Copied in rather than aliased, for the same reason as http_options above: + # the connect loop in `BaseLlmFlow.run_live` writes the newest resumption + # handle, the Vertex-only `transparent` default and + # `initial_history_in_client_content` onto whatever object it finds here, and + # the RunConfig belongs to the caller, who may reuse it for the next run. llm_request.live_connect_config.session_resumption = ( - run_config.session_resumption + run_config.session_resumption.model_copy() + if run_config.session_resumption + else None + ) + llm_request.live_connect_config.history_config = ( + run_config.history_config.model_copy() + if run_config.history_config + else None ) - llm_request.live_connect_config.history_config = run_config.history_config llm_request.live_connect_config.context_window_compression = ( run_config.context_window_compression ) diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 2511a6a09f..bd461bad45 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -1285,6 +1285,145 @@ async def mock_receive_2(): ) +@pytest.mark.asyncio +async def test_run_live_does_not_write_the_session_handle_onto_the_run_config(): + """The newest handle must not be written back onto the caller's RunConfig. + + The RunConfig belongs to the caller and may be reused for the next run. A + handle written back onto it points at the session that just ended, so the + next run would try to resume a dead session instead of starting a new one. + """ + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + async def mock_receive(): + yield LlmResponse( + live_session_resumption_update=types.LiveServerSessionResumptionUpdate( + new_handle='server_handle' + ) + ) + raise ConnectionClosed(None, None) + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + agent = Agent(name='test_agent', model=real_model) + run_config = RunConfig( + session_resumption=types.SessionResumptionConfig(handle='caller_handle') + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, run_config=run_config + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with ( + mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_basic + ), + mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock), + ): + mock_connection_2 = mock.AsyncMock() + + class NonRetryableError(Exception): + pass + + async def mock_receive_2(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise NonRetryableError('stop') + + mock_connection_2.receive = mock.Mock(side_effect=mock_receive_2) + + mock_aenter = mock.AsyncMock() + mock_aenter.side_effect = [mock_connection, mock_connection_2] + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__ = mock_aenter + + with mock.patch.object( + Gemini, + '_api_backend', + new_callable=mock.PropertyMock, + return_value=GoogleLLMVariant.VERTEX_AI, + ): + try: + async for _ in flow.run_live(invocation_context): + pass + except NonRetryableError: + pass + + # The reconnect did use the server handle. + second_request = mock_connect.call_args_list[1][0][0] + assert ( + second_request.live_connect_config.session_resumption.handle + == 'server_handle' + ) + # ...without that, or the Vertex-only transparent default, reaching the + # caller's config. + assert run_config.session_resumption == types.SessionResumptionConfig( + handle='caller_handle' + ) + + +@pytest.mark.asyncio +async def test_run_live_does_not_write_initial_history_onto_the_run_config(): + """The initial-history flag must not be written onto the caller's RunConfig.""" + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + agent = Agent(name='test_agent', model=real_model) + run_config = RunConfig(history_config=types.HistoryConfig()) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, run_config=run_config + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with ( + mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_with_history + ), + mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock), + ): + + class StopError(Exception): + pass + + async def mock_receive(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise StopError('stop') + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__.return_value = mock_connection + + try: + async for _ in flow.run_live(invocation_context): + pass + except StopError: + pass + + # The request declares the replayed history as client-provided... + connect_request = mock_connect.call_args[0][0] + assert ( + connect_request.live_connect_config.history_config.initial_history_in_client_content + ) + # ...while the caller's config keeps saying nothing about it. + assert run_config.history_config == types.HistoryConfig() + + @pytest.mark.asyncio async def test_run_live_does_not_log_http_options_headers(caplog): """run_live must not log http_options headers, which can carry secrets.""" diff --git a/tests/unittests/flows/llm_flows/test_basic_processor.py b/tests/unittests/flows/llm_flows/test_basic_processor.py index 140f0cc0a5..2221dba5d8 100644 --- a/tests/unittests/flows/llm_flows/test_basic_processor.py +++ b/tests/unittests/flows/llm_flows/test_basic_processor.py @@ -484,6 +484,34 @@ async def test_run_config_http_options_object_is_not_aliased(self): llm_request.config.http_options.headers['Injected'] = 'x' assert 'Injected' not in run_config_http_options.headers + @pytest.mark.asyncio + async def test_run_config_live_session_objects_are_not_aliased(self): + """The request must not hold the RunConfig's own live session objects.""" + agent = LlmAgent(name='test_agent', model='gemini-1.5-flash') + + invocation_context = await _create_invocation_context(agent) + session_resumption = types.SessionResumptionConfig(handle='caller_handle') + history_config = types.HistoryConfig() + invocation_context.run_config = RunConfig( + session_resumption=session_resumption, history_config=history_config + ) + llm_request = LlmRequest() + + processor = _BasicLlmRequestProcessor() + async for _ in processor.run_async(invocation_context, llm_request): + pass + + # What the connect loop in `run_live` writes as a session goes on. + live_config = llm_request.live_connect_config + live_config.session_resumption.handle = 'server_handle' + live_config.session_resumption.transparent = True + live_config.history_config.initial_history_in_client_content = True + + assert session_resumption == types.SessionResumptionConfig( + handle='caller_handle' + ) + assert history_config == types.HistoryConfig() + @pytest.mark.asyncio async def test_http_options_carrying_an_unpicklable_client_are_copied(self): """http_options can hold a live client, which no deep copy survives."""