Skip to content

Commit 12c7bed

Browse files
Somay-kousiscopybara-github
authored andcommitted
fix: pass the caller's temp: state to an AgentTool's wrapped agent
Merge #7596 The wrapped agent now sees the caller's temp: keys in its instruction, callbacks and tools; before, they were dropped when the nested session was created. ADK's own temp:_adk keys are still not passed on. The keys travel one way: dict, list and set values are shallow-copied so in-place edits inside the wrapped agent do not reach the caller's state, and temp: writes made inside the wrapped agent are not copied back. Fixes #7595 PiperOrigin-RevId: 1006234842
1 parent 25d1917 commit 12c7bed

3 files changed

Lines changed: 171 additions & 7 deletions

File tree

‎docs/guides/tools/agent_tool/index.md‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -131,9 +131,11 @@ services:
131131

132132
- The session is new, belongs to the same user as the parent's session, and
133133
starts with a copy of the parent's session state. ADK leaves out keys that
134-
begin with `_adk`, because it uses them for its own bookkeeping. Placeholders
135-
in the wrapped agent's instruction, such as `{customer_name}`, resolve against
136-
that copy.
134+
begin with `_adk` or `temp:_adk`, because it uses them for its own
135+
bookkeeping. Placeholders in the wrapped agent's instruction, such as
136+
`{customer_name}`, resolve against that copy.
137+
- The parent's `temp:` keys are passed to the wrapped agent's invocation too, so
138+
its instruction, callbacks and tools can read them.
137139
- The only message in the session is the request from the parent's model. ADK
138140
includes none of the parent's earlier turns.
139141
- The wrapped agent shares the parent's artifacts. When it loads or saves an
@@ -161,9 +163,11 @@ message from the run, or an empty string.
161163

162164
State changes travel back as well. ADK copies every change that the wrapped
163165
agent makes to session state into the parent's state while the wrapped agent
164-
runs. Setting
165-
`output_key` on the wrapped agent therefore stores its answer in the parent's
166-
session state, where later instructions and tools can read it.
166+
runs. Setting `output_key` on the wrapped agent therefore stores its answer in
167+
the parent's session state, where later instructions and tools can read it.
168+
`temp:` state only travels one way: the parent's `temp:` keys reach the wrapped
169+
agent, but `temp:` writes made inside the wrapped agent are trimmed when its
170+
events are recorded and are not copied back to the parent's state.
167171

168172
## Configuration options
169173

‎src/google/adk/tools/agent_tool.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from __future__ import annotations
1616

1717
import asyncio
18+
import copy
1819
import json
1920
import logging
2021
from typing import Any
@@ -231,6 +232,7 @@ async def run_async(
231232
) -> Any:
232233
from ..runners import Runner
233234
from ..sessions.in_memory_session_service import InMemorySessionService
235+
from ..sessions.state import State
234236

235237
if self.skip_summarization:
236238
tool_context.actions.skip_summarization = True
@@ -309,8 +311,19 @@ async def run_async(
309311
state_dict = {
310312
k: v
311313
for k, v in tool_context.state.to_dict().items()
312-
if not k.startswith('_adk') # Filter out adk internal states
314+
# Filter out adk internal states, temp-scoped ones included.
315+
if not k.startswith(('_adk', f'{State.TEMP_PREFIX}_adk'))
313316
}
317+
# create_session drops temp: keys, since they are never stored. Hand them
318+
# to the nested run as a state delta instead, the way run_async applies
319+
# temp: state for one invocation. Containers are shallow-copied so in-place
320+
# edits in the nested run do not leak back into the caller's state, while
321+
# other values such as client handles keep their identity.
322+
temp_state: dict[str, Any] = {}
323+
for k in list(state_dict):
324+
if k.startswith(State.TEMP_PREFIX):
325+
v = state_dict.pop(k)
326+
temp_state[k] = copy.copy(v) if isinstance(v, (dict, list, set)) else v
314327
session = await runner.session_service.create_session(
315328
app_name=child_app_name,
316329
user_id=tool_context._invocation_context.user_id,
@@ -383,6 +396,7 @@ async def run_async(
383396
user_id=session.user_id,
384397
session_id=session.id,
385398
new_message=content,
399+
state_delta=temp_state or None,
386400
run_config=nested_run_config,
387401
abort_signal=abort_signal,
388402
)

‎tests/unittests/tools/test_agent_tool.py‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,152 @@ def test_update_state():
557557
assert runner.session.state['state_1'] == 'changed_value'
558558

559559

560+
@mark.parametrize(
561+
'reader',
562+
[
563+
'instruction',
564+
'optional_instruction',
565+
'callback',
566+
'tool',
567+
'sequential_agent',
568+
'input_schema',
569+
],
570+
)
571+
def test_agent_tool_reads_caller_temp_state(reader):
572+
"""The wrapped agent reads the caller's temp: state, but not adk's own."""
573+
574+
class _NotJsonSerializable:
575+
"""A value only temp: state can hold, since it is never persisted."""
576+
577+
class TicketInput(BaseModel):
578+
custom_input: str
579+
580+
client = _NotJsonSerializable()
581+
seen = {}
582+
keys = ('temp:ticket', 'temp:client', 'plain', 'temp:_adk_internal')
583+
584+
def set_caller_state(callback_context: CallbackContext):
585+
callback_context.state['temp:ticket'] = 'T-42'
586+
callback_context.state['temp:client'] = client
587+
callback_context.state['plain'] = 'P'
588+
callback_context.state['temp:_adk_internal'] = 'X'
589+
590+
def read_in_callback(callback_context: CallbackContext):
591+
for key in keys:
592+
seen[key] = callback_context.state.get(key)
593+
594+
def read_in_tool(tool_context: ToolContext) -> str:
595+
for key in keys:
596+
seen[key] = tool_context.state.get(key)
597+
return 'read'
598+
599+
template = 'ticket: {temp:ticket} plain: {plain}'
600+
tool_call = (
601+
function_call_custom
602+
if reader == 'input_schema'
603+
else function_call_no_schema
604+
)
605+
responses = ['response1', 'response2']
606+
if reader == 'tool':
607+
responses.insert(0, Part.from_function_call(name='read_in_tool', args={}))
608+
mock_model = testing_utils.MockModel.create(responses=[tool_call, *responses])
609+
if reader == 'sequential_agent':
610+
tool_agent = SequentialAgent(
611+
name='tool_agent',
612+
sub_agents=[
613+
Agent(name='inner_agent', model=mock_model, instruction=template)
614+
],
615+
)
616+
elif reader == 'input_schema':
617+
tool_agent = Agent(
618+
name='tool_agent',
619+
model=mock_model,
620+
input_schema=TicketInput,
621+
instruction=template,
622+
)
623+
else:
624+
tool_agent = Agent(
625+
name='tool_agent',
626+
model=mock_model,
627+
instruction={
628+
'instruction': template,
629+
'optional_instruction': 'ticket: {temp:ticket?} plain: {plain?}',
630+
}.get(reader, 'no template'),
631+
before_agent_callback=(
632+
read_in_callback if reader == 'callback' else None
633+
),
634+
tools=[read_in_tool] if reader == 'tool' else [],
635+
)
636+
root_agent = Agent(
637+
name='root_agent',
638+
model=mock_model,
639+
before_agent_callback=set_caller_state,
640+
tools=[AgentTool(agent=tool_agent)],
641+
)
642+
643+
runner = testing_utils.InMemoryRunner(root_agent)
644+
runner.run('test1')
645+
646+
if reader in ('callback', 'tool'):
647+
assert seen['temp:ticket'] == 'T-42'
648+
assert seen['temp:client'] is client
649+
assert seen['plain'] == 'P'
650+
assert seen['temp:_adk_internal'] is None
651+
else:
652+
assert mock_model.requests[1].config.system_instruction.startswith(
653+
'ticket: T-42 plain: P'
654+
)
655+
656+
657+
def test_agent_tool_temp_state_is_not_copied_back_to_caller():
658+
"""temp: state the wrapped agent edits or writes stays in the nested run."""
659+
child_parts = []
660+
661+
def set_caller_state(callback_context: CallbackContext):
662+
callback_context.state['temp:ticket'] = 'T-42'
663+
callback_context.state['temp:parts'] = []
664+
665+
def edit_temp_state(tool_context: ToolContext) -> str:
666+
tool_context.state['temp:parts'] += ['child_part']
667+
tool_context.state['temp:ticket'] = 'changed'
668+
child_parts.append(list(tool_context.state['temp:parts']))
669+
return 'edited'
670+
671+
caller_deltas = []
672+
caller_parts = []
673+
674+
def record_caller_state(tool, args, tool_context, tool_response):
675+
caller_deltas.append(dict(tool_context.actions.state_delta))
676+
caller_parts.append(tool_context.state.get('temp:parts'))
677+
678+
mock_model = testing_utils.MockModel.create(
679+
responses=[
680+
function_call_no_schema,
681+
Part.from_function_call(name='edit_temp_state', args={}),
682+
'response1',
683+
'response2',
684+
]
685+
)
686+
tool_agent = Agent(
687+
name='tool_agent', model=mock_model, tools=[edit_temp_state]
688+
)
689+
root_agent = Agent(
690+
name='root_agent',
691+
model=mock_model,
692+
before_agent_callback=set_caller_state,
693+
after_tool_callback=record_caller_state,
694+
tools=[AgentTool(agent=tool_agent)],
695+
)
696+
697+
runner = testing_utils.InMemoryRunner(root_agent)
698+
runner.run('test1')
699+
700+
assert child_parts == [['child_part']]
701+
assert len(caller_deltas) == 1
702+
assert not any(key.startswith('temp:') for key in caller_deltas[0])
703+
assert caller_parts == [[]]
704+
705+
560706
@mark.asyncio
561707
async def test_update_artifacts():
562708
"""The agent tool can read and write artifacts."""

0 commit comments

Comments
 (0)