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

src/google/adk/agents/invocation_context.py

@@ -249,6 +249,25 @@ class InvocationContext(BaseModel):
       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 @@ class LoopAgent(BaseAgent):
           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 @@ class LoopAgent(BaseAgent):
       # 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 @@ class ParallelAgent(BaseAgent):
     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 @@ class SequentialAgent(BaseAgent):
           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 @@ class TestInvocationContextWithAppResumablity:
     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."""