@@ -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
561707async def test_update_artifacts ():
562708 """The agent tool can read and write artifacts."""
0 commit comments