fix: Resolve scheduler leakage and make scheduler instantiation explicit

- Remove implicit on-demand fabrication of DynamicNodeScheduler during child Context derivation to avoid mutating parent context state during execution.
- Always instantiate and set the DynamicNodeScheduler at the root level in the runner.
- Clear `_workflow_scheduler` on execution exit (in both runner and workflow loops) to prevent scheduler lifetime leakage.
- Update tests to explicitly set `rerun_on_resume = True` where resumption of mock nodes is expected.
- Always rerun Workflow nodes on resume to allow their internal loops to correctly drive resumption of child nodes.
- Use `weakref` for `_workflow_scheduler` in `Context` to prevent reference cycle between scheduler, tasks, and contexts.

PiperOrigin-RevId: 945945253

Author: google-genai-bot

src/google/adk/agents/context.py

@@ -20,7 +20,6 @@
 from collections.abc import Sequence
 from typing import Any
 from typing import TYPE_CHECKING
-import weakref
 
 from opentelemetry import context as context_api
 from typing_extensions import override
@@ -56,7 +55,13 @@ def _derive_scheduler(
 ) -> ScheduleDynamicNode | None:
   """Derives the dynamic node scheduler from the parent context."""
   if parent_ctx:
-    return parent_ctx._workflow_scheduler
+    scheduler = parent_ctx._workflow_scheduler
+    if scheduler is None:
+      from ..workflow._dynamic_node_scheduler import DynamicNodeScheduler
+      from ..workflow._dynamic_node_scheduler import DynamicNodeState
+
+      scheduler = DynamicNodeScheduler(state=DynamicNodeState())
+    return scheduler
   return None
 
 
@@ -117,8 +122,6 @@ class Context(ReadonlyContext):
   fields`` section are available.
   """
 
-  _workflow_scheduler_ref: weakref.ref[ScheduleDynamicNode] | None = None
-
   def __init__(
       self,
       invocation_context: InvocationContext,
@@ -200,10 +203,7 @@ def __init__(
         node=node,
     )
     self._resume_inputs = resume_inputs or {}
-    scheduler = _derive_scheduler(parent_ctx)
-    self._workflow_scheduler_ref = (
-        weakref.ref(scheduler) if scheduler is not None else None
-    )
+    self._workflow_scheduler = _derive_scheduler(parent_ctx)
     self._node_rerun_on_resume = node.rerun_on_resume if node else True
     self._child_run_counters: dict[str, int] = {}
     self._attempt_count = attempt_count
@@ -264,19 +264,6 @@ def isolation_scope(self) -> str | None:
   def isolation_scope(self, value: str | None) -> None:
     self._isolation_scope = value
 
-  @property
-  def _workflow_scheduler(self) -> ScheduleDynamicNode | None:
-    """The workflow scheduler associated with this context."""
-    if self._workflow_scheduler_ref is not None:
-      return self._workflow_scheduler_ref()
-    return None
-
-  @_workflow_scheduler.setter
-  def _workflow_scheduler(self, value: ScheduleDynamicNode | None) -> None:
-    self._workflow_scheduler_ref = (
-        weakref.ref(value) if value is not None else None
-    )
-
   @property
   def tool_confirmation(self) -> ToolConfirmation | None:
     """The tool confirmation of the current tool call."""

src/google/adk/runners.py

@@ -573,6 +573,11 @@ async def _run_node_async(
 
       root_ctx = Context(ic)
       root_agent = node or self.agent
+      is_agent = isinstance(self.agent, BaseAgent)
+      has_sub_agents = is_agent and bool(
+          getattr(self.agent, 'sub_agents', None)
+      )
+      use_scheduler = is_agent and has_sub_agents
 
       # The root chat coordinator's isolation_scope stays None: its own
       # events (FCs, text, synthesized FRs from completed task
@@ -586,10 +591,11 @@ async def _run_node_async(
 
       async def _drive_root_node():
         try:
-          # Rehydration warning: DynamicNodeScheduler relies on session.events scanning.
-          # Stateful live EUC/LRO streams may rehydrate freshly if not yet persisted.
-          scheduler = DynamicNodeScheduler(state=_LoopState())
-          root_ctx._workflow_scheduler = scheduler
+          if use_scheduler:
+            # Rehydration warning: DynamicNodeScheduler relies on session.events scanning.
+            # Stateful live EUC/LRO streams may rehydrate freshly if not yet persisted.
+            scheduler = DynamicNodeScheduler(state=_LoopState())
+            root_ctx._workflow_scheduler = scheduler
 
           try:
             await root_ctx._run_node_internal(
@@ -603,7 +609,6 @@ async def _drive_root_node():
           except DynamicNodeFailError as e:
             raise e.error
         finally:
-          root_ctx._workflow_scheduler = None
           await ic._event_queue.put((done_sentinel, None))
 
       task = asyncio.create_task(_drive_root_node())
@@ -674,6 +679,7 @@ async def _run_node_live(
     from .workflow._errors import DynamicNodeFailError
     from .workflow._errors import NodeInterruptedError
     from .workflow._workflow import _LoopState
+    from .workflow._workflow import Workflow
 
     ic = self._new_invocation_context_for_live(
         session,
@@ -684,12 +690,15 @@ async def _run_node_live(
 
     root_ctx = Context(ic)
     root_agent = self.agent
+    is_workflow = isinstance(root_agent, Workflow)
+
     done_sentinel = object()
 
     async def _drive_root_node():
       try:
-        scheduler = DynamicNodeScheduler(state=_LoopState())
-        root_ctx._workflow_scheduler = scheduler
+        if is_workflow:
+          scheduler = DynamicNodeScheduler(state=_LoopState())
+          root_ctx._workflow_scheduler = scheduler
 
         try:
           await root_ctx.run_node(
@@ -701,7 +710,6 @@ async def _drive_root_node():
         except DynamicNodeFailError as e:
           raise e.error
       finally:
-        root_ctx._workflow_scheduler = None
         await ic._event_queue.put((done_sentinel, None))
 
     task = asyncio.create_task(_drive_root_node())

src/google/adk/workflow/_workflow.py

@@ -245,7 +245,6 @@ async def _run_impl(
     try:
       await self._run_loop(loop_state, ctx)
     finally:
-      ctx._workflow_scheduler = None
       await self._cleanup_all_tasks(loop_state)
 
     if loop_state.error_shut_down:

src/google/adk/workflow/utils/_replay_interceptor.py

@@ -79,17 +79,10 @@ def check_interception(
             interrupts=set(current_run.state.interrupts),
         )
 
+  # Intercept executions based on historical session events (cross-turn replay).
   if not recovered:
     return InterceptionResult(should_run=True)
 
-  from .._workflow import Workflow
-
-  if isinstance(node, Workflow):
-    return InterceptionResult(
-        should_run=True,
-        resume_inputs=recovered.resolved_responses,
-    )
-
   unresolved = recovered.interrupt_ids - recovered.resolved_ids
 
   should_run = False

tests/unittests/agents/test_context.py

@@ -649,12 +649,13 @@ def test_derive_scheduler_with_parent_having_scheduler(self):
 
   def test_derive_scheduler_with_parent_no_scheduler(self):
     from google.adk.agents.context import _derive_scheduler
+    from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler
 
     mock_parent = MagicMock()
     mock_parent._workflow_scheduler = None
 
     scheduler = _derive_scheduler(mock_parent)
-    assert scheduler is None
+    assert isinstance(scheduler, DynamicNodeScheduler)
 
 
 class TestContextGetInvocationContext:

tests/unittests/runners/test_runner_node.py

@@ -456,7 +456,6 @@ async def test_standalone_node_resume():
   """A standalone node resumes with resume_inputs from function response."""
 
   class _Node(BaseNode):
-    rerun_on_resume: bool = True
 
     async def _run_impl(
         self, *, ctx: Context, node_input: Any
@@ -482,7 +481,6 @@ async def test_resume_preserves_original_user_content():
   """On resume, Runner passes the original text as node_input, not the FR."""
 
   class _Node(BaseNode):
-    rerun_on_resume: bool = True
 
     async def _run_impl(
         self, *, ctx: Context, node_input: Any
@@ -513,7 +511,6 @@ async def test_resume_populates_invocation_user_content():
   seen: list[Any] = []
 
   class _Node(BaseNode):
-    rerun_on_resume: bool = True
 
     async def _run_impl(
         self, *, ctx: Context, node_input: Any
@@ -540,7 +537,6 @@ async def test_resume_by_invocation_id_populates_user_content():
   seen: list[Any] = []
 
   class _Node(BaseNode):
-    rerun_on_resume: bool = True
 
     async def _run_impl(
         self, *, ctx: Context, node_input: Any
@@ -568,10 +564,7 @@ async def _run_impl(
   invocation_id = updated.events[0].invocation_id
 
   async for _ in runner.run_async(
-      user_id='u',
-      session_id=session.id,
-      invocation_id=invocation_id,
-      new_message=_make_resume_message(fc_name='tool', response={'v': 1}),
+      user_id='u', session_id=session.id, invocation_id=invocation_id
   ):
     pass