feat(plugins): Support opt-in session retention in MultimodalToolResultsPlugin

src/google/adk/plugins/multimodal_tool_results_plugin.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Literal
 from typing import Optional
 
 from google.genai import types
@@ -27,6 +28,14 @@
 from .base_plugin import BasePlugin
 
 PARTS_RETURNED_BY_TOOLS_ID = "temp:PARTS_RETURNED_BY_TOOLS_ID"
+# Deliberately NOT "temp:"-prefixed: the session layer treats "temp:" state
+# as invocation-scoped and strips it before persisting an event (see
+# BaseSessionService._trim_temp_delta_state). retention="session" needs the
+# saved parts to survive into later invocations (i.e. later conversational
+# turns), so it is stored under this session-scoped key instead.
+SESSION_PARTS_RETURNED_BY_TOOLS_ID = (
+    "multimodal_tool_results_plugin:PARTS_RETURNED_BY_TOOLS_ID"
+)
 
 
 class MultimodalToolResultsPlugin(BasePlugin):
@@ -36,13 +45,29 @@ class MultimodalToolResultsPlugin(BasePlugin):
   are supported outside of computer use tool.
   """
 
-  def __init__(self, name: str = "multimodal_tool_results_plugin"):
+  def __init__(
+      self,
+      name: str = "multimodal_tool_results_plugin",
+      retention: Literal["next_model_call", "session"] = "next_model_call",
+  ):
     """Initialize the multimodal tool results plugin.
 
     Args:
       name: The name of the plugin instance.
+      retention: How long tool-returned parts stay attached to model
+        requests. "next_model_call" (default) attaches the saved parts once
+        and then clears them. "session" keeps re-attaching the latest saved
+        parts to every subsequent model request for the rest of the
+        session, so follow-up turns can still reference them.
     """
     super().__init__(name)
+    self._retention = retention
+
+  def _state_key(self) -> str:
+    """Returns the state key parts are stored under for this retention mode."""
+    if self._retention == "session":
+      return SESSION_PARTS_RETURNED_BY_TOOLS_ID
+    return PARTS_RETURNED_BY_TOOLS_ID
 
   async def after_tool_callback(
       self,
@@ -67,11 +92,12 @@ async def after_tool_callback(
       return result
 
     parts = [result] if isinstance(result, types.Part) else result[:]
+    key = self._state_key()
 
-    if PARTS_RETURNED_BY_TOOLS_ID in tool_context.state:
-      tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] += parts
+    if key in tool_context.state:
+      tool_context.state[key] += parts
     else:
-      tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] = parts
+      tool_context.state[key] = parts
 
     return None
 
@@ -83,10 +109,10 @@ async def before_model_callback(
     if not llm_request.contents:
       return None
 
-    if saved_parts := callback_context.state.get(
-        PARTS_RETURNED_BY_TOOLS_ID, None
-    ):
+    key = self._state_key()
+    if saved_parts := callback_context.state.get(key, None):
       llm_request.contents[-1].parts += saved_parts
-      callback_context.state.update({PARTS_RETURNED_BY_TOOLS_ID: []})
+      if self._retention == "next_model_call":
+        callback_context.state.update({key: []})
 
     return None

tests/unittests/plugins/test_multimodal_tool_results_plugin.py

@@ -19,6 +19,7 @@
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
 from google.adk.models.llm_request import LlmRequest
 from google.adk.plugins.multimodal_tool_results_plugin import MultimodalToolResultsPlugin
 from google.adk.plugins.multimodal_tool_results_plugin import PARTS_RETURNED_BY_TOOLS_ID
@@ -144,6 +145,53 @@ async def test_empty_contents_leaves_saved_parts_pending(
   assert tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] == parts
 
 
+@pytest.mark.asyncio
+async def test_session_retention_reattaches_parts_across_turns():
+  """Test that retention="session" resurfaces parts in a LATER invocation.
+
+  Regression test for #6695. This must go through two separate
+  runner.run_async() calls against a real session service: the saved parts
+  are stored under a "temp:"-prefixed key by default, and that prefix is
+  stripped by BaseSessionService before an event is persisted, so a test
+  that only calls before_model_callback twice on the same in-memory State
+  object (without an intervening append_event()) cannot detect whether the
+  parts actually survive a real turn boundary.
+  """
+  file_part = types.Part(
+      file_data=types.FileData(
+          file_uri="gs://bucket/document.pdf", mime_type="application/pdf"
+      )
+  )
+
+  def get_document() -> types.Part:
+    return file_part
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(name="get_document", args={}),
+          "Here is a summary of the document.",
+          "The document says X.",
+      ]
+  )
+  agent = Agent(name="root_agent", model=mock_model, tools=[get_document])
+  runner = testing_utils.InMemoryRunner(
+      agent, plugins=[MultimodalToolResultsPlugin(retention="session")]
+  )
+
+  # Turn 1: triggers the tool call.
+  await runner.run_async("Please fetch the document")
+  # Turn 2: a NEW invocation, sharing the same session as turn 1.
+  await runner.run_async("What does the document say?")
+
+  assert len(mock_model.requests) == 3
+  # Turn 1's first request precedes the tool call: nothing attached yet.
+  assert file_part not in mock_model.requests[0].contents[-1].parts
+  # Turn 1's second request: parts attached within the same invocation.
+  assert file_part in mock_model.requests[1].contents[-1].parts
+  # Turn 2's request is a separate invocation: parts must still be attached.
+  assert file_part in mock_model.requests[2].contents[-1].parts
+
+
 @pytest.mark.asyncio
 async def test_multiple_tools_returning_parts_are_accumulated(
     plugin: ToolReturningGenAiPartsPlugin,