fix(tools): pass the caller's temp: state to an AgentTool's wrapped agent

src/google/adk/tools/agent_tool.py

@@ -229,6 +229,7 @@ async def run_async(
   ) -> Any:
     from ..runners import Runner
     from ..sessions.in_memory_session_service import InMemorySessionService
+    from ..sessions.state import State
 
     if self.skip_summarization:
       tool_context.actions.skip_summarization = True
@@ -286,6 +287,14 @@ async def run_async(
         for k, v in tool_context.state.to_dict().items()
         if not k.startswith('_adk')  # Filter out adk internal states
     }
+    # create_session drops temp: keys, since they are never stored. Hand them
+    # to the nested run as a state delta instead, the way run_async applies
+    # temp: state for one invocation.
+    temp_state = {
+        k: state_dict.pop(k)
+        for k in list(state_dict)
+        if k.startswith(State.TEMP_PREFIX)
+    }
     session = await runner.session_service.create_session(
         app_name=child_app_name,
         user_id=tool_context._invocation_context.user_id,
@@ -358,6 +367,7 @@ async def run_async(
             user_id=session.user_id,
             session_id=session.id,
             new_message=content,
+            state_delta=temp_state or None,
             run_config=nested_run_config,
             abort_signal=abort_signal,
         )

tests/unittests/tools/test_agent_tool.py

@@ -555,6 +555,111 @@ def test_update_state():
   assert runner.session.state['state_1'] == 'changed_value'
 
 
+class _NotJsonSerializable:
+  """A value only temp: state can hold, since it is never persisted."""
+
+
+@mark.parametrize(
+    'reader',
+    [
+        'instruction',
+        'optional_instruction',
+        'callback',
+        'tool',
+        'sequential_agent',
+        'input_schema',
+    ],
+)
+def test_agent_tool_reads_caller_temp_state(reader):
+  """The wrapped agent can read the caller's temp: state."""
+  client = _NotJsonSerializable()
+  seen = {}
+
+  def set_caller_state(callback_context: CallbackContext):
+    callback_context.state['temp:ticket'] = 'T-42'
+    callback_context.state['temp:client'] = client
+    callback_context.state['plain'] = 'P'
+
+  def read_in_callback(callback_context: CallbackContext):
+    for key in ('temp:ticket', 'temp:client', 'plain'):
+      seen[key] = callback_context.state.get(key)
+
+  def read_in_tool(tool_context: ToolContext) -> str:
+    for key in ('temp:ticket', 'temp:client', 'plain'):
+      seen[key] = tool_context.state.get(key)
+    return 'read'
+
+  caller_deltas = []
+
+  def record_caller_delta(tool, args, tool_context, tool_response):
+    caller_deltas.append(dict(tool_context.actions.state_delta))
+
+  class TicketInput(BaseModel):
+    custom_input: str
+
+  template = 'ticket: {temp:ticket} plain: {plain}'
+  tool_call = (
+      function_call_custom
+      if reader == 'input_schema'
+      else function_call_no_schema
+  )
+  responses = ['response1', 'response2']
+  if reader == 'tool':
+    responses.insert(0, Part.from_function_call(name='read_in_tool', args={}))
+  mock_model = testing_utils.MockModel.create(responses=[tool_call, *responses])
+  if reader == 'sequential_agent':
+    tool_agent = SequentialAgent(
+        name='tool_agent',
+        sub_agents=[
+            Agent(name='inner_agent', model=mock_model, instruction=template)
+        ],
+    )
+  elif reader == 'input_schema':
+    tool_agent = Agent(
+        name='tool_agent',
+        model=mock_model,
+        input_schema=TicketInput,
+        instruction=template,
+    )
+  else:
+    tool_agent = Agent(
+        name='tool_agent',
+        model=mock_model,
+        instruction={
+            'instruction': template,
+            'optional_instruction': 'ticket: {temp:ticket?} plain: {plain?}',
+        }.get(reader, 'no template'),
+        before_agent_callback=(
+            read_in_callback if reader == 'callback' else None
+        ),
+        tools=[read_in_tool] if reader == 'tool' else [],
+    )
+  root_agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_agent_callback=set_caller_state,
+      after_tool_callback=record_caller_delta,
+      tools=[AgentTool(agent=tool_agent)],
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent)
+  runner.run('test1')
+
+  # The seeded temp: keys are not echoed back into the caller's state by the
+  # state_delta that AgentTool forwards from the nested run.
+  assert len(caller_deltas) == 1
+  assert not any(key.startswith('temp:') for key in caller_deltas[0])
+
+  if reader in ('callback', 'tool'):
+    assert seen['temp:ticket'] == 'T-42'
+    assert seen['temp:client'] is client
+    assert seen['plain'] == 'P'
+  else:
+    assert mock_model.requests[1].config.system_instruction.startswith(
+        'ticket: T-42 plain: P'
+    )
+
+
 @mark.asyncio
 async def test_update_artifacts():
   """The agent tool can read and write artifacts."""