fix: preserve nontransparent live resumption

Author: he-yufeng

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -533,11 +533,15 @@ async def run_live(
           attempt += 1
           if not llm_request.live_connect_config:
             llm_request.live_connect_config = types.LiveConnectConfig()
-          if not llm_request.live_connect_config.session_resumption:
+          session_resumption = (
+              llm_request.live_connect_config.session_resumption
+          )
+          if not session_resumption:
+            session_resumption = types.SessionResumptionConfig()
             llm_request.live_connect_config.session_resumption = (
-                types.SessionResumptionConfig()
+                session_resumption
             )
-          llm_request.live_connect_config.session_resumption.handle = (
+          session_resumption.handle = (
               invocation_context.live_session_resumption_handle
           )
 

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -747,6 +747,78 @@ async def mock_receive_2():
       assert invocation_context.live_session_resumption_handle == 'test_handle'
 
 
+@pytest.mark.asyncio
+async def test_run_live_reconnect_preserves_nontransparent_resumption():
+  """Test that reconnect does not force transparent resumption."""
+  from google.adk.agents.live_request_queue import LiveRequestQueue
+  from websockets.exceptions import ConnectionClosed
+
+  real_model = Gemini()
+  mock_connection = mock.AsyncMock()
+
+  async def mock_receive():
+    yield LlmResponse(
+        live_session_resumption_update=types.LiveServerSessionResumptionUpdate(
+            new_handle='test_handle'
+        )
+    )
+    raise ConnectionClosed(None, None)
+
+  mock_connection.receive = mock.Mock(side_effect=mock_receive)
+
+  agent = Agent(name='test_agent', model=real_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  flow = BaseLlmFlowForTesting()
+
+  async def mock_preprocess(ctx, req):
+    req.live_connect_config.session_resumption = types.SessionResumptionConfig(
+        transparent=False
+    )
+    if False:
+      yield
+
+  with mock.patch.object(
+      flow, '_preprocess_async', side_effect=mock_preprocess
+  ):
+    with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock):
+      mock_connection_2 = mock.AsyncMock()
+
+      class StopError(Exception):
+        pass
+
+      async def mock_receive_2():
+        yield LlmResponse(
+            content=types.Content(parts=[types.Part.from_text(text='hi')])
+        )
+        raise StopError('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
+
+        try:
+          async for _ in flow.run_live(invocation_context):
+            pass
+        except StopError:
+          pass
+
+        reconnect_request = mock_connect.call_args_list[1].args[0]
+        assert (
+            reconnect_request.live_connect_config.session_resumption.transparent
+            is False
+        )
+
+
 @pytest.mark.asyncio
 async def test_run_live_skips_send_history_on_resumption():
   """Test that run_live skips send_history when resuming a session."""
@@ -1390,7 +1462,7 @@ async def mock_receive_2():
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "api_backend",
+    'api_backend',
     [
         GoogleLLMVariant.GEMINI_API,
         GoogleLLMVariant.VERTEX_AI,
@@ -1422,8 +1494,11 @@ async def mock_receive():
   flow = BaseLlmFlowForTesting()
 
   with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock):
+
     async def mock_preprocess(ctx, req):
-      req.contents = [types.Content(parts=[types.Part.from_text(text='history')])]
+      req.contents = [
+          types.Content(parts=[types.Part.from_text(text='history')])
+      ]
       yield Event(id=Event.new_id(), author='test')
 
     with mock.patch.object(
@@ -1467,7 +1542,9 @@ async def test_run_live_respects_explicit_initial_history_in_client_content_fals
   )
   invocation_context.live_request_queue = LiveRequestQueue()
   run_config = RunConfig(
-      history_config=types.HistoryConfig(initial_history_in_client_content=False)
+      history_config=types.HistoryConfig(
+          initial_history_in_client_content=False
+      )
   )
   invocation_context.run_config = run_config
 
@@ -1476,6 +1553,7 @@ async def test_run_live_respects_explicit_initial_history_in_client_content_fals
   async def mock_preprocess(ctx, req):
     req.contents = [types.Content(parts=[types.Part.from_text(text='history')])]
     from google.adk.flows.llm_flows.basic import _build_basic_request
+
     _build_basic_request(ctx, req)
     yield Event(id=Event.new_id(), author='test')
 
@@ -1509,5 +1587,7 @@ async def mock_receive():
         assert mock_connect.call_count == 1
         call_req = mock_connect.call_args[0][0]
         assert call_req.live_connect_config.history_config is not None
-        assert call_req.live_connect_config.history_config.initial_history_in_client_content is False
-
+        assert (
+            call_req.live_connect_config.history_config.initial_history_in_client_content
+            is False
+        )