fix(flows): Reset reconnect attempts on connection success

Reset the reconnect attempt counter immediately when a live connection is established successfully, rather than waiting for the first event to be received. This prevents reconnection failures if the connection succeeds but no events are received before a transient disconnect occurs.

Change-Id: I9a8002558525fad5715b9a0ac5cf2634a92b923f

Author: wukath

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

@@ -556,6 +556,8 @@ async def run_live(
             invocation_context.agent.name,
         )
         async with llm.connect(llm_request) as llm_connection:
+          # Reset attempt counter on successful connection.
+          attempt = 1
           # Skip sending history if we are resuming a session. The server
           # already has the state associated with the resumption handle.
           if (
@@ -585,8 +587,6 @@ async def run_live(
                 )
             ) as agen:
               async for event in agen:
-                # Reset attempt counter on successful communication.
-                attempt = 1
                 # Empty event means the queue is closed.
                 if not event:
                   break

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -1402,3 +1402,53 @@ async def mock_receive():
             call_req.live_connect_config.history_config.initial_history_in_client_content
             is False
         )
+
+
+@pytest.mark.asyncio
+async def test_run_live_reset_attempt_on_connection_success():
+  """Test that attempt counter is reset immediately when connection succeeds."""
+  # Arrange
+  real_model = Gemini()
+  connection_cnt = 0
+
+  class StopTestError(Exception):
+    pass
+
+  async def mock_connect_impl(*args, **kwargs):
+    nonlocal connection_cnt
+    connection_cnt += 1
+    conn = mock.AsyncMock()
+
+    async def mock_receive():
+      if connection_cnt < 8:
+        # Raise ConnectionClosed immediately without yielding any events
+        raise ConnectionClosed(None, None)
+      else:
+        raise StopTestError('success')
+
+    conn.receive = mock.Mock(side_effect=mock_receive)
+    return conn
+
+  agent = Agent(name='test_agent', model=real_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+  # Pre-populate resumption handle so the flow attempts to reconnect
+  invocation_context.live_session_resumption_handle = 'test_handle'
+
+  flow = BaseLlmFlowForTesting()
+
+  # Act & Assert
+  with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock):
+    with mock.patch(
+        'google.adk.models.google_llm.Gemini.connect'
+    ) as mock_connect:
+      mock_connect.return_value.__aenter__.side_effect = mock_connect_impl
+      try:
+        async for _ in flow.run_live(invocation_context):
+          pass
+      except StopTestError:
+        pass
+
+      assert mock_connect.call_count == 8