From db3a863c83c7f7db923c54c5a42149091c7c7f39 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 13:35:16 -0400 Subject: [PATCH 01/14] feat(workflow): add workflow script parser and AST validation --- src/pythinker_code/tools/workflow/__init__.py | 1 + src/pythinker_code/tools/workflow/engine.py | 137 ++++++++++++++++++ tests/tools/test_workflow_parser.py | 39 +++++ 3 files changed, 177 insertions(+) create mode 100644 src/pythinker_code/tools/workflow/__init__.py create mode 100644 src/pythinker_code/tools/workflow/engine.py create mode 100644 tests/tools/test_workflow_parser.py diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py new file mode 100644 index 00000000..837bf138 --- /dev/null +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -0,0 +1 @@ +# (intentionally empty until Task 5 adds the Workflow tool) diff --git a/src/pythinker_code/tools/workflow/engine.py b/src/pythinker_code/tools/workflow/engine.py new file mode 100644 index 00000000..2abfab33 --- /dev/null +++ b/src/pythinker_code/tools/workflow/engine.py @@ -0,0 +1,137 @@ +"""Runtime-agnostic dynamic-workflow engine. + +Parses a model-written Python workflow script, AST-validates it (literal `meta` +first, no imports / `time` / `random` / `datetime`), and runs the remaining +statements inside an `async def` wrapper in a restricted namespace whose +`agent()`/`parallel()`/`pipeline()` primitives orchestrate subagents through an +injected `agent_runner` callback. + +This is NOT a security sandbox: the host agent already has shell and file tools, +so the script can run nothing the model could not already run. The AST checks +exist for reproducibility and a parseable `meta`, mirroring the reference's +determinism rules — not for isolation. +""" + +from __future__ import annotations + +import ast +from typing import Any, cast + +_FORBIDDEN_NAME_LOADS = { + "__import__", + "eval", + "exec", + "compile", + "open", + "input", + "globals", + "locals", + "vars", + "time", + "random", + "datetime", + "os", + "sys", +} + + +class WorkflowScriptError(Exception): + """Raised when a workflow script is structurally invalid.""" + + +class WorkflowMetaPhase: + __slots__ = ("title", "detail", "model") + + def __init__(self, title: str, detail: str | None = None, model: str | None = None) -> None: + self.title = title + self.detail = detail + self.model = model + + +class WorkflowMeta: + __slots__ = ("name", "description", "when_to_use", "phases") + + def __init__( + self, + name: str, + description: str, + when_to_use: str | None = None, + phases: tuple[WorkflowMetaPhase, ...] = (), + ) -> None: + self.name = name + self.description = description + self.when_to_use = when_to_use + self.phases = phases + + +def parse_workflow_script(script: str) -> tuple[WorkflowMeta, list[ast.stmt]]: + """Parse + validate a workflow script. Returns (meta, post-meta statements).""" + try: + tree = ast.parse(script, filename="", mode="exec") + except SyntaxError as exc: + raise WorkflowScriptError(f"workflow script is not valid Python: {exc}") from exc + if not tree.body: + raise WorkflowScriptError("workflow script is empty") + meta = _extract_meta(tree.body[0]) + _assert_deterministic(tree) + return meta, tree.body[1:] + + +def _extract_meta(node: ast.stmt) -> WorkflowMeta: + if not ( + isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == "meta" + ): + raise WorkflowScriptError("the first statement must be `meta = { ... }`") + try: + raw = ast.literal_eval(node.value) + except (ValueError, SyntaxError, TypeError) as exc: + raise WorkflowScriptError( + "meta must be a literal dict (no function calls, names, or interpolation)" + ) from exc + return _validate_meta(raw) + + +def _validate_meta(raw: Any) -> WorkflowMeta: + if not isinstance(raw, dict): + raise WorkflowScriptError("meta must be a dict") + # isinstance narrows Any to dict[Unknown, Unknown]; cast to the correct type. + data = cast(dict[str, Any], raw) + name: Any = data.get("name") + description: Any = data.get("description") + if not isinstance(name, str) or not name.strip(): + raise WorkflowScriptError("meta.name must be a non-empty string") + if not isinstance(description, str) or not description.strip(): + raise WorkflowScriptError("meta.description must be a non-empty string") + when_to_use: Any = data.get("when_to_use") + if when_to_use is not None and not isinstance(when_to_use, str): + raise WorkflowScriptError("meta.when_to_use must be a string") + phases_raw: Any = data.get("phases", []) + if not isinstance(phases_raw, list): + raise WorkflowScriptError("meta.phases must be a list") + phases: list[WorkflowMetaPhase] = [] + for entry in cast(list[Any], phases_raw): + entry_d = cast(dict[str, Any], entry) + if not isinstance(entry, dict) or not isinstance(entry_d.get("title"), str): + raise WorkflowScriptError("each meta phase must have a title string") + phases.append( + WorkflowMetaPhase(entry_d["title"], entry_d.get("detail"), entry_d.get("model")) + ) + return WorkflowMeta(name.strip(), description.strip(), when_to_use, tuple(phases)) + + +def _assert_deterministic(tree: ast.AST) -> None: + for node in ast.walk(tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + raise WorkflowScriptError("import statements are not allowed in a workflow script") + if ( + isinstance(node, ast.Name) + and isinstance(node.ctx, ast.Load) + and node.id in _FORBIDDEN_NAME_LOADS + ): + raise WorkflowScriptError( + f"`{node.id}` is not allowed: workflow scripts must be deterministic and " + "may not import modules or read the clock / RNG" + ) diff --git a/tests/tools/test_workflow_parser.py b/tests/tools/test_workflow_parser.py new file mode 100644 index 00000000..2f2bca64 --- /dev/null +++ b/tests/tools/test_workflow_parser.py @@ -0,0 +1,39 @@ +import pytest + +from pythinker_code.tools.workflow.engine import ( + WorkflowScriptError, + parse_workflow_script, +) + +GOOD = '''meta = {"name": "inspect", "description": "Inspect repo", "phases": [{"title": "Scan"}]} +phase("Scan") +inventory = await agent("Inspect the repository.", {"label": "repo inventory"}) +return {"inventory": inventory} +''' + + +def test_parse_accepts_valid_script(): + meta, body = parse_workflow_script(GOOD) + assert meta.name == "inspect" + assert meta.description == "Inspect repo" + assert meta.phases[0].title == "Scan" + # meta statement is stripped; body keeps the rest. + assert len(body) == 3 + + +@pytest.mark.parametrize( + "script, fragment", + [ + ('phase("x")\n', "first statement"), + ('meta = compute()\nawait agent("x")\n', "literal dict"), + ('meta = {"name": "", "description": "d"}\nawait agent("x")\n', "name"), + ('meta = {"name": "n", "description": ""}\nawait agent("x")\n', "description"), + ('meta = {"name": "n", "description": "d"}\nimport os\n', "not allowed"), + ('meta = {"name": "n", "description": "d"}\nx = random.random()\n', "deterministic"), + ('meta = {"name": "n", "description": "d"}\nx = time.time()\n', "deterministic"), + ], +) +def test_parse_rejects(script, fragment): + with pytest.raises(WorkflowScriptError) as exc: + parse_workflow_script(script) + assert fragment in str(exc.value) From abd31c53ab4625e509f7a5b49c37515539a492bd Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 13:42:58 -0400 Subject: [PATCH 02/14] fix(workflow): harden phase validation and add rejection tests --- src/pythinker_code/tools/workflow/engine.py | 4 +++- tests/tools/test_workflow_parser.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/pythinker_code/tools/workflow/engine.py b/src/pythinker_code/tools/workflow/engine.py index 2abfab33..ae2d0fd3 100644 --- a/src/pythinker_code/tools/workflow/engine.py +++ b/src/pythinker_code/tools/workflow/engine.py @@ -113,8 +113,10 @@ def _validate_meta(raw: Any) -> WorkflowMeta: raise WorkflowScriptError("meta.phases must be a list") phases: list[WorkflowMetaPhase] = [] for entry in cast(list[Any], phases_raw): + if not isinstance(entry, dict): + raise WorkflowScriptError("each meta phase must have a title string") entry_d = cast(dict[str, Any], entry) - if not isinstance(entry, dict) or not isinstance(entry_d.get("title"), str): + if not isinstance(entry_d.get("title"), str): raise WorkflowScriptError("each meta phase must have a title string") phases.append( WorkflowMetaPhase(entry_d["title"], entry_d.get("detail"), entry_d.get("model")) diff --git a/tests/tools/test_workflow_parser.py b/tests/tools/test_workflow_parser.py index 2f2bca64..1518bb45 100644 --- a/tests/tools/test_workflow_parser.py +++ b/tests/tools/test_workflow_parser.py @@ -31,6 +31,9 @@ def test_parse_accepts_valid_script(): ('meta = {"name": "n", "description": "d"}\nimport os\n', "not allowed"), ('meta = {"name": "n", "description": "d"}\nx = random.random()\n', "deterministic"), ('meta = {"name": "n", "description": "d"}\nx = time.time()\n', "deterministic"), + ('meta = {"name": "n", "description": "d", "when_to_use": 42}\nawait agent("x")\n', "when_to_use"), + ('meta = {"name": "n", "description": "d", "phases": "oops"}\nawait agent("x")\n', "phases"), + ('meta = {"name": "n", "description": "d", "phases": [42]}\nawait agent("x")\n', "title"), ], ) def test_parse_rejects(script, fragment): From 69a854329d7ea415a44c9e03282ecf038ef31f88 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 13:54:20 -0400 Subject: [PATCH 03/14] feat(workflow): add sandboxed engine with agent/parallel/pipeline primitives --- src/pythinker_code/tools/workflow/engine.py | 366 ++++++++++++++++++++ tests/tools/test_workflow_engine.py | 148 ++++++++ 2 files changed, 514 insertions(+) create mode 100644 tests/tools/test_workflow_engine.py diff --git a/src/pythinker_code/tools/workflow/engine.py b/src/pythinker_code/tools/workflow/engine.py index ae2d0fd3..7e6ff695 100644 --- a/src/pythinker_code/tools/workflow/engine.py +++ b/src/pythinker_code/tools/workflow/engine.py @@ -15,6 +15,12 @@ from __future__ import annotations import ast +import asyncio +import builtins as _builtins +import inspect +import json +import math +from collections.abc import Awaitable, Callable, Sequence from typing import Any, cast _FORBIDDEN_NAME_LOADS = { @@ -137,3 +143,363 @@ def _assert_deterministic(tree: ast.AST) -> None: f"`{node.id}` is not allowed: workflow scripts must be deterministic and " "may not import modules or read the clock / RNG" ) + + +# --------------------------------------------------------------------------- +# Runtime half — execution primitives and run_workflow entry point +# --------------------------------------------------------------------------- + +AgentRunner = Callable[[str, "AgentOptions"], Awaitable[Any]] + +_SAFE_BUILTIN_NAMES = ( + "len", + "range", + "enumerate", + "list", + "dict", + "set", + "tuple", + "frozenset", + "str", + "int", + "float", + "bool", + "bytes", + "sorted", + "reversed", + "min", + "max", + "sum", + "any", + "all", + "map", + "filter", + "zip", + "abs", + "round", + "isinstance", + "repr", + "Exception", + "ValueError", + "TypeError", + "KeyError", + "IndexError", +) + + +class WorkflowRuntimeError(Exception): + """Raised when a workflow script misuses a primitive at runtime.""" + + +class AgentOptions: + __slots__ = ("label", "phase", "schema", "model", "agent_type") + + def __init__( + self, + label: str | None = None, + phase: str | None = None, + schema: dict[str, Any] | None = None, + model: str | None = None, + agent_type: str | None = None, + ) -> None: + self.label = label + self.phase = phase + self.schema = schema + self.model = model + self.agent_type = agent_type + + +class AgentStartEvent: + __slots__ = ("label", "phase", "prompt") + + def __init__(self, label: str, phase: str | None, prompt: str) -> None: + self.label = label + self.phase = phase + self.prompt = prompt + + +class AgentEndEvent: + __slots__ = ("label", "phase", "result", "error") + + def __init__( + self, label: str, phase: str | None, result: Any, error: str | None = None + ) -> None: + self.label = label + self.phase = phase + self.result = result + self.error = error + + +class RunWorkflowHooks: + __slots__ = ("on_log", "on_phase", "on_agent_start", "on_agent_end") + + def __init__( + self, + on_log: Callable[[str], None] | None = None, + on_phase: Callable[[str], None] | None = None, + on_agent_start: Callable[[AgentStartEvent], None] | None = None, + on_agent_end: Callable[[AgentEndEvent], None] | None = None, + ) -> None: + self.on_log = on_log + self.on_phase = on_phase + self.on_agent_start = on_agent_start + self.on_agent_end = on_agent_end + + +class WorkflowRunResult: + __slots__ = ("meta", "result", "logs", "phases", "agent_count") + + def __init__( + self, + meta: WorkflowMeta, + result: Any, + logs: list[str], + phases: list[str], + agent_count: int, + ) -> None: + self.meta = meta + self.result = result + self.logs = logs + self.phases = phases + self.agent_count = agent_count + + +def _require_str(value: Any, name: str) -> str: + if not isinstance(value, str): + raise WorkflowRuntimeError(f"{name} must be a string") + return value + + +def _normalize_options(value: Any) -> AgentOptions: + if value is None: + return AgentOptions() + if not isinstance(value, dict): + raise WorkflowRuntimeError("agent options must be a dict") + d = cast(dict[str, Any], value) + return AgentOptions( + label=d.get("label"), + phase=d.get("phase"), + schema=d.get("schema"), + model=d.get("model"), + agent_type=d.get("agent_type") or d.get("agentType"), + ) + + +def _default_label(phase: str | None, index: int) -> str: + return f"{phase} agent {index}" if phase else f"agent {index}" + + +def _estimate_tokens(value: Any) -> int: + try: + text = json.dumps(value, default=str) + except (TypeError, ValueError): + text = str(value) + return math.ceil(len(text) / 4) + + +def _close_coroutines(value: Any) -> None: + if inspect.iscoroutine(value): + value.close() + elif isinstance(value, dict): + for item in cast(dict[str, Any], value).values(): + _close_coroutines(item) + elif isinstance(value, (list, tuple, set)): + for item in cast(list[Any], value): + _close_coroutines(item) + + +def _assert_no_coroutines(value: Any) -> None: + found = _contains_coroutine(value) + if found: + _close_coroutines(value) + raise WorkflowRuntimeError( + "workflow result contains a coroutine; did you forget to await " + "agent(), parallel(), or pipeline()?" + ) + + +def _contains_coroutine(value: Any) -> bool: + if inspect.iscoroutine(value): + return True + if isinstance(value, dict): + return any(_contains_coroutine(v) for v in cast(dict[str, Any], value).values()) + if isinstance(value, (list, tuple, set)): + return any(_contains_coroutine(v) for v in cast(list[Any], value)) + return False + + +async def run_workflow( + script: str, + *, + agent_runner: AgentRunner, + args: Any = None, + cwd: str = ".", + concurrency: int = 8, + token_budget: int | None = None, + hooks: RunWorkflowHooks | None = None, +) -> WorkflowRunResult: + meta, body = parse_workflow_script(script) + hooks = hooks or RunWorkflowHooks() + logs: list[str] = [] + phases: list[str] = [] + state: dict[str, Any] = {"current_phase": None, "agent_count": 0, "spent": 0} + semaphore = asyncio.Semaphore(max(1, concurrency)) + + def log(message: Any) -> None: + text = str(message) + logs.append(text) + if hooks.on_log: + hooks.on_log(text) + + def phase(title: Any) -> None: + text = _require_str(title, "phase title") + state["current_phase"] = text + if text not in phases: + phases.append(text) + if hooks.on_phase: + hooks.on_phase(text) + + class _Budget: + total = token_budget + + @staticmethod + def spent() -> int: + return state["spent"] + + @staticmethod + def remaining() -> float: + if token_budget is None: + return math.inf + return max(0, token_budget - state["spent"]) + + budget = _Budget() + + async def agent(prompt: Any, options: Any = None) -> Any: + task_prompt = _require_str(prompt, "agent prompt") + opts = _normalize_options(options) + assigned_phase = opts.phase or state["current_phase"] + if token_budget is not None and budget.remaining() <= 0: + raise WorkflowRuntimeError("workflow token budget exhausted") + async with semaphore: + state["agent_count"] += 1 + label = (opts.label or "").strip() or _default_label( + assigned_phase, state["agent_count"] + ) + opts.label = label + opts.phase = assigned_phase + if hooks.on_agent_start: + hooks.on_agent_start(AgentStartEvent(label, assigned_phase, task_prompt)) + try: + result = await agent_runner(task_prompt, opts) + except asyncio.CancelledError: + if hooks.on_agent_end: + hooks.on_agent_end( + AgentEndEvent(label, assigned_phase, None, error="cancelled") + ) + raise + except Exception as exc: # noqa: BLE001 - reference parity: branch fails to None + log(f"agent {label} failed: {exc}") + if hooks.on_agent_end: + hooks.on_agent_end(AgentEndEvent(label, assigned_phase, None, error=str(exc))) + return None + state["spent"] += _estimate_tokens(result) + if hooks.on_agent_end: + hooks.on_agent_end(AgentEndEvent(label, assigned_phase, result)) + return result + + async def parallel(items: Sequence[Any]) -> list[Any]: + if not isinstance(items, (list, tuple)): + raise WorkflowRuntimeError("parallel() expects a list of awaitables") + for item in items: + if callable(item) and not inspect.isawaitable(item): + raise WorkflowRuntimeError( + "parallel() expects awaitables, not functions: " + "use parallel([agent('...'), agent('...')])" + ) + + async def guarded(index: int, item: Any) -> Any: + try: + return await item + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 - reference parity + log(f"parallel[{index}] failed: {exc}") + return None + + return list(await asyncio.gather(*(guarded(i, it) for i, it in enumerate(items)))) + + async def pipeline(items: Any, *stages: Any) -> list[Any]: + if not isinstance(items, (list, tuple)): + raise WorkflowRuntimeError("pipeline() expects a list as the first argument") + for stage in stages: + if not callable(stage): + raise WorkflowRuntimeError("pipeline() stages must be callables") + typed_items: list[Any] = list(cast(Any, items)) + + async def run_item(index: int, item: Any) -> Any: + value: Any = item + for stage in stages: + try: + produced = stage(value, item, index) + value = await produced if inspect.isawaitable(produced) else produced + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 - reference parity + log(f"pipeline[{index}] failed: {exc}") + return None + return value + + return list(await asyncio.gather(*(run_item(i, it) for i, it in enumerate(typed_items)))) + + def _print(*a: Any, **_kw: Any) -> None: + log(" ".join(str(x) for x in a)) + + namespace: dict[str, Any] = { + "__builtins__": {name: getattr(_builtins, name) for name in _SAFE_BUILTIN_NAMES}, + "agent": agent, + "parallel": parallel, + "pipeline": pipeline, + "phase": phase, + "log": log, + "print": _print, + "budget": budget, + "args": args, + "cwd": cwd, + "json": json, + "math": math, + } + + result = await _execute(body, namespace) + _assert_no_coroutines(result) + return WorkflowRunResult( + meta=meta, + result=result, + logs=logs, + phases=phases, + agent_count=int(state["agent_count"]), + ) + + +async def _execute(body: list[ast.stmt], namespace: dict[str, Any]) -> Any: + func = ast.AsyncFunctionDef( + name="__workflow_main__", + args=ast.arguments( + posonlyargs=[], + args=[], + vararg=None, + kwonlyargs=[], + kw_defaults=[], + kwarg=None, + defaults=[], + ), + body=body or [ast.Pass()], + decorator_list=[], + returns=None, + type_comment=None, + type_params=[], + ) + module = ast.Module(body=[func], type_ignores=[]) + ast.fix_missing_locations(module) + code = compile(module, filename="", mode="exec") + exec(code, namespace) # noqa: S102 - not a security boundary; see module docstring + return await namespace["__workflow_main__"]() diff --git a/tests/tools/test_workflow_engine.py b/tests/tools/test_workflow_engine.py new file mode 100644 index 00000000..97ac1fce --- /dev/null +++ b/tests/tools/test_workflow_engine.py @@ -0,0 +1,148 @@ +# tests/tools/test_workflow_engine.py +import asyncio + +import pytest + +from pythinker_code.tools.workflow.engine import ( + AgentOptions, + RunWorkflowHooks, + WorkflowRuntimeError, + run_workflow, +) + + +def make_runner(delay: float = 0.0): + calls: list[str] = [] + + async def runner(prompt: str, opts: AgentOptions): + calls.append(prompt) + if delay: + await asyncio.sleep(delay) + return f"result:{prompt}" + + return runner, calls + + +@pytest.mark.asyncio +async def test_single_agent_and_return(): + runner, calls = make_runner() + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("hello", {"label": "L"})\n' + "return {\"r\": r}\n" + ) + out = await run_workflow(script, agent_runner=runner, cwd=".") + assert out.result == {"r": "result:hello"} + assert out.agent_count == 1 + assert calls == ["hello"] + + +@pytest.mark.asyncio +async def test_multiline_prompt_string_preserved(): + # Regression guard: the AST-wrap must NOT alter triple-quoted strings the way + # textwrap.indent would (it would inject leading spaces on continuation lines). + seen: list[str] = [] + + async def runner(prompt: str, opts): + seen.append(prompt) + return "ok" + + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'p = """line one\n' + "line two\n" + ' indented body"""\n' + "await agent(p)\n" + "return None\n" + ) + await run_workflow(script, agent_runner=runner, cwd=".") + assert seen == ["line one\nline two\n indented body"] + + +@pytest.mark.asyncio +async def test_parallel_preserves_order(): + runner, _ = make_runner() + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'rs = await parallel([agent("a"), agent("b"), agent("c")])\n' + "return rs\n" + ) + out = await run_workflow(script, agent_runner=runner, cwd=".") + assert out.result == ["result:a", "result:b", "result:c"] + assert out.agent_count == 3 + + +@pytest.mark.asyncio +async def test_pipeline_stages_receive_prev_original_index(): + async def runner(prompt, opts): + return prompt.upper() + + script = ( + 'meta = {"name": "n", "description": "d"}\n' + "rs = await pipeline(\n" + ' ["x", "y"],\n' + " lambda prev, orig, i: agent(prev),\n" + ' lambda prev, orig, i: prev + ":" + orig + ":" + str(i),\n' + ")\n" + "return rs\n" + ) + out = await run_workflow(script, agent_runner=runner, cwd=".") + assert out.result == ["X:x:0", "Y:y:1"] + + +@pytest.mark.asyncio +async def test_failed_branch_returns_none_and_logs(): + async def runner(prompt, opts): + raise RuntimeError("boom") + + logs: list[str] = [] + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("x", {"label": "bad"})\n' + 'return {"r": r}\n' + ) + out = await run_workflow( + script, + agent_runner=runner, + cwd=".", + hooks=RunWorkflowHooks(on_log=logs.append), + ) + assert out.result == {"r": None} + assert any("bad" in m and "boom" in m for m in logs) + + +@pytest.mark.asyncio +async def test_unawaited_coroutine_in_result_raises(): + runner, _ = make_runner() + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'return {"r": agent("x")}\n' # not awaited + ) + with pytest.raises(WorkflowRuntimeError) as exc: + await run_workflow(script, agent_runner=runner, cwd=".") + assert "await" in str(exc.value) + + +@pytest.mark.asyncio +async def test_cancellation_marks_running_skipped_and_reraises(): + runner, _ = make_runner(delay=10.0) + skipped: list[str] = [] + started: list[str] = [] + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'rs = await parallel([agent("a", {"label": "a"}), agent("b", {"label": "b"})])\n' + "return rs\n" + ) + hooks = RunWorkflowHooks( + on_agent_start=lambda e: started.append(e.label), + on_agent_end=lambda e: skipped.append(e.label) if e.error == "cancelled" else None, + ) + task = asyncio.create_task( + run_workflow(script, agent_runner=runner, cwd=".", concurrency=4, hooks=hooks) + ) + await asyncio.sleep(0.05) # let both agents start + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert set(started) == {"a", "b"} + assert set(skipped) == {"a", "b"} # in-flight agents reported as cancelled, none completed From c2319e9394450ac731dfb0460026157ab47b3219 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 14:00:29 -0400 Subject: [PATCH 04/14] feat(workflow): add progress snapshot and renderer --- src/pythinker_code/tools/workflow/display.py | 99 ++++++++++++++++++++ tests/tools/test_workflow_display.py | 28 ++++++ 2 files changed, 127 insertions(+) create mode 100644 src/pythinker_code/tools/workflow/display.py create mode 100644 tests/tools/test_workflow_display.py diff --git a/src/pythinker_code/tools/workflow/display.py b/src/pythinker_code/tools/workflow/display.py new file mode 100644 index 00000000..126558f8 --- /dev/null +++ b/src/pythinker_code/tools/workflow/display.py @@ -0,0 +1,99 @@ +"""Compact progress snapshot + renderer for the Workflow tool. + +A port of the reference display module, trimmed to a single text block suitable +for a `ProgressNote` wire event. +""" + +from __future__ import annotations + +_STATUS_ICON = {"running": "●", "done": "✓", "error": "✗", "skipped": "-"} + + +class AgentSnapshot: + __slots__ = ("id", "label", "phase", "status") + + def __init__(self, id: int, label: str, phase: str | None) -> None: + self.id = id + self.label = label + self.phase = phase + self.status = "running" + + +class WorkflowSnapshot: + __slots__ = ("name", "description", "phases", "current_phase", "agents", "logs") + + def __init__(self, name: str, description: str = "") -> None: + self.name = name + self.description = description + self.phases: list[str] = [] + self.current_phase: str | None = None + self.agents: list[AgentSnapshot] = [] + self.logs: list[str] = [] + + def add_phase(self, title: str | None) -> None: + if not title: + return + self.current_phase = title + if title not in self.phases: + self.phases.append(title) + + def start_agent(self, label: str, phase: str | None) -> AgentSnapshot: + self.add_phase(phase) + agent = AgentSnapshot(len(self.agents) + 1, label, phase) + self.agents.append(agent) + return agent + + def end_agent(self, label: str, *, error: str | None = None) -> None: + for agent in reversed(self.agents): + if agent.label == label and agent.status == "running": + agent.status = "error" if error else "done" + return + + def mark_running_skipped(self) -> None: + for agent in self.agents: + if agent.status == "running": + agent.status = "skipped" + + @property + def running_count(self) -> int: + return sum(1 for a in self.agents if a.status == "running") + + @property + def done_count(self) -> int: + return sum(1 for a in self.agents if a.status == "done") + + @property + def error_count(self) -> int: + return sum(1 for a in self.agents if a.status == "error") + + @property + def skipped_count(self) -> int: + return sum(1 for a in self.agents if a.status == "skipped") + + +def render_progress(snapshot: WorkflowSnapshot, max_agents: int = 6) -> str: + state = "" + if snapshot.error_count: + state = f", {snapshot.error_count} errors" + elif snapshot.running_count: + state = f", {snapshot.running_count} running" + lines = [ + f"◆ Workflow: {snapshot.name} ({snapshot.done_count}/{len(snapshot.agents)} done{state})" + ] + rendered: set[int] = set() + phase_order = list(snapshot.phases) + if snapshot.current_phase and snapshot.current_phase not in phase_order: + phase_order.append(snapshot.current_phase) + for phase in phase_order: + agents = [a for a in snapshot.agents if a.phase == phase] + if not agents and snapshot.current_phase != phase: + continue + done = sum(1 for a in agents if a.status == "done") + lines.append(f" {phase} {done}/{len(agents)}") + for agent in agents[-max_agents:]: + rendered.add(agent.id) + lines.append(f" #{agent.id} {_STATUS_ICON.get(agent.status, '?')} {agent.label}") + unphased = [a for a in snapshot.agents if a.id not in rendered] + for agent in unphased[-max_agents:]: + lines.append(f" #{agent.id} {_STATUS_ICON.get(agent.status, '?')} {agent.label}") + return "\n".join(lines) diff --git a/tests/tools/test_workflow_display.py b/tests/tools/test_workflow_display.py new file mode 100644 index 00000000..940ab9a9 --- /dev/null +++ b/tests/tools/test_workflow_display.py @@ -0,0 +1,28 @@ +from pythinker_code.tools.workflow.display import WorkflowSnapshot, render_progress + + +def test_snapshot_lifecycle_and_render(): + snap = WorkflowSnapshot(name="inspect", description="d") + snap.add_phase("Scan") + a = snap.start_agent("repo inventory", "Scan") + assert a.status == "running" + assert snap.running_count == 1 + snap.end_agent("repo inventory") + assert snap.done_count == 1 and snap.running_count == 0 + + snap.add_phase("Analyze") + snap.start_agent("modules", "Analyze") + text = render_progress(snap) + assert "Workflow: inspect" in text + assert "Scan" in text and "Analyze" in text + assert "repo inventory" in text + + +def test_mark_running_skipped(): + snap = WorkflowSnapshot(name="n", description="d") + snap.start_agent("a", None) + snap.start_agent("b", None) + snap.end_agent("a") + snap.mark_running_skipped() + assert snap.done_count == 1 + assert snap.skipped_count == 1 From 3d0b47453dfb4a7d3eaeb56f9953fa071eb14523 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 14:10:47 -0400 Subject: [PATCH 05/14] fix(workflow): gate token budget check inside the agent semaphore agent() checked budget.remaining() before acquiring the concurrency semaphore, so every agent() call dispatched together (e.g. via parallel()) saw the same stale state["spent"] and passed regardless of the configured concurrency, only catching the overrun after the fact. Move the check inside async with semaphore so it re-evaluates spend once a slot is actually granted, bounding worst-case overshoot to one concurrency batch instead of the whole dispatched set. --- src/pythinker_code/tools/workflow/engine.py | 12 +++++-- tests/tools/test_workflow_engine.py | 37 ++++++++++++++++++++- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/pythinker_code/tools/workflow/engine.py b/src/pythinker_code/tools/workflow/engine.py index 7e6ff695..32a3b256 100644 --- a/src/pythinker_code/tools/workflow/engine.py +++ b/src/pythinker_code/tools/workflow/engine.py @@ -378,9 +378,17 @@ async def agent(prompt: Any, options: Any = None) -> Any: task_prompt = _require_str(prompt, "agent prompt") opts = _normalize_options(options) assigned_phase = opts.phase or state["current_phase"] - if token_budget is not None and budget.remaining() <= 0: - raise WorkflowRuntimeError("workflow token budget exhausted") async with semaphore: + # Checked AFTER acquiring the semaphore, not before: agent() calls + # dispatched together (e.g. via parallel()) all reach this point + # before any of them has recorded real spend, so the check would + # otherwise let every dispatched agent through regardless of + # concurrency, deferring rejection past the budget. Gating inside + # the semaphore re-evaluates state["spent"] once a slot actually + # frees, bounding the worst-case overshoot to one concurrency + # batch instead of the whole dispatched set. + if token_budget is not None and budget.remaining() <= 0: + raise WorkflowRuntimeError("workflow token budget exhausted") state["agent_count"] += 1 label = (opts.label or "").strip() or _default_label( assigned_phase, state["agent_count"] diff --git a/tests/tools/test_workflow_engine.py b/tests/tools/test_workflow_engine.py index 97ac1fce..6b5d6f7c 100644 --- a/tests/tools/test_workflow_engine.py +++ b/tests/tools/test_workflow_engine.py @@ -7,6 +7,7 @@ AgentOptions, RunWorkflowHooks, WorkflowRuntimeError, + _estimate_tokens, run_workflow, ) @@ -29,7 +30,7 @@ async def test_single_agent_and_return(): script = ( 'meta = {"name": "n", "description": "d"}\n' 'r = await agent("hello", {"label": "L"})\n' - "return {\"r\": r}\n" + 'return {"r": r}\n' ) out = await run_workflow(script, agent_runner=runner, cwd=".") assert out.result == {"r": "result:hello"} @@ -123,6 +124,40 @@ async def test_unawaited_coroutine_in_result_raises(): assert "await" in str(exc.value) +@pytest.mark.asyncio +async def test_budget_check_does_not_race_past_semaphore(): + # Regression guard: agent() calls dispatched together (via parallel()) all + # reach the budget check before any of them has recorded real spend. If the + # check runs before the semaphore is acquired, every dispatched agent sees + # the same stale state["spent"] and passes, regardless of `concurrency` — + # so a 6-item parallel() with a budget for ~1 result would let all 6 run. + # The check must be re-evaluated once a semaphore slot is actually granted, + # bounding the overshoot to one concurrency batch instead of every item. + result_payload = "x" * 200 + one_result_cost = _estimate_tokens(result_payload) + + async def runner(prompt: str, opts: AgentOptions) -> str: + await asyncio.sleep(0.01) # force a real suspension so calls overlap + return result_payload + + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'rs = await parallel([agent("a"), agent("b"), agent("c"), ' + 'agent("d"), agent("e"), agent("f")])\n' + "return rs\n" + ) + out = await run_workflow( + script, + agent_runner=runner, + cwd=".", + concurrency=2, + token_budget=one_result_cost + 1, + ) + succeeded = [r for r in out.result if r == result_payload] + assert len(succeeded) <= 2 # bounded by concurrency, not by the 6 dispatched items + assert len(succeeded) >= 1 # the first batch must still get through + + @pytest.mark.asyncio async def test_cancellation_marks_running_skipped_and_reraises(): runner, _ = make_runner(delay=10.0) From 797ceb1275d1b8ec9df80e80fe3921bd214b7dcb Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 14:26:55 -0400 Subject: [PATCH 06/14] feat(workflow): add Workflow tool over AgentTool with approval and schema support --- src/pythinker_code/tools/workflow/__init__.py | 228 +++++++++++++++++- .../tools/workflow/description.md | 17 ++ tests/tools/test_workflow_tool.py | 145 +++++++++++ tests/utils/test_pyinstaller_utils.py | 7 + 4 files changed, 396 insertions(+), 1 deletion(-) create mode 100644 src/pythinker_code/tools/workflow/description.md create mode 100644 tests/tools/test_workflow_tool.py diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py index 837bf138..e597b314 100644 --- a/src/pythinker_code/tools/workflow/__init__.py +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -1 +1,227 @@ -# (intentionally empty until Task 5 adds the Workflow tool) +import hashlib +import json +import re +from pathlib import Path +from typing import Any, override + +import jsonschema +from pydantic import BaseModel, Field +from pythinker_core.message import ContentPart, TextPart +from pythinker_core.tooling import CallableTool2, ToolError, ToolOk, ToolReturnValue + +from pythinker_code.soul import wire_send +from pythinker_code.soul.agent import Runtime +from pythinker_code.tools.agent import AgentTool +from pythinker_code.tools.agent import Params as AgentParams +from pythinker_code.tools.utils import load_desc +from pythinker_code.tools.workflow.display import WorkflowSnapshot, render_progress +from pythinker_code.tools.workflow.engine import ( + AgentEndEvent, + AgentOptions, + AgentStartEvent, + RunWorkflowHooks, + WorkflowScriptError, + parse_workflow_script, + run_workflow, +) +from pythinker_code.utils.logging import logger +from pythinker_code.wire.types import ProgressNote + +_FENCE_RE = re.compile(r"^```(?:py|python)?\s*\n([\s\S]*?)\n```$", re.IGNORECASE) +_JSON_BLOCK_RE = re.compile(r"```json\s*\n([\s\S]*?)\n```", re.IGNORECASE) +_SUMMARY_SEP = "\n[summary]\n" + + +class Params(BaseModel): + script: str = Field( + description=( + "Raw Python workflow script (no Markdown fences). First statement must be " + '`meta = {"name": ..., "description": ...}`. Must call agent() at least once ' + "and return a JSON-serializable value." + ) + ) + args: Any | None = Field( + default=None, + description="Optional JSON value exposed to the script as the global `args`.", + ) + + +class Workflow(CallableTool2[Params]): + name: str = "Workflow" + params: type[Params] = Params + emits_tool_execution_started_after_approval = True + + def __init__(self, runtime: Runtime) -> None: + super().__init__(description=load_desc(Path(__file__).parent / "description.md")) + self._runtime = runtime + self._agent_tool = AgentTool(runtime) + + @override + async def __call__(self, params: Params) -> ToolReturnValue: + if self._runtime.role != "root": + return ToolError( + message="Workflow can only run from the root agent; subagents may not " + "launch nested workflows.", + brief="Not allowed in subagent", + ) + + script = _strip_fences(params.script) + try: + meta, _ = parse_workflow_script(script) + except WorkflowScriptError as exc: + return ToolError(message=str(exc), brief="Invalid workflow script") + + fingerprint = _fingerprint(script, params.args) + if self._runtime.approval.is_orchestration_approved(fingerprint): + from pythinker_code.soul.toolset import emit_current_tool_execution_started + + emit_current_tool_execution_started() + else: + approval = await self._runtime.approval.request( + self.name, + "run workflow orchestration", + f"Run workflow `{meta.name}`: {meta.description}", + ) + if not approval: + return approval.rejection_error() + self._runtime.approval.approve_orchestration(fingerprint) + + snapshot = WorkflowSnapshot(name=meta.name, description=meta.description) + + def emit() -> None: + wire_send(ProgressNote(title=f"Workflow {meta.name}", body=render_progress(snapshot))) + + def on_phase(title: str) -> None: + snapshot.add_phase(title) + emit() + + def on_agent_start(event: AgentStartEvent) -> None: + snapshot.start_agent(event.label, event.phase) + emit() + + def on_agent_end(event: AgentEndEvent) -> None: + snapshot.end_agent(event.label, error=event.error) + emit() + + hooks = RunWorkflowHooks( + on_log=lambda _m: emit(), + on_phase=on_phase, + on_agent_start=on_agent_start, + on_agent_end=on_agent_end, + ) + + try: + result = await run_workflow( + script, + agent_runner=self._agent_runner, + args=params.args, + cwd=str(self._runtime.work_dir), + concurrency=max(1, self._runtime.config.background.max_running_tasks), + hooks=hooks, + ) + except WorkflowScriptError as exc: + return ToolError(message=str(exc), brief="Invalid workflow script") + except Exception as exc: # noqa: BLE001 - converted to a typed ToolError below + # asyncio.CancelledError is a BaseException in 3.12+ and is NOT caught here; + # it propagates so the host can abort the tool. The engine already fires + # on_agent_end(..., error="cancelled") for in-flight agents before + # re-raising, so the incrementally-emitted snapshot already shows them. + logger.exception("Workflow run failed") + return ToolError(message=f"workflow failed: {exc}", brief="Workflow failed") + + if result.agent_count == 0: + return ToolError( + message="workflow scripts must call agent() at least once; this workflow " + "declared phases but ran no subagents.", + brief="No agents run", + ) + + output = json.dumps(result.result, indent=2, default=str) + return ToolOk( + output=output, + message=f"Workflow {meta.name} completed with {result.agent_count} agent(s).", + brief=f"{result.agent_count} agent(s), {len(result.phases)} phase(s)", + ) + + async def _agent_runner(self, prompt: str, opts: AgentOptions) -> Any: + if opts.schema is None: + text, is_error, message = await self._run_child(prompt, opts) + if is_error: + raise RuntimeError(message or "subagent failed") + return text + + # Schema mode: instruct, parse, validate; one retry with the error fed back. + instructed = prompt + "\n\n" + _schema_instructions(opts.schema) + last_error = "" + for _attempt in range(2): + text, is_error, message = await self._run_child(instructed, opts) + if is_error: + raise RuntimeError(message or "subagent failed") + value, err = _parse_and_validate(text, opts.schema) + if err is None: + return value + last_error = err + instructed = ( + prompt + + "\n\n" + + _schema_instructions(opts.schema) + + f"\n\nYour previous output was invalid: {err}\nReturn a valid ```json block." + ) + raise RuntimeError(f"subagent never produced schema-valid output: {last_error}") + + async def _run_child(self, prompt: str, opts: AgentOptions) -> tuple[str, bool, str]: + child_params = AgentParams( + description=opts.label or "workflow agent", + prompt=prompt, + subagent_type=opts.agent_type or "coder", + model=opts.model, + run_in_background=False, + ) + ret = await self._agent_tool(child_params) + output = ret.output if isinstance(ret.output, str) else _content_text(ret.output) + return _extract_summary(output), ret.is_error, ret.message + + +def _strip_fences(script: str) -> str: + text = script.strip() + match = _FENCE_RE.match(text) + return match.group(1).strip() if match else text + + +def _fingerprint(script: str, args: Any) -> str: + payload = json.dumps({"script": script, "args": args}, sort_keys=True, default=str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _extract_summary(output: str) -> str: + if _SUMMARY_SEP in output: + return output.split(_SUMMARY_SEP, 1)[1].strip() + return output.strip() + + +def _content_text(parts: list[ContentPart]) -> str: + # AgentTool returns a string; this is a defensive fallback for ContentPart lists. + return "".join(part.text for part in parts if isinstance(part, TextPart)) + + +def _schema_instructions(schema: dict[str, Any]) -> str: + return ( + "Final output contract: your final message MUST end with a single fenced " + "```json block whose contents validate against this JSON Schema:\n" + + json.dumps(schema, indent=2) + + "\nDo not write prose after the json block." + ) + + +def _parse_and_validate(text: str, schema: dict[str, Any]) -> tuple[Any, str | None]: + blocks = _JSON_BLOCK_RE.findall(text) + raw = blocks[-1] if blocks else text + try: + value = json.loads(raw) + except json.JSONDecodeError as exc: + return None, f"output was not valid JSON ({exc})" + try: + jsonschema.validate(value, schema) + except jsonschema.ValidationError as exc: + return None, f"output did not match schema ({exc.message})" + return value, None diff --git a/src/pythinker_code/tools/workflow/description.md b/src/pythinker_code/tools/workflow/description.md new file mode 100644 index 00000000..69a6b49a --- /dev/null +++ b/src/pythinker_code/tools/workflow/description.md @@ -0,0 +1,17 @@ +Execute a deterministic Python workflow that orchestrates multiple subagents with `agent()`, `parallel()`, and `pipeline()`. Use this only when the user explicitly asks for a workflow, fan-out, or multi-agent orchestration, or when a task decomposes into many independent subtasks worth running concurrently and synthesizing. + +`script` is required raw Python (no Markdown fences). Rules: + +- The first statement MUST be `meta = {"name": "short_snake_case", "description": "non-empty description"}`. `meta` must be a literal dict (no function calls or interpolation). `meta["phases"]` is optional documentation; live progress is driven by `phase(title)` at runtime. +- After `meta`, write plain Python. Do NOT use `import`, `time`, `random`, `datetime`, `os`, `sys`, `open`, `eval`, or `exec` — scripts must be deterministic. +- The script must call `agent()` at least once. End by `return`-ing a compact JSON-serializable value. + +Available globals: + +- `agent(prompt, opts=None)` — spawn one subagent; returns its final text, or a validated dict when `opts={"schema": }`. Other opts: `label` (short, 2-5 words), `phase`, `model`, `agent_type` (a built-in subagent type, e.g. `"explore"`, `"coder"`, `"review"`). Always `await` it. A failed agent returns `None` and logs — check for `None` before synthesizing. +- `parallel([awaitables])` — run awaitables concurrently, results in input order: `await parallel([agent("a"), agent("b")])`. Pass awaitables, NOT functions. +- `pipeline(items, *stages)` — run each item through sequential stages while items fan out. Each stage is called `(prev, original, index)` and may be sync or async. +- `phase(title)` — start a progress group. Names may be conditional or built in a loop; do not predeclare speculative phases. +- `log(message)`, `args` (the optional JSON `args` input), `cwd`, `budget` (`.total`, `.spent()`, `.remaining()`). + +Include enough context and file paths in each `agent()` prompt — subagents do NOT inherit the parent conversation. Add a final synthesis `agent()` (or a plain return) that combines results into `{ "ok": ..., ... }`. diff --git a/tests/tools/test_workflow_tool.py b/tests/tools/test_workflow_tool.py new file mode 100644 index 00000000..c912df49 --- /dev/null +++ b/tests/tools/test_workflow_tool.py @@ -0,0 +1,145 @@ +import json + +import pytest + +from pythinker_code.tools.workflow import Workflow + + +class FakeApproval: + def __init__(self, approved=True): + self._approved = approved + self._orchestrations = set() + self.requests = 0 + + def is_orchestration_approved(self, fp): + return fp in self._orchestrations + + def approve_orchestration(self, fp): + self._orchestrations.add(fp) + + async def request(self, sender, action, description): + self.requests += 1 + return _ApprovalResult(self._approved) + + +class _ApprovalResult: + def __init__(self, approved): + self._approved = approved + + def __bool__(self): + return self._approved + + def rejection_error(self): + from pythinker_core.tooling import ToolError + + return ToolError(message="rejected", brief="rejected") + + +class FakeAgentTool: + """Stands in for AgentTool: returns the runner-format `[summary]` output.""" + + def __init__(self, responder): + self._responder = responder + self.calls = [] + + async def __call__(self, params): + from pythinker_core.tooling import ToolOk + + self.calls.append(params) + body = self._responder(params) + return ToolOk(output=f"agent_id: x\nstatus: completed\n\n[summary]\n{body}") + + +def make_tool(monkeypatch, *, responder, approved=True, role="root"): + approval = FakeApproval(approved) + + class FakeConfig: + class background: + max_running_tasks = 4 + + class FakeRuntime: + def __init__(self): + self.role = role + self.approval = approval + self.config = FakeConfig() + + class _WD: + def __str__(self): + return "." + + self.work_dir = _WD() + + runtime = FakeRuntime() + # Patch AgentTool so the Workflow tool wires our fake instead of the real one. + import pythinker_code.tools.workflow as mod + + monkeypatch.setattr(mod, "AgentTool", lambda rt: FakeAgentTool(responder)) + # Suppress wire emission (no Wire ContextVar in a unit test). + monkeypatch.setattr(mod, "wire_send", lambda *a, **k: None) + tool = Workflow(runtime) + return tool, approval + + +@pytest.mark.asyncio +async def test_runs_and_returns_result(monkeypatch): + tool, approval = make_tool(monkeypatch, responder=lambda p: "SUMMARY_TEXT") + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("inspect repo", {"label": "L"})\n' + 'return {"r": r}\n' + ) + res = await tool(tool.params(script=script)) + assert not res.is_error + assert json.loads(res.output) == {"r": "SUMMARY_TEXT"} + assert approval.requests == 1 # approval requested exactly once + + +@pytest.mark.asyncio +async def test_root_only_guard(monkeypatch): + tool, _ = make_tool(monkeypatch, responder=lambda p: "x", role="subagent") + res = await tool( + tool.params(script='meta = {"name": "n", "description": "d"}\nawait agent("x")\n') + ) + assert res.is_error + assert "root" in res.message.lower() + + +@pytest.mark.asyncio +async def test_rejected_approval_short_circuits(monkeypatch): + tool, _ = make_tool(monkeypatch, responder=lambda p: "x", approved=False) + res = await tool( + tool.params(script='meta = {"name": "n", "description": "d"}\nawait agent("x")\n') + ) + assert res.is_error + + +@pytest.mark.asyncio +async def test_no_agent_call_errors(monkeypatch): + tool, _ = make_tool(monkeypatch, responder=lambda p: "x") + res = await tool( + tool.params(script='meta = {"name": "n", "description": "d"}\nreturn {"x": 1}\n') + ) + assert res.is_error + assert "agent()" in res.message + + +@pytest.mark.asyncio +async def test_schema_validation_and_retry(monkeypatch): + attempts = {"n": 0} + + def responder(params): + attempts["n"] += 1 + if attempts["n"] == 1: + return "not json" # first attempt fails validation + return '```json\n{"paths": ["a.py"]}\n```' + + tool, _ = make_tool(monkeypatch, responder=responder) + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("find files", {"schema": {"type": "object", ' + '"properties": {"paths": {"type": "array"}}, "required": ["paths"]}})\n' + "return r\n" + ) + res = await tool(tool.params(script=script)) + assert json.loads(res.output) == {"paths": ["a.py"]} + assert attempts["n"] == 2 # retried once diff --git a/tests/utils/test_pyinstaller_utils.py b/tests/utils/test_pyinstaller_utils.py index 2e89e860..d4053591 100644 --- a/tests/utils/test_pyinstaller_utils.py +++ b/tests/utils/test_pyinstaller_utils.py @@ -280,6 +280,10 @@ def test_pyinstaller_datas(): "src/pythinker_code/tools/web/search.md", "pythinker_code/tools/web", ), + ( + "src/pythinker_code/tools/workflow/description.md", + "pythinker_code/tools/workflow", + ), ( "src/pythinker_code/tools/worktree/enter_worktree.md", "pythinker_code/tools/worktree", @@ -373,6 +377,9 @@ def test_pyinstaller_hiddenimports(): "pythinker_code.tools.web._allowlist", "pythinker_code.tools.web.fetch", "pythinker_code.tools.web.search", + "pythinker_code.tools.workflow", + "pythinker_code.tools.workflow.display", + "pythinker_code.tools.workflow.engine", "pythinker_code.tools.worktree", "setproctitle", ] From 6c97265ffe6016f3bb5a33299ba6ca6c6e43f09d Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 14:38:55 -0400 Subject: [PATCH 07/14] feat(workflow): register Workflow tool in the default agent spec --- src/pythinker_code/agents/default/agent.yaml | 1 + tests/tools/test_workflow_registration.py | 27 ++++++++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 tests/tools/test_workflow_registration.py diff --git a/src/pythinker_code/agents/default/agent.yaml b/src/pythinker_code/agents/default/agent.yaml index 8dee0aba..6898ab37 100644 --- a/src/pythinker_code/agents/default/agent.yaml +++ b/src/pythinker_code/agents/default/agent.yaml @@ -24,6 +24,7 @@ agent: - "pythinker_code.tools.worktree:ExitWorktree" - "pythinker_code.tools.goal:UpdateGoal" - "pythinker_code.tools.progress:Progress" + - "pythinker_code.tools.workflow:Workflow" - "pythinker_code.tools.suggest:Suggest" - "pythinker_code.tools.memory:Memory" - "pythinker_code.tools.recall:Recall" diff --git a/tests/tools/test_workflow_registration.py b/tests/tools/test_workflow_registration.py new file mode 100644 index 00000000..a86bc488 --- /dev/null +++ b/tests/tools/test_workflow_registration.py @@ -0,0 +1,27 @@ +import importlib + + +def test_workflow_tool_string_resolves(): + mod = importlib.import_module("pythinker_code.tools.workflow") + assert hasattr(mod, "Workflow") + assert mod.Workflow.name == "Workflow" + + +def test_workflow_registered_in_default_spec(): + from pathlib import Path + + import yaml + + spec_path = ( + Path(__file__).resolve().parents[2] + / "src/pythinker_code/agents/default/agent.yaml" + ) + spec = yaml.safe_load(spec_path.read_text(encoding="utf-8")) + tools = spec["agent"]["tools"] + assert "pythinker_code.tools.workflow:Workflow" in tools + # Must NOT be wired into child specs (no nesting). + coder = yaml.safe_load( + (spec_path.parent / "coder.yaml").read_text(encoding="utf-8") + ) + coder_tools = coder.get("agent", {}).get("tools", []) if isinstance(coder, dict) else [] + assert "pythinker_code.tools.workflow:Workflow" not in coder_tools From dadc3c8732e1797c808fee65a1275fdd463a55dc Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 14:42:59 -0400 Subject: [PATCH 08/14] test(workflow): assert Workflow absence from coder's resolved tool list --- tests/tools/test_workflow_registration.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/tools/test_workflow_registration.py b/tests/tools/test_workflow_registration.py index a86bc488..7db13911 100644 --- a/tests/tools/test_workflow_registration.py +++ b/tests/tools/test_workflow_registration.py @@ -19,9 +19,15 @@ def test_workflow_registered_in_default_spec(): spec = yaml.safe_load(spec_path.read_text(encoding="utf-8")) tools = spec["agent"]["tools"] assert "pythinker_code.tools.workflow:Workflow" in tools - # Must NOT be wired into child specs (no nesting). - coder = yaml.safe_load( - (spec_path.parent / "coder.yaml").read_text(encoding="utf-8") + # Must NOT be wired into child specs (no nesting). Resolve the coder subagent + # spec through the real agentspec loader (extend + allowed_tools applied), + # not raw YAML — coder.yaml has no literal `tools:` key, so a raw-YAML + # check would vacuously pass regardless of what allowed_tools contains. + from pythinker_code.agentspec import DEFAULT_AGENT_FILE, load_agent_spec + + root_spec = load_agent_spec(DEFAULT_AGENT_FILE) + coder_spec = load_agent_spec(root_spec.subagents["coder"].path) + effective_coder_tools = ( + coder_spec.allowed_tools if coder_spec.allowed_tools is not None else coder_spec.tools ) - coder_tools = coder.get("agent", {}).get("tools", []) if isinstance(coder, dict) else [] - assert "pythinker_code.tools.workflow:Workflow" not in coder_tools + assert "pythinker_code.tools.workflow:Workflow" not in effective_coder_tools From 218c771dc2686f07d7086ecacc8c676119c60388 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 16:56:41 -0400 Subject: [PATCH 09/14] chore(workflow): changelog entry, abort-teardown test, and snapshot updates Adds the Unreleased changelog bullet for the Workflow tool and a real abort-teardown integration test proving cancellation propagates through AgentTool -> ForegroundSubagentRunner and marks the child instance "killed" (not just the engine's own bookkeeping). Registering Workflow in agent.yaml grew the default agent spec's tools list, which moved pinned tool-list snapshots in test_agent_spec.py (root spec + 5 inheriting child specs) and test_default_agent.py (the loaded toolset's tool names) -- all fixed as a single additive line each. test_workflow_parser.py/test_workflow_registration.py/test_workflow_tool.py also picked up ruff-format reflow from the full-repo gate. --- CHANGELOG.md | 4 + tests/core/test_agent_spec.py | 6 ++ tests/core/test_default_agent.py | 1 + .../tools/test_workflow_abort_integration.py | 92 +++++++++++++++++++ tests/tools/test_workflow_parser.py | 14 ++- tests/tools/test_workflow_registration.py | 5 +- tests/tools/test_workflow_tool.py | 8 +- 7 files changed, 121 insertions(+), 9 deletions(-) create mode 100644 tests/tools/test_workflow_abort_integration.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b3085f82..0ead08fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,10 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- Add a built-in `Workflow` tool that runs a deterministic Python script orchestrating + multiple subagents via `agent()`, `parallel()`, and `pipeline()`, with live progress, + structured-schema output, and one-shot approval. + ## 0.53.0 (2026-06-23) - **GPT/ChatGPT session re-authentication crash fixed.** A fatal "Could not diff --git a/tests/core/test_agent_spec.py b/tests/core/test_agent_spec.py index 2732f8ed..35cdd433 100644 --- a/tests/core/test_agent_spec.py +++ b/tests/core/test_agent_spec.py @@ -42,6 +42,7 @@ def test_load_default_agent_spec(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", @@ -255,6 +256,7 @@ def test_load_default_agent_spec(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", @@ -389,6 +391,7 @@ def test_load_default_agent_spec(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", @@ -533,6 +536,7 @@ def test_load_default_agent_spec(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", @@ -662,6 +666,7 @@ def test_load_default_agent_spec(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", @@ -839,6 +844,7 @@ def test_load_agent_spec_default_extension(): "pythinker_code.tools.worktree:ExitWorktree", "pythinker_code.tools.goal:UpdateGoal", "pythinker_code.tools.progress:Progress", + "pythinker_code.tools.workflow:Workflow", "pythinker_code.tools.suggest:Suggest", "pythinker_code.tools.memory:Memory", "pythinker_code.tools.recall:Recall", diff --git a/tests/core/test_default_agent.py b/tests/core/test_default_agent.py index a3e77d6f..7cc2352b 100644 --- a/tests/core/test_default_agent.py +++ b/tests/core/test_default_agent.py @@ -318,6 +318,7 @@ async def test_default_agent_background_bash_guardrails(runtime: Runtime): "ExitWorktree", "UpdateGoal", "Progress", + "Workflow", "Suggest", "Memory", "Recall", diff --git a/tests/tools/test_workflow_abort_integration.py b/tests/tools/test_workflow_abort_integration.py new file mode 100644 index 00000000..3fb0e989 --- /dev/null +++ b/tests/tools/test_workflow_abort_integration.py @@ -0,0 +1,92 @@ +"""Integration test proving Workflow's cancellation tears down a real child agent. + +``test_cancellation_marks_running_skipped_and_reraises`` in +``tests/tools/test_workflow_engine.py`` only proves the engine's own bookkeeping with a +fake ``agent_runner``. This test exercises the real spawn layer one level up: cancelling +the ``Workflow`` tool's task must propagate through ``AgentTool`` -> +``ForegroundSubagentRunner`` and mark the child instance "killed" in the subagent store, +the same pattern used by +``test_agent_tool.py::test_agent_tool_marks_instance_killed_when_initial_run_is_cancelled``. +""" + +from __future__ import annotations + +import asyncio + +import pytest +from pythinker_core.tooling.empty import EmptyToolset + +from pythinker_code.soul.agent import Agent as SoulAgent +from pythinker_code.subagents.models import AgentTypeDefinition, ToolPolicy +from pythinker_code.tools.workflow import Workflow +from tests.conftest import tool_call_context + + +class _RecordingWire: + """Minimal wire stub so Workflow's progress emits (wire_send) don't assert-fail. + + Workflow emits ProgressNote updates via wire_send on phase/agent-start/agent-end + hooks, which asserts a wire is set. Outside the full soul loop there is none, so + this stub (the same pattern as test_agent_tool.py's ``_RecordingWire``) stands in. + """ + + def __init__(self) -> None: + self.soul_side = self + + def send(self, msg: object) -> None: + pass + + +@pytest.mark.asyncio +async def test_workflow_cancel_tears_down_real_child_through_agent_tool(runtime, monkeypatch): + monkeypatch.setattr("pythinker_code.soul.get_wire_or_none", lambda: _RecordingWire()) + runtime.labor_market.add_builtin_type( + AgentTypeDefinition( + name="coder", + description="Good at general software engineering tasks.", + agent_file=runtime.subagent_store.root / "coder.yaml", + tool_policy=ToolPolicy(mode="inherit"), + ) + ) + + async def fake_load_agent(agent_file, runtime, *, mcp_configs, start_mcp_loading=True): + return SoulAgent( + name=agent_file.stem, + system_prompt="Subagent system prompt", + toolset=EmptyToolset(), + runtime=runtime, + ) + + started = asyncio.Event() + + async def fake_run_soul( + soul, user_input, ui_loop_fn, cancel_event, wire_file=None, runtime=None + ): + started.set() + await asyncio.sleep(10.0) # never resolves on its own; must be cancelled + raise AssertionError("fake_run_soul should have been cancelled, not completed") + + monkeypatch.setattr("pythinker_code.subagents.builder.load_agent", fake_load_agent) + monkeypatch.setattr("pythinker_code.subagents.runner.run_soul", fake_run_soul) + + tool = Workflow(runtime) + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("investigate bug", {"label": "L"})\n' + 'return {"r": r}\n' + ) + with tool_call_context("Workflow"): + task = asyncio.create_task(tool(tool.params(script=script))) + await asyncio.wait_for(started.wait(), timeout=2.0) # let the real child actually start + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + records = [ + r + for r in runtime.subagent_store.list_instances() + if r.description in ("L", "investigate bug") + ] + assert len(records) == 1 + assert records[0].status == "killed" # real AgentTool/ForegroundSubagentRunner teardown ran, + # not just the engine's own bookkeeping — this is the load-bearing assertion. diff --git a/tests/tools/test_workflow_parser.py b/tests/tools/test_workflow_parser.py index 1518bb45..eea5e4ca 100644 --- a/tests/tools/test_workflow_parser.py +++ b/tests/tools/test_workflow_parser.py @@ -5,11 +5,11 @@ parse_workflow_script, ) -GOOD = '''meta = {"name": "inspect", "description": "Inspect repo", "phases": [{"title": "Scan"}]} +GOOD = """meta = {"name": "inspect", "description": "Inspect repo", "phases": [{"title": "Scan"}]} phase("Scan") inventory = await agent("Inspect the repository.", {"label": "repo inventory"}) return {"inventory": inventory} -''' +""" def test_parse_accepts_valid_script(): @@ -31,8 +31,14 @@ def test_parse_accepts_valid_script(): ('meta = {"name": "n", "description": "d"}\nimport os\n', "not allowed"), ('meta = {"name": "n", "description": "d"}\nx = random.random()\n', "deterministic"), ('meta = {"name": "n", "description": "d"}\nx = time.time()\n', "deterministic"), - ('meta = {"name": "n", "description": "d", "when_to_use": 42}\nawait agent("x")\n', "when_to_use"), - ('meta = {"name": "n", "description": "d", "phases": "oops"}\nawait agent("x")\n', "phases"), + ( + 'meta = {"name": "n", "description": "d", "when_to_use": 42}\nawait agent("x")\n', + "when_to_use", + ), + ( + 'meta = {"name": "n", "description": "d", "phases": "oops"}\nawait agent("x")\n', + "phases", + ), ('meta = {"name": "n", "description": "d", "phases": [42]}\nawait agent("x")\n', "title"), ], ) diff --git a/tests/tools/test_workflow_registration.py b/tests/tools/test_workflow_registration.py index 7db13911..9720476a 100644 --- a/tests/tools/test_workflow_registration.py +++ b/tests/tools/test_workflow_registration.py @@ -12,10 +12,7 @@ def test_workflow_registered_in_default_spec(): import yaml - spec_path = ( - Path(__file__).resolve().parents[2] - / "src/pythinker_code/agents/default/agent.yaml" - ) + spec_path = Path(__file__).resolve().parents[2] / "src/pythinker_code/agents/default/agent.yaml" spec = yaml.safe_load(spec_path.read_text(encoding="utf-8")) tools = spec["agent"]["tools"] assert "pythinker_code.tools.workflow:Workflow" in tools diff --git a/tests/tools/test_workflow_tool.py b/tests/tools/test_workflow_tool.py index c912df49..c3ff3d5d 100644 --- a/tests/tools/test_workflow_tool.py +++ b/tests/tools/test_workflow_tool.py @@ -1,7 +1,11 @@ +from __future__ import annotations + import json +from typing import cast import pytest +from pythinker_code.soul.agent import Runtime from pythinker_code.tools.workflow import Workflow @@ -76,7 +80,7 @@ def __str__(self): monkeypatch.setattr(mod, "AgentTool", lambda rt: FakeAgentTool(responder)) # Suppress wire emission (no Wire ContextVar in a unit test). monkeypatch.setattr(mod, "wire_send", lambda *a, **k: None) - tool = Workflow(runtime) + tool = Workflow(cast(Runtime, runtime)) return tool, approval @@ -90,6 +94,7 @@ async def test_runs_and_returns_result(monkeypatch): ) res = await tool(tool.params(script=script)) assert not res.is_error + assert isinstance(res.output, str) assert json.loads(res.output) == {"r": "SUMMARY_TEXT"} assert approval.requests == 1 # approval requested exactly once @@ -141,5 +146,6 @@ def responder(params): "return r\n" ) res = await tool(tool.params(script=script)) + assert isinstance(res.output, str) assert json.loads(res.output) == {"paths": ["a.py"]} assert attempts["n"] == 2 # retried once From f82abeeddd4457e10bcf9fbff229101dbde03495 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 17:23:40 -0400 Subject: [PATCH 10/14] feat(workflow): add settable token_budget param so budget actually enforces --- src/pythinker_code/tools/workflow/__init__.py | 13 +++++ .../tools/workflow/description.md | 9 +++- tests/tools/test_workflow_engine.py | 2 +- tests/tools/test_workflow_tool.py | 48 +++++++++++++++++++ 4 files changed, 69 insertions(+), 3 deletions(-) diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py index e597b314..5af29355 100644 --- a/src/pythinker_code/tools/workflow/__init__.py +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -44,6 +44,18 @@ class Params(BaseModel): default=None, description="Optional JSON value exposed to the script as the global `args`.", ) + token_budget: int | None = Field( + default=None, + description=( + "Optional cap on estimated total tokens spent by spawned subagents. When set, " + "the script's `budget.remaining()` reaches 0 once the cap is hit and further " + "agent() calls raise inside the engine: within parallel(), that one call is " + "caught and returns None (logged), but a bare `await agent(...)` propagates " + "and fails the whole workflow run. Leave unset for no cap (budget.remaining() " + "stays unbounded)." + ), + ge=1, + ) class Workflow(CallableTool2[Params]): @@ -117,6 +129,7 @@ def on_agent_end(event: AgentEndEvent) -> None: args=params.args, cwd=str(self._runtime.work_dir), concurrency=max(1, self._runtime.config.background.max_running_tasks), + token_budget=params.token_budget, hooks=hooks, ) except WorkflowScriptError as exc: diff --git a/src/pythinker_code/tools/workflow/description.md b/src/pythinker_code/tools/workflow/description.md index 69a6b49a..2332acf3 100644 --- a/src/pythinker_code/tools/workflow/description.md +++ b/src/pythinker_code/tools/workflow/description.md @@ -1,6 +1,11 @@ Execute a deterministic Python workflow that orchestrates multiple subagents with `agent()`, `parallel()`, and `pipeline()`. Use this only when the user explicitly asks for a workflow, fan-out, or multi-agent orchestration, or when a task decomposes into many independent subtasks worth running concurrently and synthesizing. -`script` is required raw Python (no Markdown fences). Rules: +`script` is required raw Python (no Markdown fences). `args` (optional) is a JSON value exposed to +the script as the global `args`. `token_budget` (optional) caps estimated total token spend across +all spawned subagents — set it when you want the workflow to self-throttle via +`budget.remaining()`; leave it unset for no cap. + +Rules: - The first statement MUST be `meta = {"name": "short_snake_case", "description": "non-empty description"}`. `meta` must be a literal dict (no function calls or interpolation). `meta["phases"]` is optional documentation; live progress is driven by `phase(title)` at runtime. - After `meta`, write plain Python. Do NOT use `import`, `time`, `random`, `datetime`, `os`, `sys`, `open`, `eval`, or `exec` — scripts must be deterministic. @@ -12,6 +17,6 @@ Available globals: - `parallel([awaitables])` — run awaitables concurrently, results in input order: `await parallel([agent("a"), agent("b")])`. Pass awaitables, NOT functions. - `pipeline(items, *stages)` — run each item through sequential stages while items fan out. Each stage is called `(prev, original, index)` and may be sync or async. - `phase(title)` — start a progress group. Names may be conditional or built in a loop; do not predeclare speculative phases. -- `log(message)`, `args` (the optional JSON `args` input), `cwd`, `budget` (`.total`, `.spent()`, `.remaining()`). +- `log(message)`, `args` (the optional JSON `args` input), `cwd`, `budget` (`.total`, `.spent()`, `.remaining()` — reflects the `token_budget` tool parameter; unbounded if it was not set). Include enough context and file paths in each `agent()` prompt — subagents do NOT inherit the parent conversation. Add a final synthesis `agent()` (or a plain return) that combines results into `{ "ok": ..., ... }`. diff --git a/tests/tools/test_workflow_engine.py b/tests/tools/test_workflow_engine.py index 6b5d6f7c..c37c75a6 100644 --- a/tests/tools/test_workflow_engine.py +++ b/tests/tools/test_workflow_engine.py @@ -159,7 +159,7 @@ async def runner(prompt: str, opts: AgentOptions) -> str: @pytest.mark.asyncio -async def test_cancellation_marks_running_skipped_and_reraises(): +async def test_cancellation_marks_running_cancelled_and_reraises(): runner, _ = make_runner(delay=10.0) skipped: list[str] = [] started: list[str] = [] diff --git a/tests/tools/test_workflow_tool.py b/tests/tools/test_workflow_tool.py index c3ff3d5d..3030664d 100644 --- a/tests/tools/test_workflow_tool.py +++ b/tests/tools/test_workflow_tool.py @@ -128,6 +128,54 @@ async def test_no_agent_call_errors(monkeypatch): assert "agent()" in res.message +@pytest.mark.asyncio +async def test_token_budget_param_reaches_run_workflow(monkeypatch): + captured: dict[str, object] = {} + import pythinker_code.tools.workflow as mod + + real_run_workflow = mod.run_workflow + + async def spy_run_workflow(*args, **kwargs): + captured.update(kwargs) + return await real_run_workflow(*args, **kwargs) + + monkeypatch.setattr(mod, "run_workflow", spy_run_workflow) + + tool, _ = make_tool(monkeypatch, responder=lambda p: "ok") + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("x", {"label": "L"})\n' + 'return {"r": r}\n' + ) + res = await tool(tool.params(script=script, token_budget=12345)) + assert not res.is_error + assert captured.get("token_budget") == 12345 + + +@pytest.mark.asyncio +async def test_token_budget_defaults_to_none(monkeypatch): + captured: dict[str, object] = {} + import pythinker_code.tools.workflow as mod + + real_run_workflow = mod.run_workflow + + async def spy_run_workflow(*args, **kwargs): + captured.update(kwargs) + return await real_run_workflow(*args, **kwargs) + + monkeypatch.setattr(mod, "run_workflow", spy_run_workflow) + + tool, _ = make_tool(monkeypatch, responder=lambda p: "ok") + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r = await agent("x", {"label": "L"})\n' + 'return {"r": r}\n' + ) + res = await tool(tool.params(script=script)) + assert not res.is_error + assert captured.get("token_budget") is None + + @pytest.mark.asyncio async def test_schema_validation_and_retry(monkeypatch): attempts = {"n": 0} From c689dbc689248299ed34165a59b0b00252dc4349 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 17:24:11 -0400 Subject: [PATCH 11/14] docs(workflow): fix stale test-name reference in abort-integration test docstring --- tests/tools/test_workflow_abort_integration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tools/test_workflow_abort_integration.py b/tests/tools/test_workflow_abort_integration.py index 3fb0e989..ebf152de 100644 --- a/tests/tools/test_workflow_abort_integration.py +++ b/tests/tools/test_workflow_abort_integration.py @@ -1,6 +1,6 @@ """Integration test proving Workflow's cancellation tears down a real child agent. -``test_cancellation_marks_running_skipped_and_reraises`` in +``test_cancellation_marks_running_cancelled_and_reraises`` in ``tests/tools/test_workflow_engine.py`` only proves the engine's own bookkeeping with a fake ``agent_runner``. This test exercises the real spawn layer one level up: cancelling the ``Workflow`` tool's task must propagate through ``AgentTool`` -> From 44b7cff79702258d2dc1bb4907e9f9e174d5e842 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 17:28:49 -0400 Subject: [PATCH 12/14] docs(workflow): note pipeline() shares parallel()'s budget-exhaustion catch --- src/pythinker_code/tools/workflow/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py index 5af29355..93327d13 100644 --- a/src/pythinker_code/tools/workflow/__init__.py +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -49,8 +49,8 @@ class Params(BaseModel): description=( "Optional cap on estimated total tokens spent by spawned subagents. When set, " "the script's `budget.remaining()` reaches 0 once the cap is hit and further " - "agent() calls raise inside the engine: within parallel(), that one call is " - "caught and returns None (logged), but a bare `await agent(...)` propagates " + "agent() calls raise inside the engine: within parallel()/pipeline(), that one " + "call is caught and returns None (logged), but a bare `await agent(...)` propagates " "and fails the whole workflow run. Leave unset for no cap (budget.remaining() " "stays unbounded)." ), From 8ec4e8db5902149ff3ca0edc9709b26a8492d2a5 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 17:41:09 -0400 Subject: [PATCH 13/14] fix(workflow): drop inaccurate 'declared phases' claim from no-agents error --- src/pythinker_code/tools/workflow/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py index 93327d13..ae5696bd 100644 --- a/src/pythinker_code/tools/workflow/__init__.py +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -145,7 +145,7 @@ def on_agent_end(event: AgentEndEvent) -> None: if result.agent_count == 0: return ToolError( message="workflow scripts must call agent() at least once; this workflow " - "declared phases but ran no subagents.", + "ran no subagents.", brief="No agents run", ) From 9553cb21a75141e2073c11c2ce04d82b4688dc6a Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 30 Jun 2026 17:57:33 -0400 Subject: [PATCH 14/14] fix: address PR review findings --- src/pythinker_code/tools/workflow/__init__.py | 10 ++- src/pythinker_code/tools/workflow/display.py | 18 +++-- src/pythinker_code/tools/workflow/engine.py | 55 ++++++++++----- .../tools/test_workflow_abort_integration.py | 2 +- tests/tools/test_workflow_display.py | 46 +++++++++++-- tests/tools/test_workflow_engine.py | 69 +++++++++++++++++++ tests/tools/test_workflow_tool.py | 62 +++++++++-------- 7 files changed, 200 insertions(+), 62 deletions(-) diff --git a/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py index ae5696bd..76dac86f 100644 --- a/src/pythinker_code/tools/workflow/__init__.py +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -103,20 +103,24 @@ async def __call__(self, params: Params) -> ToolReturnValue: def emit() -> None: wire_send(ProgressNote(title=f"Workflow {meta.name}", body=render_progress(snapshot))) + def on_log(message: str) -> None: + snapshot.logs.append(message) + emit() + def on_phase(title: str) -> None: snapshot.add_phase(title) emit() def on_agent_start(event: AgentStartEvent) -> None: - snapshot.start_agent(event.label, event.phase) + snapshot.start_agent(event.agent_id, event.label, event.phase) emit() def on_agent_end(event: AgentEndEvent) -> None: - snapshot.end_agent(event.label, error=event.error) + snapshot.end_agent(event.agent_id, error=event.error) emit() hooks = RunWorkflowHooks( - on_log=lambda _m: emit(), + on_log=on_log, on_phase=on_phase, on_agent_start=on_agent_start, on_agent_end=on_agent_end, diff --git a/src/pythinker_code/tools/workflow/display.py b/src/pythinker_code/tools/workflow/display.py index 126558f8..abd1e8e6 100644 --- a/src/pythinker_code/tools/workflow/display.py +++ b/src/pythinker_code/tools/workflow/display.py @@ -37,15 +37,19 @@ def add_phase(self, title: str | None) -> None: if title not in self.phases: self.phases.append(title) - def start_agent(self, label: str, phase: str | None) -> AgentSnapshot: + def start_agent(self, agent_id: int, label: str, phase: str | None) -> AgentSnapshot: self.add_phase(phase) - agent = AgentSnapshot(len(self.agents) + 1, label, phase) + agent = AgentSnapshot(agent_id, label, phase) self.agents.append(agent) return agent - def end_agent(self, label: str, *, error: str | None = None) -> None: - for agent in reversed(self.agents): - if agent.label == label and agent.status == "running": + def end_agent(self, agent_id: int, *, error: str | None = None) -> None: + # Matched by the engine's stable dispatch id, not by label: two + # concurrent agents may share the same caller-supplied label, and a + # label-based reverse search can mark the wrong entry done/error if + # they finish out of dispatch order. + for agent in self.agents: + if agent.id == agent_id and agent.status == "running": agent.status = "error" if error else "done" return @@ -71,7 +75,7 @@ def skipped_count(self) -> int: return sum(1 for a in self.agents if a.status == "skipped") -def render_progress(snapshot: WorkflowSnapshot, max_agents: int = 6) -> str: +def render_progress(snapshot: WorkflowSnapshot, max_agents: int = 6, max_logs: int = 3) -> str: state = "" if snapshot.error_count: state = f", {snapshot.error_count} errors" @@ -96,4 +100,6 @@ def render_progress(snapshot: WorkflowSnapshot, max_agents: int = 6) -> str: unphased = [a for a in snapshot.agents if a.id not in rendered] for agent in unphased[-max_agents:]: lines.append(f" #{agent.id} {_STATUS_ICON.get(agent.status, '?')} {agent.label}") + for message in snapshot.logs[-max_logs:]: + lines.append(f" log: {message}") return "\n".join(lines) diff --git a/src/pythinker_code/tools/workflow/engine.py b/src/pythinker_code/tools/workflow/engine.py index 32a3b256..461ffa35 100644 --- a/src/pythinker_code/tools/workflow/engine.py +++ b/src/pythinker_code/tools/workflow/engine.py @@ -210,20 +210,27 @@ def __init__( class AgentStartEvent: - __slots__ = ("label", "phase", "prompt") + __slots__ = ("agent_id", "label", "phase", "prompt") - def __init__(self, label: str, phase: str | None, prompt: str) -> None: + def __init__(self, agent_id: int, label: str, phase: str | None, prompt: str) -> None: + self.agent_id = agent_id self.label = label self.phase = phase self.prompt = prompt class AgentEndEvent: - __slots__ = ("label", "phase", "result", "error") + __slots__ = ("agent_id", "label", "phase", "result", "error") def __init__( - self, label: str, phase: str | None, result: Any, error: str | None = None + self, + agent_id: int, + label: str, + phase: str | None, + result: Any, + error: str | None = None, ) -> None: + self.agent_id = agent_id self.label = label self.phase = phase self.result = result @@ -270,18 +277,28 @@ def _require_str(value: Any, name: str) -> str: return value +def _optional_str_field(d: dict[str, Any], name: str) -> str | None: + value = d.get(name) + if value is not None and not isinstance(value, str): + raise WorkflowRuntimeError(f"agent() option '{name}' must be a string or None") + return value + + def _normalize_options(value: Any) -> AgentOptions: if value is None: return AgentOptions() if not isinstance(value, dict): raise WorkflowRuntimeError("agent options must be a dict") d = cast(dict[str, Any], value) + schema = d.get("schema") + if schema is not None and not isinstance(schema, dict): + raise WorkflowRuntimeError("agent() option 'schema' must be a dict or None") return AgentOptions( - label=d.get("label"), - phase=d.get("phase"), - schema=d.get("schema"), - model=d.get("model"), - agent_type=d.get("agent_type") or d.get("agentType"), + label=_optional_str_field(d, "label"), + phase=_optional_str_field(d, "phase"), + schema=cast("dict[str, Any] | None", schema), + model=_optional_str_field(d, "model"), + agent_type=_optional_str_field(d, "agent_type") or _optional_str_field(d, "agentType"), ) @@ -390,29 +407,35 @@ async def agent(prompt: Any, options: Any = None) -> Any: if token_budget is not None and budget.remaining() <= 0: raise WorkflowRuntimeError("workflow token budget exhausted") state["agent_count"] += 1 - label = (opts.label or "").strip() or _default_label( - assigned_phase, state["agent_count"] - ) + # Captured into a local now, before the first `await` below: other + # concurrently-dispatched agent() calls can advance + # state["agent_count"] further while this call awaits + # agent_runner, so re-reading the shared counter later would pick + # up the wrong (now-higher) value. + agent_id = state["agent_count"] + label = (opts.label or "").strip() or _default_label(assigned_phase, agent_id) opts.label = label opts.phase = assigned_phase if hooks.on_agent_start: - hooks.on_agent_start(AgentStartEvent(label, assigned_phase, task_prompt)) + hooks.on_agent_start(AgentStartEvent(agent_id, label, assigned_phase, task_prompt)) try: result = await agent_runner(task_prompt, opts) except asyncio.CancelledError: if hooks.on_agent_end: hooks.on_agent_end( - AgentEndEvent(label, assigned_phase, None, error="cancelled") + AgentEndEvent(agent_id, label, assigned_phase, None, error="cancelled") ) raise except Exception as exc: # noqa: BLE001 - reference parity: branch fails to None log(f"agent {label} failed: {exc}") if hooks.on_agent_end: - hooks.on_agent_end(AgentEndEvent(label, assigned_phase, None, error=str(exc))) + hooks.on_agent_end( + AgentEndEvent(agent_id, label, assigned_phase, None, error=str(exc)) + ) return None state["spent"] += _estimate_tokens(result) if hooks.on_agent_end: - hooks.on_agent_end(AgentEndEvent(label, assigned_phase, result)) + hooks.on_agent_end(AgentEndEvent(agent_id, label, assigned_phase, result)) return result async def parallel(items: Sequence[Any]) -> list[Any]: diff --git a/tests/tools/test_workflow_abort_integration.py b/tests/tools/test_workflow_abort_integration.py index ebf152de..657d0cb9 100644 --- a/tests/tools/test_workflow_abort_integration.py +++ b/tests/tools/test_workflow_abort_integration.py @@ -39,7 +39,7 @@ def send(self, msg: object) -> None: @pytest.mark.asyncio async def test_workflow_cancel_tears_down_real_child_through_agent_tool(runtime, monkeypatch): - monkeypatch.setattr("pythinker_code.soul.get_wire_or_none", lambda: _RecordingWire()) + monkeypatch.setattr("pythinker_code.soul.get_wire_or_none", _RecordingWire) runtime.labor_market.add_builtin_type( AgentTypeDefinition( name="coder", diff --git a/tests/tools/test_workflow_display.py b/tests/tools/test_workflow_display.py index 940ab9a9..fedb2cd0 100644 --- a/tests/tools/test_workflow_display.py +++ b/tests/tools/test_workflow_display.py @@ -4,14 +4,14 @@ def test_snapshot_lifecycle_and_render(): snap = WorkflowSnapshot(name="inspect", description="d") snap.add_phase("Scan") - a = snap.start_agent("repo inventory", "Scan") + a = snap.start_agent(1, "repo inventory", "Scan") assert a.status == "running" assert snap.running_count == 1 - snap.end_agent("repo inventory") + snap.end_agent(1) assert snap.done_count == 1 and snap.running_count == 0 snap.add_phase("Analyze") - snap.start_agent("modules", "Analyze") + snap.start_agent(2, "modules", "Analyze") text = render_progress(snap) assert "Workflow: inspect" in text assert "Scan" in text and "Analyze" in text @@ -20,9 +20,43 @@ def test_snapshot_lifecycle_and_render(): def test_mark_running_skipped(): snap = WorkflowSnapshot(name="n", description="d") - snap.start_agent("a", None) - snap.start_agent("b", None) - snap.end_agent("a") + snap.start_agent(1, "a", None) + snap.start_agent(2, "b", None) + snap.end_agent(1) snap.mark_running_skipped() assert snap.done_count == 1 assert snap.skipped_count == 1 + + +def test_end_agent_does_not_swap_status_on_duplicate_labels(): + # Regression guard for CodeRabbit finding: two concurrent agents sharing the + # same explicit label must not have their completion statuses swapped when + # the SECOND-started one finishes (errors) before the FIRST-started one. + snap = WorkflowSnapshot(name="n", description="d") + first = snap.start_agent(1, "scan", None) + second = snap.start_agent(2, "scan", None) + snap.end_agent(2, error="boom") # the second-started agent fails first + assert first.status == "running" + assert second.status == "error" + snap.end_agent(1) # the first-started agent finishes after + assert first.status == "done" + assert second.status == "error" + + +def test_render_progress_shows_recent_log_messages(): + snap = WorkflowSnapshot(name="n", description="d") + snap.logs.append("first checkpoint") + snap.logs.append("second checkpoint") + text = render_progress(snap) + assert "first checkpoint" in text + assert "second checkpoint" in text + + +def test_render_progress_truncates_to_max_logs(): + snap = WorkflowSnapshot(name="n", description="d") + for i in range(5): + snap.logs.append(f"log {i}") + text = render_progress(snap, max_logs=2) + assert "log 3" in text + assert "log 4" in text + assert "log 0" not in text diff --git a/tests/tools/test_workflow_engine.py b/tests/tools/test_workflow_engine.py index c37c75a6..f9c240eb 100644 --- a/tests/tools/test_workflow_engine.py +++ b/tests/tools/test_workflow_engine.py @@ -124,6 +124,75 @@ async def test_unawaited_coroutine_in_result_raises(): assert "await" in str(exc.value) +@pytest.mark.asyncio +async def test_malformed_agent_option_types_raise_workflow_runtime_error(): + # Regression guard: a malformed agent() option (e.g. a non-string label) + # must surface as a clean WorkflowRuntimeError, not an internal + # AttributeError leaking from deep inside the engine. + runner, _ = make_runner() + script = ( + 'meta = {"name": "n", "description": "d"}\nr = await agent("x", {"label": 1})\nreturn r\n' + ) + with pytest.raises(WorkflowRuntimeError) as exc: + await run_workflow(script, agent_runner=runner, cwd=".") + assert "label" in str(exc.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "options, bad_field", + [ + ({"phase": 1}, "phase"), + ({"model": 1}, "model"), + ({"agent_type": 1}, "agent_type"), + ({"schema": "not-a-dict"}, "schema"), + ], +) +async def test_malformed_agent_option_fields_raise_workflow_runtime_error(options, bad_field): + runner, _ = make_runner() + script = ( + 'meta = {"name": "n", "description": "d"}\n' + f"r = await agent('x', {options!r})\n" + "return r\n" + ) + with pytest.raises(WorkflowRuntimeError) as exc: + await run_workflow(script, agent_runner=runner, cwd=".") + assert bad_field in str(exc.value) + + +@pytest.mark.asyncio +async def test_duplicate_labels_do_not_swap_completion_status(): + # Regression guard for a CodeRabbit finding: two agent() calls dispatched + # with the SAME explicit label, where the second-started one finishes + # first (errors), must not have their on_agent_end statuses swapped. + async def runner(prompt: str, opts: AgentOptions) -> str: + if prompt == "second": + raise RuntimeError("boom") # second-started agent fails immediately + await asyncio.sleep(0.05) # first-started agent finishes later + return "ok" + + statuses: dict[int, str | None] = {} + + def on_agent_end(event): + statuses[event.agent_id] = event.error + + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'rs = await parallel([agent("first", {"label": "scan"}), ' + 'agent("second", {"label": "scan"})])\n' + "return rs\n" + ) + out = await run_workflow( + script, + agent_runner=runner, + cwd=".", + concurrency=4, + hooks=RunWorkflowHooks(on_agent_end=on_agent_end), + ) + assert out.result == ["ok", None] + assert statuses == {1: None, 2: "boom"} # id 1 (first) succeeded, id 2 (second) errored + + @pytest.mark.asyncio async def test_budget_check_does_not_race_past_semaphore(): # Regression guard: agent() calls dispatched together (via parallel()) all diff --git a/tests/tools/test_workflow_tool.py b/tests/tools/test_workflow_tool.py index 3030664d..7b94e6f0 100644 --- a/tests/tools/test_workflow_tool.py +++ b/tests/tools/test_workflow_tool.py @@ -7,6 +7,7 @@ from pythinker_code.soul.agent import Runtime from pythinker_code.tools.workflow import Workflow +from pythinker_code.wire.types import ProgressNote class FakeApproval: @@ -75,11 +76,11 @@ def __str__(self): runtime = FakeRuntime() # Patch AgentTool so the Workflow tool wires our fake instead of the real one. - import pythinker_code.tools.workflow as mod - - monkeypatch.setattr(mod, "AgentTool", lambda rt: FakeAgentTool(responder)) + monkeypatch.setattr( + "pythinker_code.tools.workflow.AgentTool", lambda rt: FakeAgentTool(responder) + ) # Suppress wire emission (no Wire ContextVar in a unit test). - monkeypatch.setattr(mod, "wire_send", lambda *a, **k: None) + monkeypatch.setattr("pythinker_code.tools.workflow.wire_send", lambda *a, **k: None) tool = Workflow(cast(Runtime, runtime)) return tool, approval @@ -129,51 +130,52 @@ async def test_no_agent_call_errors(monkeypatch): @pytest.mark.asyncio -async def test_token_budget_param_reaches_run_workflow(monkeypatch): - captured: dict[str, object] = {} - import pythinker_code.tools.workflow as mod - - real_run_workflow = mod.run_workflow - - async def spy_run_workflow(*args, **kwargs): - captured.update(kwargs) - return await real_run_workflow(*args, **kwargs) +async def test_token_budget_exhausted_fails_the_workflow(monkeypatch): + # Drives the tool through its public ToolOk/ToolError surface rather than + # spying on run_workflow's internal kwargs: a tiny token_budget is exhausted + # by the first agent() call's estimated spend, so the second (bare, not + # inside parallel()) agent() call propagates "budget exhausted" and fails + # the whole tool call. + tool, _ = make_tool(monkeypatch, responder=lambda p: "x" * 200) + script = ( + 'meta = {"name": "n", "description": "d"}\n' + 'r1 = await agent("x", {"label": "L1"})\n' + 'r2 = await agent("y", {"label": "L2"})\n' + 'return {"r1": r1, "r2": r2}\n' + ) + res = await tool(tool.params(script=script, token_budget=1)) + assert res.is_error + assert "budget" in res.message.lower() - monkeypatch.setattr(mod, "run_workflow", spy_run_workflow) - tool, _ = make_tool(monkeypatch, responder=lambda p: "ok") +@pytest.mark.asyncio +async def test_token_budget_unset_allows_large_results(monkeypatch): + # No token_budget set: even a large agent() result must not trip any cap, + # proving the default is genuinely unbounded (observable via ToolOk). + tool, _ = make_tool(monkeypatch, responder=lambda p: "x" * 5000) script = ( 'meta = {"name": "n", "description": "d"}\n' 'r = await agent("x", {"label": "L"})\n' 'return {"r": r}\n' ) - res = await tool(tool.params(script=script, token_budget=12345)) + res = await tool(tool.params(script=script)) assert not res.is_error - assert captured.get("token_budget") == 12345 @pytest.mark.asyncio -async def test_token_budget_defaults_to_none(monkeypatch): - captured: dict[str, object] = {} - import pythinker_code.tools.workflow as mod - - real_run_workflow = mod.run_workflow - - async def spy_run_workflow(*args, **kwargs): - captured.update(kwargs) - return await real_run_workflow(*args, **kwargs) - - monkeypatch.setattr(mod, "run_workflow", spy_run_workflow) - +async def test_log_calls_reach_the_progress_note(monkeypatch): tool, _ = make_tool(monkeypatch, responder=lambda p: "ok") + notes: list[ProgressNote] = [] + monkeypatch.setattr("pythinker_code.tools.workflow.wire_send", notes.append) script = ( 'meta = {"name": "n", "description": "d"}\n' + 'log("custom checkpoint reached")\n' 'r = await agent("x", {"label": "L"})\n' 'return {"r": r}\n' ) res = await tool(tool.params(script=script)) assert not res.is_error - assert captured.get("token_budget") is None + assert any("custom checkpoint reached" in note.body for note in notes) @pytest.mark.asyncio