@samitouri / QOSami_Agents_Zero-2 / commits / f6879249

Fix parallel await timeouts

Keep running parallel jobs alive when an await call reaches its timeout so agents can await the same job ids again instead of cancelling child work. Distinguish direct background tool workers from call_subordinate child chats so nested subordinate chats can use parallel normally while true worker recursion remains blocked. Update the parallel prompt, DOX notes, and regressions for non-destructive timeout and non-blocking collect semantics.

Alessandro committed Jun 13, 2026 at 18:36 UTC f68792496b38b8a0d7a5f38f833edb602bbdd4fb
6 files changed +205 -21
helpers/parallel_tools.py
+37 -13
@@ -20,6 +20,7 @@ if TYPE_CHECKING:
20 PARALLEL_JOBS_KEY = "_parallel_jobs"
21 PARALLEL_WORKER_PARENT_CONTEXT_KEY = "_parallel_parent_context_id"
22 PARALLEL_WORKER_JOB_KEY = "_parallel_job_id"
23 +PARALLEL_WORKER_KIND_KEY = "_parallel_worker_kind"
24
25 CHILD_PARENT_CONTEXT_ID_KEY = "parent_context_id"
26 CHILD_PARENT_CONTEXT_KIND_KEY = "parent_context_kind"
@@ -139,11 +140,20 @@ def coerce_timeout(value: Any) -> int:
140 return timeout
141
142
142 -def is_parallel_worker(agent: "Agent | None") -> bool:
143 +def _parallel_worker_kind(agent: "Agent | None") -> JobKind | None:
144 context = getattr(agent, "context", None)
145 if not context:
145 - return False
146 - return bool(context.get_data(PARALLEL_WORKER_JOB_KEY))
146 + return None
147 + kind = context.get_data(PARALLEL_WORKER_KIND_KEY)
148 + if kind in {"tool", "subordinate"}:
149 + return kind
150 + if context.get_data(PARALLEL_WORKER_JOB_KEY):
151 + return "tool"
152 + return None
153 +
154 +
155 +def is_parallel_worker(agent: "Agent | None") -> bool:
156 + return _parallel_worker_kind(agent) == "tool"
157
158
159 def _jobs_for_context(context: "AgentContext") -> dict[str, ParallelJob]:
@@ -209,12 +219,14 @@ async def await_parallel_jobs(
219 timeout: int = DEFAULT_TIMEOUT_SECONDS,
220 *,
221 collect: bool = True,
222 + wait: bool = True,
223 ) -> list[dict[str, Any]]:
224 if not job_ids:
225 raise ValueError("No `job_ids` were provided to await.")
226
227 deadline = time.time() + timeout
228 known_job_ids = set(job_ids)
229 + wait_timed_out_job_ids: set[str] = set()
230 while True:
231 await refresh_parallel_jobs(agent)
232 jobs = [_jobs_for_context(agent.context).get(job_id) for job_id in job_ids]
@@ -223,12 +235,11 @@ async def await_parallel_jobs(
235 raise ValueError(f"Unknown parallel job id(s): {', '.join(missing)}")
236
237 active = [job for job in jobs if job and job.state not in TERMINAL_STATES]
226 - if not active:
238 + if not wait or not active:
239 break
240
241 if time.time() >= deadline:
230 - for job in active:
231 - await _timeout_job(job)
242 + wait_timed_out_job_ids = {job.id for job in active}
243 break
244
245 await asyncio.sleep(POLL_INTERVAL_SECONDS)
@@ -237,7 +248,10 @@ async def await_parallel_jobs(
248 for job_id in job_ids:
249 job = _jobs_for_context(agent.context).get(job_id)
250 if job:
240 - snapshots.append(_job_snapshot(job, include_result=True))
251 + snapshot = _job_snapshot(job, include_result=True)
252 + if job.id in wait_timed_out_job_ids and job.state not in TERMINAL_STATES:
253 + snapshot["wait_timed_out"] = True
254 + snapshots.append(snapshot)
255
256 if collect:
257 for job_id in known_job_ids:
@@ -338,8 +352,14 @@ def format_started_jobs(jobs: list[ParallelJob]) -> str:
352
353 def format_parallel_results(results: list[dict[str, Any]]) -> str:
354 states = [result.get("state") for result in results]
341 - if states and all(state == "success" for state in states):
355 + has_active_jobs = any(state not in TERMINAL_STATES for state in states)
356 + wait_timed_out = any(result.get("wait_timed_out") for result in results)
357 + if has_active_jobs:
358 + status = "waiting" if wait_timed_out else "running"
359 + elif states and all(state == "success" for state in states):
360 status = "success"
361 + elif states and all(state == "cancelled" for state in states):
362 + status = "cancelled"
363 elif any(state == "success" for state in states):
364 status = "partial"
365 else:
@@ -349,6 +369,13 @@ def format_parallel_results(results: list[dict[str, Any]]) -> str:
369 "status": status,
370 "jobs": results,
371 }
372 + if wait_timed_out:
373 + payload["wait_timeout"] = True
374 + if has_active_jobs:
375 + payload["instruction"] = (
376 + "Some jobs are still running. Call `parallel` with `action: \"await\"` "
377 + "and the listed `job_ids` to wait again, or `action: \"cancel\"` to stop them."
378 + )
379 return json.dumps(payload, indent=2, ensure_ascii=False)
380
381
@@ -399,6 +426,7 @@ async def _run_subordinate_context_job(parent_context_id: str, job: ParallelJob)
426
427 worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context.id)
428 worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id)
429 + worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind)
430 worker_context.set_output_data(CHILD_PARENT_CONTEXT_ID_KEY, parent_context.id)
431 worker_context.set_output_data(CHILD_PARENT_CONTEXT_KIND_KEY, "parallel")
432 worker_context.set_output_data(CHILD_PARENT_CONTEXT_LABEL_KEY, child_name)
@@ -441,6 +469,7 @@ async def _run_direct_tool_job(parent_context_id: str, job: ParallelJob) -> str:
469 )
470 worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context_id)
471 worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id)
472 + worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind)
473 job.worker_context_id = worker_context.id
474 _copy_project(parent_context, worker_context)
475
@@ -503,10 +532,6 @@ async def execute_tool_call(agent: "Agent", tool_name: str, tool_args: dict[str,
532 agent.loop_data.current_tool = None
533
534
506 -async def _timeout_job(job: ParallelJob) -> None:
507 - await _cancel_job(job, state="timeout", message="Parallel job timed out.")
508 -
509 -
535 async def _cancel_job(
536 job: ParallelJob,
537 *,
@@ -631,7 +656,6 @@ def _subordinate_worker_system_prompt(profile: str) -> str:
656 lines = [
657 "You are running as an isolated parallel worker for a parent Agent Zero chat.",
658 "Return a concise final textual summary for the parent. Artifacts and files are supplementary, not a substitute for the textual result.",
634 - "Do not call the `parallel` tool from this worker.",
659 ]
660 if profile:
661 lines.append(f"Act with the `{profile}` profile's expertise and priorities.")
helpers/parallel_tools.py.dox.md
+2 -2
@@ -24,7 +24,7 @@
24 - Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together.
25 - Wrapped tool-call items must use the same shape as normal tool calls: a tool name plus arguments.
26 - Normalization accepts full agent-reply-shaped objects when `tool_name` and `tool_args` are present; non-contract planning fields such as `thoughts` or `headline` are ignored.
27 -- `call_subordinate` jobs run in isolated child chat contexts tagged with parent-chat metadata; they must not be added to the scheduler task list.
27 +- `call_subordinate` jobs run in isolated child chat contexts tagged with parent-chat metadata; they must not be added to the scheduler task list and may use normal child-chat tools, including `parallel`.
28 - Direct tool jobs run in isolated background contexts and are blocked from recursively invoking `parallel`.
29 - Parent-visible child log items are created for each wrapped call so the WebUI can inspect concurrent children separately while the wrapper result remains model-history-only.
30 - Child tool logs mirror normal tool-call visible args; job ids remain available through wrapper results and prompt extras rather than visible process-step args.
@@ -34,7 +34,7 @@
34 ## Key Concepts
35
36 - The parent context stores in-flight jobs under a private data key; collected terminal jobs are removed from that registry.
37 -- `wait=True` starts jobs and awaits them before returning; `wait=False` returns job IDs immediately.
37 +- `wait=True` starts jobs and awaits them before returning until all requested jobs finish or the wait timeout is reached; the timeout stops waiting but does not cancel running jobs.
38 - `collect` returns already-finished job results without waiting; `await` waits for requested job IDs.
39 - Canceled jobs should be marked terminal and should stop their background `DeferredTask` when cancellation is possible.
40
prompts/agent.system.tool.parallel.md
+1
@@ -11,6 +11,7 @@ Rules:
11 - `call_subordinate` inside `parallel` starts an isolated child chat under the parent chat, not a scheduler task
12 - use `wait: false` only when you will collect results later with `job_ids`
13 - if extras list running or ready parallel jobs, collect them before final synthesis
14 +- `timeout` only limits how long this call waits; running jobs continue and can be awaited again by `job_ids`
15
16 Args: `tool_calls`, `job_ids`, `wait` default `true`, `action` as `start|await|collect|cancel`, `timeout`.
17
tests/test_parallel_tool.py
+146
@@ -1,5 +1,6 @@
1 from __future__ import annotations
2
3 +import json
4 import time
5 import sys
6 from types import SimpleNamespace
@@ -61,6 +62,27 @@ class _FakeAgent:
62 self.agent_name = "A0"
63
64
65 +class _FakeDeferredTask:
66 + def __init__(self, *, ready: bool = False, alive: bool = True, result=None) -> None:
67 + self.ready = ready
68 + self.alive = alive
69 + self._result = result
70 + self.killed = 0
71 +
72 + def is_ready(self):
73 + return self.ready
74 +
75 + def is_alive(self):
76 + return self.alive
77 +
78 + async def result(self):
79 + return self._result
80 +
81 + def kill(self):
82 + self.killed += 1
83 + self.alive = False
84 +
85 +
86 def test_normalize_parallel_tool_calls_accepts_normal_tool_request_shapes() -> None:
87 calls = parallel_tools.normalize_parallel_tool_calls(
88 [
@@ -125,6 +147,130 @@ async def test_parallel_jobs_extras_lists_running_and_ready_jobs() -> None:
147 assert "ready to collect" in extras
148
149
150 +@pytest.mark.asyncio
151 +async def test_parallel_await_timeout_keeps_running_jobs_awaitable(monkeypatch) -> None:
152 + agent = _FakeAgent()
153 + task = _FakeDeferredTask(alive=True)
154 + job = parallel_tools.ParallelJob(
155 + id="wait-1234abcd",
156 + parent_context_id="ctx",
157 + index=0,
158 + tool_name="wait",
159 + tool_args={"seconds": 60},
160 + kind="tool",
161 + state="running",
162 + created_at=99.0,
163 + started_at=99.0,
164 + deferred_task=task, # type: ignore[arg-type]
165 + )
166 + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
167 + times = iter([100.0, 102.0])
168 + monkeypatch.setattr(
169 + parallel_tools.time,
170 + "time",
171 + lambda: next(times, 102.0),
172 + )
173 +
174 + results = await parallel_tools.await_parallel_jobs( # type: ignore[arg-type]
175 + agent,
176 + [job.id],
177 + timeout=1,
178 + collect=True,
179 + wait=True,
180 + )
181 + payload = json.loads(parallel_tools.format_parallel_results(results))
182 +
183 + assert results[0]["state"] == "running"
184 + assert results[0]["wait_timed_out"] is True
185 + assert payload["status"] == "waiting"
186 + assert payload["wait_timeout"] is True
187 + assert "await" in payload["instruction"]
188 + assert task.killed == 0
189 + assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
190 +
191 +
192 +@pytest.mark.asyncio
193 +async def test_parallel_collect_returns_running_jobs_without_waiting_or_canceling() -> None:
194 + from tools.parallel import ParallelTool
195 +
196 + agent = _FakeAgent()
197 + task = _FakeDeferredTask(alive=True)
198 + job = parallel_tools.ParallelJob(
199 + id="wait-collect",
200 + parent_context_id="ctx",
201 + index=0,
202 + tool_name="wait",
203 + tool_args={"seconds": 60},
204 + kind="tool",
205 + state="running",
206 + deferred_task=task, # type: ignore[arg-type]
207 + )
208 + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
209 + tool = ParallelTool(
210 + agent, # type: ignore[arg-type]
211 + "parallel",
212 + None,
213 + {"action": "collect", "job_ids": [job.id]},
214 + "",
215 + None,
216 + )
217 +
218 + response = await tool.execute(**tool.args)
219 + payload = json.loads(response.message)
220 +
221 + assert payload["status"] == "running"
222 + assert payload["jobs"][0]["job_id"] == job.id
223 + assert "wait_timeout" not in payload
224 + assert task.killed == 0
225 + assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
226 +
227 +
228 +@pytest.mark.asyncio
229 +async def test_parallel_cancel_still_stops_and_removes_running_jobs() -> None:
230 + agent = _FakeAgent()
231 + task = _FakeDeferredTask(alive=True)
232 + job = parallel_tools.ParallelJob(
233 + id="wait-cancel",
234 + parent_context_id="ctx",
235 + index=0,
236 + tool_name="wait",
237 + tool_args={"seconds": 60},
238 + kind="tool",
239 + state="running",
240 + deferred_task=task, # type: ignore[arg-type]
241 + )
242 + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
243 +
244 + results = await parallel_tools.cancel_parallel_jobs(agent, [job.id]) # type: ignore[arg-type]
245 + payload = json.loads(parallel_tools.format_parallel_results(results))
246 +
247 + assert results[0]["state"] == "cancelled"
248 + assert payload["status"] == "cancelled"
249 + assert task.killed == 1
250 + assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
251 +
252 +
253 +@pytest.mark.asyncio
254 +async def test_parallel_recursion_guard_allows_subordinate_children_but_blocks_tool_workers() -> None:
255 + from extensions.python.tool_execute_before._20_block_parallel_recursion import (
256 + BlockParallelRecursion,
257 + )
258 + from helpers.errors import RepairableException
259 +
260 + agent = _FakeAgent()
261 + agent.context.set_data(parallel_tools.PARALLEL_WORKER_JOB_KEY, "legacy-job")
262 + assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
263 +
264 + agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "subordinate")
265 + assert parallel_tools.is_parallel_worker(agent) is False # type: ignore[arg-type]
266 + await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
267 +
268 + agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "tool")
269 + assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
270 + with pytest.raises(RepairableException, match="cannot be used inside a parallel worker"):
271 + await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
272 +
273 +
274 @pytest.mark.asyncio
275 async def test_parallel_subordinate_jobs_are_visible_child_logs_not_scheduler_tasks(monkeypatch) -> None:
276 class FakeDeferredTask:
tools/parallel.py
+17 -4
@@ -48,19 +48,31 @@ class ParallelTool(Tool):
48 break_loop=False,
49 )
50
51 - wait_default = action not in {"start", "background"}
51 + wait_default = action not in {"start", "background", "collect"}
52 wait = parallel_tools.coerce_bool(args.get("wait"), wait_default)
53 - if action in {"await", "wait", "collect"}:
53 + if action in {"await", "wait"}:
54 wait = True
55
56 if not wait:
57 - if not started_jobs:
57 + if not started_jobs and not job_ids:
58 return Response(
59 message="Error: `wait: false` requires `tool_calls` to start new jobs.",
60 break_loop=False,
61 )
62 + if not job_ids:
63 + return Response(
64 + message=parallel_tools.format_started_jobs(started_jobs),
65 + break_loop=False,
66 + )
67 + results = await parallel_tools.await_parallel_jobs(
68 + self.agent,
69 + all_job_ids,
70 + timeout=timeout,
71 + collect=True,
72 + wait=False,
73 + )
74 return Response(
63 - message=parallel_tools.format_started_jobs(started_jobs),
75 + message=parallel_tools.format_parallel_results(results),
76 break_loop=False,
77 )
78
@@ -69,6 +81,7 @@ class ParallelTool(Tool):
81 all_job_ids,
82 timeout=timeout,
83 collect=True,
84 + wait=True,
85 )
86 return Response(
87 message=parallel_tools.format_parallel_results(results),
tools/parallel.py.dox.md
+2 -2
@@ -23,10 +23,10 @@
23 - The tool is intended for independent calls only; dependent operations remain sequential.
24 - Independent calls should share one batch even when they use different tools; split only for dependencies, ordering, shared mutable state, or parent-context state/tool-availability changes.
25 - `action="start"` starts calls and optionally waits according to `wait`.
26 -- `action="await"` waits for requested job IDs.
26 +- `action="await"` waits for requested job IDs until completion or `timeout`; timeout returns running job handles without canceling them.
27 - `action="collect"` returns completed job results without waiting.
28 - `action="cancel"` requests cancellation for requested job IDs.
29 -- Recursive use of `parallel` from inside a parallel worker is blocked before execution.
29 +- Recursive use of `parallel` from inside a direct background tool worker is blocked before execution; subordinate child chats started by `call_subordinate` can use normal child-chat tools, including `parallel`.
30 - The wrapper tool does not create its own visible process-step log; each wrapped child call owns the visible log row, and the wrapper result is recorded only in model history.
31
32 ## Key Concepts