diff --git a/helpers/parallel_tools.py b/helpers/parallel_tools.py index c1514a025..48bd08f97 100644 --- a/helpers/parallel_tools.py +++ b/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.") diff --git a/helpers/parallel_tools.py.dox.md b/helpers/parallel_tools.py.dox.md index 4f65fed81..9ce49d53b 100644 --- a/helpers/parallel_tools.py.dox.md +++ b/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. diff --git a/prompts/agent.system.tool.parallel.md b/prompts/agent.system.tool.parallel.md index fd374f29c..44cbf97aa 100644 --- a/prompts/agent.system.tool.parallel.md +++ b/prompts/agent.system.tool.parallel.md @@ -11,6 +11,7 @@ Rules: - `call_subordinate` inside `parallel` starts an isolated child chat under the parent chat, not a scheduler task - use `wait: false` only when you will collect results later with `job_ids` - if extras list running or ready parallel jobs, collect them before final synthesis +- `timeout` only limits how long this call waits; running jobs continue and can be awaited again by `job_ids` Args: `tool_calls`, `job_ids`, `wait` default `true`, `action` as `start|await|collect|cancel`, `timeout`. diff --git a/tests/test_parallel_tool.py b/tests/test_parallel_tool.py index 40a1459fd..d0e9d30cc 100644 --- a/tests/test_parallel_tool.py +++ b/tests/test_parallel_tool.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import time import sys from types import SimpleNamespace @@ -61,6 +62,27 @@ class _FakeAgent: self.agent_name = "A0" +class _FakeDeferredTask: + def __init__(self, *, ready: bool = False, alive: bool = True, result=None) -> None: + self.ready = ready + self.alive = alive + self._result = result + self.killed = 0 + + def is_ready(self): + return self.ready + + def is_alive(self): + return self.alive + + async def result(self): + return self._result + + def kill(self): + self.killed += 1 + self.alive = False + + def test_normalize_parallel_tool_calls_accepts_normal_tool_request_shapes() -> None: calls = parallel_tools.normalize_parallel_tool_calls( [ @@ -125,6 +147,130 @@ async def test_parallel_jobs_extras_lists_running_and_ready_jobs() -> None: assert "ready to collect" in extras +@pytest.mark.asyncio +async def test_parallel_await_timeout_keeps_running_jobs_awaitable(monkeypatch) -> None: + agent = _FakeAgent() + task = _FakeDeferredTask(alive=True) + job = parallel_tools.ParallelJob( + id="wait-1234abcd", + parent_context_id="ctx", + index=0, + tool_name="wait", + tool_args={"seconds": 60}, + kind="tool", + state="running", + created_at=99.0, + started_at=99.0, + deferred_task=task, # type: ignore[arg-type] + ) + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job}) + times = iter([100.0, 102.0]) + monkeypatch.setattr( + parallel_tools.time, + "time", + lambda: next(times, 102.0), + ) + + results = await parallel_tools.await_parallel_jobs( # type: ignore[arg-type] + agent, + [job.id], + timeout=1, + collect=True, + wait=True, + ) + payload = json.loads(parallel_tools.format_parallel_results(results)) + + assert results[0]["state"] == "running" + assert results[0]["wait_timed_out"] is True + assert payload["status"] == "waiting" + assert payload["wait_timeout"] is True + assert "await" in payload["instruction"] + assert task.killed == 0 + assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job + + +@pytest.mark.asyncio +async def test_parallel_collect_returns_running_jobs_without_waiting_or_canceling() -> None: + from tools.parallel import ParallelTool + + agent = _FakeAgent() + task = _FakeDeferredTask(alive=True) + job = parallel_tools.ParallelJob( + id="wait-collect", + parent_context_id="ctx", + index=0, + tool_name="wait", + tool_args={"seconds": 60}, + kind="tool", + state="running", + deferred_task=task, # type: ignore[arg-type] + ) + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job}) + tool = ParallelTool( + agent, # type: ignore[arg-type] + "parallel", + None, + {"action": "collect", "job_ids": [job.id]}, + "", + None, + ) + + response = await tool.execute(**tool.args) + payload = json.loads(response.message) + + assert payload["status"] == "running" + assert payload["jobs"][0]["job_id"] == job.id + assert "wait_timeout" not in payload + assert task.killed == 0 + assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job + + +@pytest.mark.asyncio +async def test_parallel_cancel_still_stops_and_removes_running_jobs() -> None: + agent = _FakeAgent() + task = _FakeDeferredTask(alive=True) + job = parallel_tools.ParallelJob( + id="wait-cancel", + parent_context_id="ctx", + index=0, + tool_name="wait", + tool_args={"seconds": 60}, + kind="tool", + state="running", + deferred_task=task, # type: ignore[arg-type] + ) + agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job}) + + results = await parallel_tools.cancel_parallel_jobs(agent, [job.id]) # type: ignore[arg-type] + payload = json.loads(parallel_tools.format_parallel_results(results)) + + assert results[0]["state"] == "cancelled" + assert payload["status"] == "cancelled" + assert task.killed == 1 + assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY) + + +@pytest.mark.asyncio +async def test_parallel_recursion_guard_allows_subordinate_children_but_blocks_tool_workers() -> None: + from extensions.python.tool_execute_before._20_block_parallel_recursion import ( + BlockParallelRecursion, + ) + from helpers.errors import RepairableException + + agent = _FakeAgent() + agent.context.set_data(parallel_tools.PARALLEL_WORKER_JOB_KEY, "legacy-job") + assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type] + + agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "subordinate") + assert parallel_tools.is_parallel_worker(agent) is False # type: ignore[arg-type] + await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type] + + agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "tool") + assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type] + with pytest.raises(RepairableException, match="cannot be used inside a parallel worker"): + await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type] + + @pytest.mark.asyncio async def test_parallel_subordinate_jobs_are_visible_child_logs_not_scheduler_tasks(monkeypatch) -> None: class FakeDeferredTask: diff --git a/tools/parallel.py b/tools/parallel.py index abb783a6c..b10b9580f 100644 --- a/tools/parallel.py +++ b/tools/parallel.py @@ -48,19 +48,31 @@ class ParallelTool(Tool): break_loop=False, ) - wait_default = action not in {"start", "background"} + wait_default = action not in {"start", "background", "collect"} wait = parallel_tools.coerce_bool(args.get("wait"), wait_default) - if action in {"await", "wait", "collect"}: + if action in {"await", "wait"}: wait = True if not wait: - if not started_jobs: + if not started_jobs and not job_ids: return Response( message="Error: `wait: false` requires `tool_calls` to start new jobs.", break_loop=False, ) + if not job_ids: + return Response( + message=parallel_tools.format_started_jobs(started_jobs), + break_loop=False, + ) + results = await parallel_tools.await_parallel_jobs( + self.agent, + all_job_ids, + timeout=timeout, + collect=True, + wait=False, + ) return Response( - message=parallel_tools.format_started_jobs(started_jobs), + message=parallel_tools.format_parallel_results(results), break_loop=False, ) @@ -69,6 +81,7 @@ class ParallelTool(Tool): all_job_ids, timeout=timeout, collect=True, + wait=True, ) return Response( message=parallel_tools.format_parallel_results(results), diff --git a/tools/parallel.py.dox.md b/tools/parallel.py.dox.md index 7a25edd5e..988a75e05 100644 --- a/tools/parallel.py.dox.md +++ b/tools/parallel.py.dox.md @@ -23,10 +23,10 @@ - The tool is intended for independent calls only; dependent operations remain sequential. - 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. - `action="start"` starts calls and optionally waits according to `wait`. -- `action="await"` waits for requested job IDs. +- `action="await"` waits for requested job IDs until completion or `timeout`; timeout returns running job handles without canceling them. - `action="collect"` returns completed job results without waiting. - `action="cancel"` requests cancellation for requested job IDs. -- Recursive use of `parallel` from inside a parallel worker is blocked before execution. +- 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`. - 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. ## Key Concepts