Skip to content

Commit 4e5b02a

Browse files
committed
fix(agents): add branch_idle_timeout_secs to ParallelAgent to detect silent stalls
When an upstream model stream (LiteLLM/vLLM or any backend) silently stops producing chunks without closing the connection, process_an_agent in _merge_agent_run blocks forever on __anext__() with no exception or EOF. The queue.get() in the merge loop then also waits indefinitely, leaving the /run_sse HTTP stream open until an external gateway times out (issue #5455). Root cause: there is no per-chunk idle timeout anywhere in the parallel branch merge loop, so a silently stalled model connection is indistinguishable from a legitimately slow one. Fix: add _iter_with_idle_timeout(), which wraps each __anext__() call with asyncio.wait_for(). If no event arrives within branch_idle_timeout_secs, it logs a warning and raises asyncio.TimeoutError, unblocking both the branch task and the merge loop so the SSE stream can close with an explicit error rather than hanging. Additional fix (pre-3.11 path): _merge_agent_run_pre_3_11 could silently swallow a branch exception when the branch put its sentinel into the queue before propagate_exceptions() ran. Added a final propagate_exceptions() call after the while loop exits to catch that race. Also: changed await queue.put(sentinel) → queue.put_nowait(sentinel) in all finally blocks. The queue is unbounded so put_nowait never blocks, and it cannot be interrupted by CancelledError during task cancellation. Usage: ParallelAgent( name='recommend', sub_agents=[workout_agent, sleep_agent], branch_idle_timeout_secs=120.0, # raise if no event for 2 min ) Fixes #5455
1 parent d3c21d7 commit 4e5b02a

2 files changed

Lines changed: 540 additions & 17 deletions

File tree

src/google/adk/agents/parallel_agent.py

Lines changed: 108 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,11 @@
1717
from __future__ import annotations
1818

1919
import asyncio
20+
import logging
2021
import sys
2122
from typing import AsyncGenerator
2223
from typing import ClassVar
24+
from typing import Optional
2325

2426
from typing_extensions import deprecated
2527
from typing_extensions import override
@@ -32,6 +34,8 @@
3234
from .invocation_context import InvocationContext
3335
from .parallel_agent_config import ParallelAgentConfig
3436

37+
logger = logging.getLogger('google_adk.' + __name__)
38+
3539

3640
def _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+
5289
async 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.
90140
async 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

Comments
 (0)