1717from __future__ import annotations
1818
1919import asyncio
20+ import logging
2021import sys
2122from typing import AsyncGenerator
2223from typing import ClassVar
24+ from typing import Optional
2325
2426from typing_extensions import deprecated
2527from typing_extensions import override
3234from .invocation_context import InvocationContext
3335from .parallel_agent_config import ParallelAgentConfig
3436
37+ logger = logging .getLogger ('google_adk.' + __name__ )
38+
3539
3640def _create_branch_ctx_for_sub_agent (
3741 agent : BaseAgent ,
@@ -49,29 +53,75 @@ def _create_branch_ctx_for_sub_agent(
4953 return invocation_context
5054
5155
56+ async def _iter_with_idle_timeout (
57+ gen : AsyncGenerator [Event , None ],
58+ timeout_secs : float ,
59+ branch_name : str ,
60+ ) -> AsyncGenerator [Event , None ]:
61+ """Wrap *gen*, raising TimeoutError if no event arrives within *timeout_secs*.
62+
63+ Uses asyncio.wait_for on each __anext__ call so that a branch whose upstream
64+ model stream silently stalls (connection open, no chunks) is detected and
65+ cancelled rather than hanging the parent ParallelAgent indefinitely.
66+ """
67+ while True :
68+ try :
69+ event = await asyncio .wait_for (
70+ gen .__anext__ (),
71+ timeout = timeout_secs ,
72+ )
73+ except StopAsyncIteration :
74+ return
75+ except asyncio .TimeoutError as exc :
76+ logger .warning (
77+ 'ParallelAgent branch %r has not produced an event for %.1fs. '
78+ 'The upstream model stream may have stalled. Cancelling the branch '
79+ 'so the parent agent is not blocked indefinitely.' ,
80+ branch_name ,
81+ timeout_secs ,
82+ )
83+ raise asyncio .TimeoutError (
84+ f'Branch { branch_name !r} idle for >{ timeout_secs } s without an event'
85+ ) from exc
86+ yield event
87+
88+
5289async def _merge_agent_run (
5390 agent_runs : list [AsyncGenerator [Event , None ]],
91+ * ,
92+ branch_names : list [str ] | None = None ,
93+ branch_idle_timeout_secs : float | None = None ,
5494) -> AsyncGenerator [Event , None ]:
5595 """Merges agent runs using asyncio.TaskGroup on Python 3.11+."""
5696 sentinel = object ()
57- queue = asyncio .Queue ()
97+ queue : asyncio .Queue = asyncio .Queue ()
98+ names = branch_names or [f'branch-{ i } ' for i in range (len (agent_runs ))]
5899
59100 # Agents are processed in parallel.
60101 # Events for each agent are put on queue sequentially.
61- async def process_an_agent (events_for_one_agent ):
102+ async def process_an_agent (events_for_one_agent , branch_name : str ):
62103 try :
63- async for event in events_for_one_agent :
104+ gen = (
105+ _iter_with_idle_timeout (
106+ events_for_one_agent , branch_idle_timeout_secs , branch_name
107+ )
108+ if branch_idle_timeout_secs is not None
109+ else events_for_one_agent
110+ )
111+ async for event in gen :
64112 resume_signal = asyncio .Event ()
65- await queue .put ((event , resume_signal ))
113+ # put_nowait: the queue is unbounded so this never blocks, and it is
114+ # safe to call from a finally block that may run during cancellation.
115+ queue .put_nowait ((event , resume_signal ))
66116 # Wait for upstream to consume event before generating new events.
67117 await resume_signal .wait ()
68118 finally :
69- # Mark agent as finished.
70- await queue .put ((sentinel , None ))
119+ # Mark agent as finished. put_nowait is cancellation-safe (see above).
120+ queue .put_nowait ((sentinel , None ))
71121
72122 async with asyncio .TaskGroup () as tg :
73- for events_for_one_agent in agent_runs :
74- tg .create_task (process_an_agent (events_for_one_agent ))
123+ for events_for_one_agent , name in zip ( agent_runs , names ) :
124+ tg .create_task (process_an_agent (events_for_one_agent , name ))
75125
76126 sentinel_count = 0
77127 # Run until all agents finished processing.
@@ -89,6 +139,9 @@ async def process_an_agent(events_for_one_agent):
89139# TODO - remove once Python <3.11 is no longer supported.
90140async def _merge_agent_run_pre_3_11 (
91141 agent_runs : list [AsyncGenerator [Event , None ]],
142+ * ,
143+ branch_names : list [str ] | None = None ,
144+ branch_idle_timeout_secs : float | None = None ,
92145) -> AsyncGenerator [Event , None ]:
93146 """Merges agent runs for Python 3.10 without asyncio.TaskGroup.
94147
@@ -97,12 +150,16 @@ async def _merge_agent_run_pre_3_11(
97150
98151 Args:
99152 agent_runs: Async generators that yield events from each agent.
153+ branch_names: Optional names for each branch, used in log messages.
154+ branch_idle_timeout_secs: If set, cancel a branch that produces no event
155+ for this many seconds (guards against silently stalled model streams).
100156
101157 Yields:
102158 Event: The next event from the merged generator.
103159 """
104160 sentinel = object ()
105- queue = asyncio .Queue ()
161+ queue : asyncio .Queue = asyncio .Queue ()
162+ names = branch_names or [f'branch-{ i } ' for i in range (len (agent_runs ))]
106163
107164 def propagate_exceptions (tasks ):
108165 # Propagate exceptions and errors from tasks.
@@ -114,21 +171,31 @@ def propagate_exceptions(tasks):
114171
115172 # Agents are processed in parallel.
116173 # Events for each agent are put on queue sequentially.
117- async def process_an_agent (events_for_one_agent ):
174+ async def process_an_agent (events_for_one_agent , branch_name : str ):
118175 try :
119- async for event in events_for_one_agent :
176+ gen = (
177+ _iter_with_idle_timeout (
178+ events_for_one_agent , branch_idle_timeout_secs , branch_name
179+ )
180+ if branch_idle_timeout_secs is not None
181+ else events_for_one_agent
182+ )
183+ async for event in gen :
120184 resume_signal = asyncio .Event ()
121- await queue .put ((event , resume_signal ))
185+ queue .put_nowait ((event , resume_signal ))
122186 # Wait for upstream to consume event before generating new events.
123187 await resume_signal .wait ()
124188 finally :
125- # Mark agent as finished.
126- await queue .put ((sentinel , None ))
189+ # put_nowait is cancellation-safe: the queue is unbounded so it never
190+ # blocks, and it will not raise even if the task is being cancelled.
191+ queue .put_nowait ((sentinel , None ))
127192
128193 tasks = []
129194 try :
130- for events_for_one_agent in agent_runs :
131- tasks .append (asyncio .create_task (process_an_agent (events_for_one_agent )))
195+ for events_for_one_agent , name in zip (agent_runs , names ):
196+ tasks .append (
197+ asyncio .create_task (process_an_agent (events_for_one_agent , name ))
198+ )
132199
133200 sentinel_count = 0
134201 # Run until all agents finished processing.
@@ -143,6 +210,10 @@ async def process_an_agent(events_for_one_agent):
143210 # Signal to agent that event has been processed by runner and it can
144211 # continue now.
145212 resume_signal .set ()
213+ # A task may have put its sentinel AND raised an exception; if the sentinel
214+ # was consumed before propagate_exceptions ran in the loop body, the error
215+ # would be silently lost. Check once more after the loop to surface it.
216+ propagate_exceptions (tasks )
146217 finally :
147218 for task in tasks :
148219 task .cancel ()
@@ -165,6 +236,15 @@ class ParallelAgent(BaseAgent):
165236 .. deprecated::
166237 ParallelAgent is deprecated and will be removed in future versions.
167238 Please use Workflow instead.
239+
240+ Attributes:
241+ branch_idle_timeout_secs: Optional per-branch idle timeout in seconds.
242+ When set, any branch that produces no event for this many seconds is
243+ cancelled and raises ``asyncio.TimeoutError``, which unblocks the
244+ parent agent instead of hanging indefinitely. This guards against
245+ upstream model streams that stall silently (connection open, no
246+ chunks arriving). ``None`` (the default) disables the timeout and
247+ preserves the original unbounded-wait behaviour.
168248 """
169249
170250 config_type : ClassVar [type [BaseAgentConfig ]] = ParallelAgentConfig
@@ -174,6 +254,9 @@ class ParallelAgent(BaseAgent):
174254 version, along with the AgentConfig YAML loader.
175255 """
176256
257+ branch_idle_timeout_secs : Optional [float ] = None
258+ """Per-branch idle timeout in seconds; None disables the guard."""
259+
177260 @override
178261 async def _run_async_impl (
179262 self , ctx : InvocationContext
@@ -187,13 +270,15 @@ async def _run_async_impl(
187270 yield self ._create_agent_state_event (ctx )
188271
189272 agent_runs = []
273+ branch_names = []
190274 # Prepare and collect async generators for each sub-agent.
191275 for sub_agent in self .sub_agents :
192276 sub_agent_ctx = _create_branch_ctx_for_sub_agent (self , sub_agent , ctx )
193277
194278 # Only include sub-agents that haven't finished in a previous run.
195279 if not sub_agent_ctx .end_of_agents .get (sub_agent .name ):
196280 agent_runs .append (sub_agent .run_async (sub_agent_ctx ))
281+ branch_names .append (sub_agent .name )
197282
198283 pause_invocation = False
199284 try :
@@ -202,7 +287,13 @@ async def _run_async_impl(
202287 if sys .version_info >= (3 , 11 )
203288 else _merge_agent_run_pre_3_11
204289 )
205- async with Aclosing (merge_func (agent_runs )) as agen :
290+ async with Aclosing (
291+ merge_func (
292+ agent_runs ,
293+ branch_names = branch_names ,
294+ branch_idle_timeout_secs = self .branch_idle_timeout_secs ,
295+ )
296+ ) as agen :
206297 async for event in agen :
207298 yield event
208299 if ctx .should_pause_invocation (event ):
0 commit comments