chore: Delegate the agent state reset logic to LoopAgent

This is so we don't need to worry about side effect of Loop in all agent type. Custom agent should do the same if there exists loop inside.

PiperOrigin-RevId: 818766305

Author: DeanChensj

src/google/adk/agents/invocation_context.py

@@ -249,6 +249,25 @@ def set_agent_state(
       self.end_of_agents.pop(agent_name, None)
       self.agent_states.pop(agent_name, None)
 
+  def reset_sub_agent_states(
+      self,
+      agent_name: str,
+  ) -> None:
+    """Resets the state of all sub-agents of the given agent in this invocation.
+
+    Args:
+      agent_name: The name of the agent whose sub-agent states need to be reset.
+    """
+    agent = self.agent.find_agent(agent_name)
+    if not agent:
+      return
+
+    for sub_agent in agent.sub_agents:
+      # Reset the sub-agent's state in the context to ensure that each
+      # sub-agent starts fresh.
+      self.set_agent_state(sub_agent.name)
+      self.reset_sub_agent_states(sub_agent.name)
+
   def populate_invocation_agent_states(self) -> None:
     """Populates agent states for the current invocation if it is resumable.
 

src/google/adk/agents/loop_agent.py

@@ -94,10 +94,6 @@ async def _run_async_impl(
           ctx.set_agent_state(self.name, agent_state=agent_state)
           yield self._create_agent_state_event(ctx)
 
-        # Reset the sub-agent's state in the context to ensure that each
-        # sub-agent starts fresh.
-        if not is_resuming_at_current_agent:
-          ctx.set_agent_state(sub_agent.name)
         is_resuming_at_current_agent = False
 
         async with Aclosing(sub_agent.run_async(ctx)) as agen:
@@ -114,6 +110,8 @@ async def _run_async_impl(
       # Restart from the beginning of the loop.
       start_index = 0
       times_looped += 1
+      # Reset the state of all sub-agents in the loop.
+      ctx.reset_sub_agent_states(self.name)
 
     # If the invocation is paused, we should not yield the end of agent event.
     if pause_invocation:

src/google/adk/agents/parallel_agent.py

@@ -187,10 +187,6 @@ async def _run_async_impl(
     agent_runs = []
     # Prepare and collect async generators for each sub-agent.
     for sub_agent in self.sub_agents:
-      if agent_state is None:
-        # Reset sub-agent state to make sure each sub-agent starts fresh.
-        ctx.set_agent_state(sub_agent.name)
-
       sub_agent_ctx = _create_branch_ctx_for_sub_agent(self, sub_agent, ctx)
 
       # Only include sub-agents that haven't finished in a previous run.

src/google/adk/agents/sequential_agent.py

@@ -73,10 +73,6 @@ async def _run_async_impl(
           ctx.set_agent_state(self.name, agent_state=agent_state)
           yield self._create_agent_state_event(ctx)
 
-        # Reset the sub-agent's state in the context to ensure that each
-        # sub-agent starts fresh.
-        ctx.set_agent_state(sub_agent.name)
-
       async with Aclosing(sub_agent.run_async(ctx)) as agen:
         async for event in agen:
           yield event

tests/unittests/agents/test_invocation_context.py

@@ -390,6 +390,40 @@ def test_reset_agent_state(self):
     assert 'agent1' not in invocation_context.agent_states
     assert 'agent1' not in invocation_context.end_of_agents
 
+  def test_reset_sub_agent_states(self):
+    """Tests that reset_sub_agent_states resets sub-agent states."""
+    sub_sub_agent_1 = BaseAgent(name='sub_sub_agent_1')
+    sub_agent_1 = BaseAgent(name='sub_agent_1', sub_agents=[sub_sub_agent_1])
+    sub_agent_2 = BaseAgent(name='sub_agent_2')
+    root_agent = BaseAgent(
+        name='root_agent', sub_agents=[sub_agent_1, sub_agent_2]
+    )
+
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    invocation_context.agent = root_agent
+    invocation_context.set_agent_state(
+        'sub_agent_1', agent_state=BaseAgentState()
+    )
+    invocation_context.set_agent_state('sub_agent_2', end_of_agent=True)
+    invocation_context.set_agent_state(
+        'sub_sub_agent_1', agent_state=BaseAgentState()
+    )
+
+    assert 'sub_agent_1' in invocation_context.agent_states
+    assert 'sub_agent_2' in invocation_context.end_of_agents
+    assert 'sub_sub_agent_1' in invocation_context.agent_states
+
+    invocation_context.reset_sub_agent_states('root_agent')
+
+    assert 'sub_agent_1' not in invocation_context.agent_states
+    assert 'sub_agent_1' not in invocation_context.end_of_agents
+    assert 'sub_agent_2' not in invocation_context.agent_states
+    assert 'sub_agent_2' not in invocation_context.end_of_agents
+    assert 'sub_sub_agent_1' not in invocation_context.agent_states
+    assert 'sub_sub_agent_1' not in invocation_context.end_of_agents
+
 
 class TestFindMatchingFunctionCall:
   """Test suite for find_matching_function_call."""