Merge branch 'main' into wjy/fix-interrupted

Wu-Jiayang · web-flow · commit 7251ed3f8a29 · 2026-04-15T09:52:04.000+08:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,11 +44,11 @@ dependencies = [
   "google-cloud-spanner>=3.56.0, <4.0.0",                    # For Spanner database
   "google-cloud-speech>=2.30.0, <3.0.0",                     # For Audio Transcription
   "google-cloud-storage>=2.18.0, <4.0.0",                    # For GCS Artifact service
-  "google-genai>=1.64.0, <2.0.0",                            # Google GenAI SDK
+  "google-genai>=1.72.0, <2.0.0",                            # Google GenAI SDK
   "graphviz>=0.20.2, <1.0.0",                                # Graphviz for graph rendering
   "httpx>=0.27.0, <1.0.0",                                   # HTTP client library
   "jsonschema>=4.23.0, <5.0.0",                              # Agent Builder config validation
-  "mcp>=1.23.0, <2.0.0",                                     # For MCP Toolset
+  "mcp>=1.24.0, <2.0.0",                                     # For MCP Toolset
   "opentelemetry-api>=1.36.0, <1.39.0",                      # OpenTelemetry - keep below 1.39.0 due to current agent_engines exporter constraints.
   "opentelemetry-exporter-gcp-logging>=1.9.0a0, <2.0.0",
   "opentelemetry-exporter-gcp-monitoring>=1.9.0a0, <2.0.0",
@@ -123,7 +123,8 @@ test = [
   "a2a-sdk>=0.3.0,<0.4.0",
   "anthropic>=0.43.0",                            # For anthropic model tests
   "crewai[tools];python_version>='3.11' and python_version<'3.12'",         # For CrewaiTool tests; chromadb/pypika fail on 3.12+
-  "google-cloud-firestore>=2.11.0",
+  "google-cloud-firestore>=2.11.0, <3.0.0",
+  "google-cloud-iamconnectorcredentials>=0.1.0, <0.2.0",
   "google-cloud-parametermanager>=0.4.0, <1.0.0",
   "kubernetes>=29.0.0",                           # For GkeCodeExecutor
   "langchain-community>=0.3.17",
@@ -159,7 +160,7 @@ extensions = [
   "beautifulsoup4>=3.2.2",                      # For load_web_page tool.
   "crewai[tools];python_version>='3.11' and python_version<'3.12'",       # For CrewaiTool; chromadb/pypika fail on 3.12+
   "docker>=7.0.0",                              # For ContainerCodeExecutor
-  "google-cloud-firestore>=2.11.0",             # For Firestore services
+  "google-cloud-firestore>=2.11.0, <3.0.0",             # For Firestore services
   "google-cloud-parametermanager>=0.4.0, <1.0.0",
   "kubernetes>=29.0.0",                         # For GkeCodeExecutor
   "k8s-agent-sandbox>=0.1.1.post3",             # For GkeCodeExecutor sandbox mode
@@ -178,6 +179,10 @@ toolbox = ["toolbox-adk>=1.0.0, <2.0.0"]
 
 slack = ["slack-bolt>=1.22.0"]
 
+agent-identity = [
+  "google-cloud-iamconnectorcredentials>=0.1.0, <0.2.0",
+]
+
 [tool.pyink]
 # Format py files following Google style-guide
 line-length = 80
diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py
@@ -81,11 +81,9 @@ async def send_history(self, history: list[types.Content]):
 
     if contents:
       logger.debug('Sending history to live connection: %s', contents)
-      await self._gemini_session.send(
-          input=types.LiveClientContent(
-              turns=contents,
-              turn_complete=contents[-1].role == 'user',
-          ),
+      await self._gemini_session.send_client_content(
+          turns=contents,
+          turn_complete=contents[-1].role == 'user',
       )
     else:
       logger.info('no content is sent')
@@ -105,10 +103,8 @@ async def send_content(self, content: types.Content):
       # All parts have to be function responses.
       function_responses = [part.function_response for part in content.parts]
       logger.debug('Sending LLM function response: %s', function_responses)
-      await self._gemini_session.send(
-          input=types.LiveClientToolResponse(
-              function_responses=function_responses
-          ),
+      await self._gemini_session.send_tool_response(
+          function_responses=function_responses
       )
     else:
       logger.debug('Sending LLM new content %s', content)
diff --git a/src/google/adk/plugins/base_plugin.py b/src/google/adk/plugins/base_plugin.py
@@ -155,10 +155,10 @@ async def before_run_callback(
   async def on_event_callback(
       self, *, invocation_context: InvocationContext, event: Event
   ) -> Optional[Event]:
-    """Callback executed after an event is yielded from runner.
+    """Callback executed when the runner produces an event.
 
-    This is the ideal place to make modification to the event before the event
-    is handled by the underlying agent app.
+    This is the ideal place to modify the event before it is persisted to the
+    session service and yielded to the caller.
 
     Args:
       invocation_context: The context for the entire invocation.
diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py
@@ -791,6 +791,34 @@ def _should_append_event(self, event: Event, is_live_call: bool) -> bool:
       return False
     return True
 
+  def _get_output_event(
+      self,
+      *,
+      original_event: Event,
+      modified_event: Event | None,
+      run_config: RunConfig | None,
+  ) -> Event:
+    """Returns the event that should be persisted and yielded.
+
+    Plugins may return a replacement event that only overrides a subset of
+    fields. Merge those changes onto the original event so the streamed event
+    and the persisted event stay aligned without losing the original event
+    identity.
+    """
+    if modified_event is None:
+      return original_event
+
+    _apply_run_config_custom_metadata(modified_event, run_config)
+    update = {}
+    for field_name in modified_event.model_fields_set:
+      if field_name in {'id', 'invocation_id', 'timestamp'}:
+        continue
+      update[field_name] = modified_event.__dict__[field_name]
+    output_event = original_event.model_copy(update=update)
+    if not output_event.author:
+      output_event.author = original_event.author
+    return output_event
+
   async def _exec_with_plugin(
       self,
       invocation_context: InvocationContext,
@@ -854,13 +882,24 @@ async def _exec_with_plugin(
           _apply_run_config_custom_metadata(
               event, invocation_context.run_config
           )
+          # Step 3: Run the on_event callbacks before persisting so callback
+          # changes are stored in the session and match the streamed event.
+          modified_event = await plugin_manager.run_on_event_callback(
+              invocation_context=invocation_context, event=event
+          )
+          output_event = self._get_output_event(
+              original_event=event,
+              modified_event=modified_event,
+              run_config=invocation_context.run_config,
+          )
+
           if is_live_call:
             if event.partial and _is_transcription(event):
               is_transcribing = True
             if is_transcribing and _is_tool_call_or_response(event):
               # only buffer function call and function response event which is
               # non-partial
-              buffered_events.append(event)
+              buffered_events.append(output_event)
               continue
             # Note for live/bidi: for audio response, it's considered as
             # non-partial event(event.partial=None)
@@ -881,7 +920,7 @@ async def _exec_with_plugin(
                 )
                 if self._should_append_event(event, is_live_call):
                   await self.session_service.append_event(
-                      session=session, event=event
+                      session=session, event=output_event
                   )
 
                 for buffered_event in buffered_events:
@@ -897,25 +936,15 @@ async def _exec_with_plugin(
                 if self._should_append_event(event, is_live_call):
                   logger.debug('Appending non-buffered event: %s', event)
                   await self.session_service.append_event(
-                      session=session, event=event
+                      session=session, event=output_event
                   )
           else:
             if event.partial is not True:
               await self.session_service.append_event(
-                  session=session, event=event
+                  session=session, event=output_event
               )
 
-          # Step 3: Run the on_event callbacks to optionally modify the event.
-          modified_event = await plugin_manager.run_on_event_callback(
-              invocation_context=invocation_context, event=event
-          )
-          if modified_event:
-            _apply_run_config_custom_metadata(
-                modified_event, invocation_context.run_config
-            )
-            yield modified_event
-          else:
-            yield event
+          yield output_event
 
     # Step 4: Run the after_run callbacks to perform global cleanup tasks or
     # finalizing logs and metrics data.
diff --git a/tests/unittests/models/test_gemini_llm_connection.py b/tests/unittests/models/test_gemini_llm_connection.py
@@ -81,11 +81,11 @@ async def test_send_history(gemini_connection, mock_gemini_session):
 
   await gemini_connection.send_history(history)
 
-  mock_gemini_session.send.assert_called_once()
-  call_args = mock_gemini_session.send.call_args[1]
-  assert 'input' in call_args
-  assert call_args['input'].turns == history
-  assert call_args['input'].turn_complete is False  # Last message is from model
+  mock_gemini_session.send_client_content.assert_called_once()
+  call_args = mock_gemini_session.send_client_content.call_args[1]
+  assert 'turns' in call_args
+  assert call_args['turns'] == history
+  assert call_args['turn_complete'] is False  # Last message is from model
 
 
 @pytest.mark.asyncio
@@ -118,10 +118,10 @@ async def test_send_content_function_response(
 
   await gemini_connection.send_content(content)
 
-  mock_gemini_session.send.assert_called_once()
-  call_args = mock_gemini_session.send.call_args[1]
-  assert 'input' in call_args
-  assert call_args['input'].function_responses == [function_response]
+  mock_gemini_session.send_tool_response.assert_called_once()
+  call_args = mock_gemini_session.send_tool_response.call_args[1]
+  assert 'function_responses' in call_args
+  assert call_args['function_responses'] == [function_response]
 
 
 @pytest.mark.asyncio
@@ -668,9 +668,9 @@ async def test_send_history_filters_audio(mock_gemini_session, audio_part):
 
   await connection.send_history(history)
 
-  mock_gemini_session.send.assert_called_once()
-  call_args = mock_gemini_session.send.call_args[1]
-  sent_contents = call_args['input'].turns
+  mock_gemini_session.send_client_content.assert_called_once()
+  call_args = mock_gemini_session.send_client_content.call_args[1]
+  sent_contents = call_args['turns']
   # Only the model response should be sent (user audio filtered out)
   assert len(sent_contents) == 1
   assert sent_contents[0].role == 'model'
@@ -696,9 +696,9 @@ async def test_send_history_keeps_image_data(mock_gemini_session):
 
   await connection.send_history(history)
 
-  mock_gemini_session.send.assert_called_once()
-  call_args = mock_gemini_session.send.call_args[1]
-  sent_contents = call_args['input'].turns
+  mock_gemini_session.send_client_content.assert_called_once()
+  call_args = mock_gemini_session.send_client_content.call_args[1]
+  sent_contents = call_args['turns']
   # Both contents should be sent (image is not filtered)
   assert len(sent_contents) == 2
   assert sent_contents[0].parts[0].inline_data == image_blob
@@ -728,9 +728,9 @@ async def test_send_history_mixed_content_filters_only_audio(
 
   await connection.send_history(history)
 
-  mock_gemini_session.send.assert_called_once()
-  call_args = mock_gemini_session.send.call_args[1]
-  sent_contents = call_args['input'].turns
+  mock_gemini_session.send_client_content.assert_called_once()
+  call_args = mock_gemini_session.send_client_content.call_args[1]
+  sent_contents = call_args['turns']
   # Content should be sent but only with the text part
   assert len(sent_contents) == 1
   assert len(sent_contents[0].parts) == 1
diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py
@@ -139,6 +139,7 @@ class MockPlugin(BasePlugin):
       "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
   )
   ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
+  ON_EVENT_CALLBACK_METADATA = {"plugin_key": "plugin_value"}
 
   def __init__(self):
     super().__init__(name="mock_plugin")
@@ -184,6 +185,7 @@ async def on_event_callback(
             ],
             role=event.content.role,
         ),
+        custom_metadata=self.ON_EVENT_CALLBACK_METADATA,
     )
 
 
@@ -359,6 +361,60 @@ async def test_run_live_auto_create_session():
   assert session is not None
 
 
+@pytest.mark.asyncio
+async def test_run_live_persists_event_callback_modifications():
+  """run_live should persist the same event it streams after callback changes."""
+  session_service = InMemorySessionService()
+  artifact_service = InMemoryArtifactService()
+  plugin = MockPlugin()
+  plugin.enable_event_callback = True
+  runner = Runner(
+      app_name="live_app",
+      agent=MockLiveAgent("live_agent"),
+      session_service=session_service,
+      artifact_service=artifact_service,
+      plugins=[plugin],
+  )
+  await session_service.create_session(
+      app_name="live_app", user_id="user", session_id="live_session"
+  )
+
+  from google.adk.agents.live_request_queue import LiveRequestQueue
+
+  live_queue = LiveRequestQueue()
+  agen = runner.run_live(
+      user_id="user",
+      session_id="live_session",
+      live_request_queue=live_queue,
+  )
+
+  streamed_event = await agen.__anext__()
+  await agen.aclose()
+
+  session = await session_service.get_session(
+      app_name="live_app", user_id="user", session_id="live_session"
+  )
+  persisted_event = session.events[0]
+
+  assert streamed_event.author == "live_agent"
+  assert streamed_event.invocation_id
+  assert streamed_event.content.parts[0].text == (
+      MockPlugin.ON_EVENT_CALLBACK_MSG
+  )
+  assert streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+
+  assert persisted_event.id == streamed_event.id
+  assert persisted_event.timestamp == streamed_event.timestamp
+  assert persisted_event.author == streamed_event.author
+  assert persisted_event.invocation_id == streamed_event.invocation_id
+  assert persisted_event.content.parts[0].text == (
+      MockPlugin.ON_EVENT_CALLBACK_MSG
+  )
+  assert (
+      persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+  )
+
+
 @pytest.mark.asyncio
 async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch):
   project_root = tmp_path / "workspace"
@@ -747,6 +803,39 @@ async def test_runner_modifies_event_after_execution(self):
 
     assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG
 
+  @pytest.mark.asyncio
+  async def test_runner_persists_event_callback_modifications(self):
+    """Event callback output should be persisted, not only streamed."""
+    self.plugin.enable_event_callback = True
+
+    events = await self.run_test()
+    streamed_event = events[0]
+
+    session = await self.session_service.get_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    persisted_event = session.events[1]
+
+    assert streamed_event.author == "test_agent"
+    assert streamed_event.invocation_id
+    assert streamed_event.content.parts[0].text == (
+        MockPlugin.ON_EVENT_CALLBACK_MSG
+    )
+    assert (
+        streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+    )
+
+    assert persisted_event.id == streamed_event.id
+    assert persisted_event.timestamp == streamed_event.timestamp
+    assert persisted_event.author == streamed_event.author
+    assert persisted_event.invocation_id == streamed_event.invocation_id
+    assert persisted_event.content.parts[0].text == (
+        MockPlugin.ON_EVENT_CALLBACK_MSG
+    )
+    assert (
+        persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+    )
+
   @pytest.mark.asyncio
   async def test_runner_close_calls_plugin_close(self):
     """Test that runner.close() calls plugin manager close."""
