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.

helpers/parallel_tools.py

@@ -20,6 +20,7 @@ if TYPE_CHECKING:
 PARALLEL_JOBS_KEY = "_parallel_jobs"
 PARALLEL_WORKER_PARENT_CONTEXT_KEY = "_parallel_parent_context_id"
 PARALLEL_WORKER_JOB_KEY = "_parallel_job_id"
+PARALLEL_WORKER_KIND_KEY = "_parallel_worker_kind"
 
 CHILD_PARENT_CONTEXT_ID_KEY = "parent_context_id"
 CHILD_PARENT_CONTEXT_KIND_KEY = "parent_context_kind"
@@ -139,11 +140,20 @@ def coerce_timeout(value: Any) -> int:
     return timeout
 
 
-def is_parallel_worker(agent: "Agent | None") -> bool:
+def _parallel_worker_kind(agent: "Agent | None") -> JobKind | None:
     context = getattr(agent, "context", None)
     if not context:
-        return False
-    return bool(context.get_data(PARALLEL_WORKER_JOB_KEY))
+        return None
+    kind = context.get_data(PARALLEL_WORKER_KIND_KEY)
+    if kind in {"tool", "subordinate"}:
+        return kind
+    if context.get_data(PARALLEL_WORKER_JOB_KEY):
+        return "tool"
+    return None
+
+
+def is_parallel_worker(agent: "Agent | None") -> bool:
+    return _parallel_worker_kind(agent) == "tool"
 
 
 def _jobs_for_context(context: "AgentContext") -> dict[str, ParallelJob]:
@@ -209,12 +219,14 @@ async def await_parallel_jobs(
     timeout: int = DEFAULT_TIMEOUT_SECONDS,
     *,
     collect: bool = True,
+    wait: bool = True,
 ) -> list[dict[str, Any]]:
     if not job_ids:
         raise ValueError("No `job_ids` were provided to await.")
 
     deadline = time.time() + timeout
     known_job_ids = set(job_ids)
+    wait_timed_out_job_ids: set[str] = set()
     while True:
         await refresh_parallel_jobs(agent)
         jobs = [_jobs_for_context(agent.context).get(job_id) for job_id in job_ids]
@@ -223,12 +235,11 @@ async def await_parallel_jobs(
             raise ValueError(f"Unknown parallel job id(s): {', '.join(missing)}")
 
         active = [job for job in jobs if job and job.state not in TERMINAL_STATES]
-        if not active:
+        if not wait or not active:
             break
 
         if time.time() >= deadline:
-            for job in active:
-                await _timeout_job(job)
+            wait_timed_out_job_ids = {job.id for job in active}
             break
 
         await asyncio.sleep(POLL_INTERVAL_SECONDS)
@@ -237,7 +248,10 @@ async def await_parallel_jobs(
     for job_id in job_ids:
         job = _jobs_for_context(agent.context).get(job_id)
         if job:
-            snapshots.append(_job_snapshot(job, include_result=True))
+            snapshot = _job_snapshot(job, include_result=True)
+            if job.id in wait_timed_out_job_ids and job.state not in TERMINAL_STATES:
+                snapshot["wait_timed_out"] = True
+            snapshots.append(snapshot)
 
     if collect:
         for job_id in known_job_ids:
@@ -338,8 +352,14 @@ def format_started_jobs(jobs: list[ParallelJob]) -> str:
 
 def format_parallel_results(results: list[dict[str, Any]]) -> str:
     states = [result.get("state") for result in results]
-    if states and all(state == "success" for state in states):
+    has_active_jobs = any(state not in TERMINAL_STATES for state in states)
+    wait_timed_out = any(result.get("wait_timed_out") for result in results)
+    if has_active_jobs:
+        status = "waiting" if wait_timed_out else "running"
+    elif states and all(state == "success" for state in states):
         status = "success"
+    elif states and all(state == "cancelled" for state in states):
+        status = "cancelled"
     elif any(state == "success" for state in states):
         status = "partial"
     else:
@@ -349,6 +369,13 @@ def format_parallel_results(results: list[dict[str, Any]]) -> str:
         "status": status,
         "jobs": results,
     }
+    if wait_timed_out:
+        payload["wait_timeout"] = True
+    if has_active_jobs:
+        payload["instruction"] = (
+            "Some jobs are still running. Call `parallel` with `action: \"await\"` "
+            "and the listed `job_ids` to wait again, or `action: \"cancel\"` to stop them."
+        )
     return json.dumps(payload, indent=2, ensure_ascii=False)
 
 
@@ -399,6 +426,7 @@ async def _run_subordinate_context_job(parent_context_id: str, job: ParallelJob)
 
     worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context.id)
     worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id)
+    worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind)
     worker_context.set_output_data(CHILD_PARENT_CONTEXT_ID_KEY, parent_context.id)
     worker_context.set_output_data(CHILD_PARENT_CONTEXT_KIND_KEY, "parallel")
     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:
         )
         worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context_id)
         worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id)
+        worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind)
         job.worker_context_id = worker_context.id
         _copy_project(parent_context, worker_context)
 
@@ -503,10 +532,6 @@ async def execute_tool_call(agent: "Agent", tool_name: str, tool_args: dict[str,
         agent.loop_data.current_tool = None
 
 
-async def _timeout_job(job: ParallelJob) -> None:
-    await _cancel_job(job, state="timeout", message="Parallel job timed out.")
-
-
 async def _cancel_job(
     job: ParallelJob,
     *,
@@ -631,7 +656,6 @@ def _subordinate_worker_system_prompt(profile: str) -> str:
     lines = [
         "You are running as an isolated parallel worker for a parent Agent Zero chat.",
         "Return a concise final textual summary for the parent. Artifacts and files are supplementary, not a substitute for the textual result.",
-        "Do not call the `parallel` tool from this worker.",
     ]
     if profile:
         lines.append(f"Act with the `{profile}` profile's expertise and priorities.")

helpers/parallel_tools.py.dox.md

@@ -24,7 +24,7 @@
 - Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together.
 - Wrapped tool-call items must use the same shape as normal tool calls: a tool name plus arguments.
 - 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.
-- `call_subordinate` jobs run in isolated child chat contexts tagged with parent-chat metadata; they must not be added to the scheduler task list.
+- `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`.
 - Direct tool jobs run in isolated background contexts and are blocked from recursively invoking `parallel`.
 - 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.
 - 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 @@
 ## Key Concepts
 
 - The parent context stores in-flight jobs under a private data key; collected terminal jobs are removed from that registry.
-- `wait=True` starts jobs and awaits them before returning; `wait=False` returns job IDs immediately.
+- `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.
 - `collect` returns already-finished job results without waiting; `await` waits for requested job IDs.
 - Canceled jobs should be marked terminal and should stop their background `DeferredTask` when cancellation is possible.