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/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/src/pythinker_code/tools/workflow/__init__.py b/src/pythinker_code/tools/workflow/__init__.py new file mode 100644 index 00000000..76dac86f --- /dev/null +++ b/src/pythinker_code/tools/workflow/__init__.py @@ -0,0 +1,244 @@ +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`.", + ) + 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()/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)." + ), + ge=1, + ) + + +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_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.agent_id, event.label, event.phase) + emit() + + def on_agent_end(event: AgentEndEvent) -> None: + snapshot.end_agent(event.agent_id, error=event.error) + emit() + + hooks = RunWorkflowHooks( + on_log=on_log, + 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), + token_budget=params.token_budget, + 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 " + "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..2332acf3 --- /dev/null +++ b/src/pythinker_code/tools/workflow/description.md @@ -0,0 +1,22 @@ +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). `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. +- 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()` — 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/src/pythinker_code/tools/workflow/display.py b/src/pythinker_code/tools/workflow/display.py new file mode 100644 index 00000000..abd1e8e6 --- /dev/null +++ b/src/pythinker_code/tools/workflow/display.py @@ -0,0 +1,105 @@ +"""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, agent_id: int, label: str, phase: str | None) -> AgentSnapshot: + self.add_phase(phase) + agent = AgentSnapshot(agent_id, label, phase) + self.agents.append(agent) + return agent + + 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 + + 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, max_logs: int = 3) -> 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}") + 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 new file mode 100644 index 00000000..461ffa35 --- /dev/null +++ b/src/pythinker_code/tools/workflow/engine.py @@ -0,0 +1,536 @@ +"""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 +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 = { + "__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): + 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_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" + ) + + +# --------------------------------------------------------------------------- +# 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__ = ("agent_id", "label", "phase", "prompt") + + 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__ = ("agent_id", "label", "phase", "result", "error") + + def __init__( + 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 + 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 _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=_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"), + ) + + +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"] + 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 + # 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(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(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(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(agent_id, 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/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..657d0cb9 --- /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_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`` -> +``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", _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_display.py b/tests/tools/test_workflow_display.py new file mode 100644 index 00000000..fedb2cd0 --- /dev/null +++ b/tests/tools/test_workflow_display.py @@ -0,0 +1,62 @@ +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(1, "repo inventory", "Scan") + assert a.status == "running" + assert snap.running_count == 1 + snap.end_agent(1) + assert snap.done_count == 1 and snap.running_count == 0 + + snap.add_phase("Analyze") + snap.start_agent(2, "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(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 new file mode 100644 index 00000000..f9c240eb --- /dev/null +++ b/tests/tools/test_workflow_engine.py @@ -0,0 +1,252 @@ +# tests/tools/test_workflow_engine.py +import asyncio + +import pytest + +from pythinker_code.tools.workflow.engine import ( + AgentOptions, + RunWorkflowHooks, + WorkflowRuntimeError, + _estimate_tokens, + 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_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 + # 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_cancelled_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 diff --git a/tests/tools/test_workflow_parser.py b/tests/tools/test_workflow_parser.py new file mode 100644 index 00000000..eea5e4ca --- /dev/null +++ b/tests/tools/test_workflow_parser.py @@ -0,0 +1,48 @@ +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"), + ( + '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): + with pytest.raises(WorkflowScriptError) as exc: + parse_workflow_script(script) + assert fragment in str(exc.value) diff --git a/tests/tools/test_workflow_registration.py b/tests/tools/test_workflow_registration.py new file mode 100644 index 00000000..9720476a --- /dev/null +++ b/tests/tools/test_workflow_registration.py @@ -0,0 +1,30 @@ +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). 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 + ) + assert "pythinker_code.tools.workflow:Workflow" not in effective_coder_tools diff --git a/tests/tools/test_workflow_tool.py b/tests/tools/test_workflow_tool.py new file mode 100644 index 00000000..7b94e6f0 --- /dev/null +++ b/tests/tools/test_workflow_tool.py @@ -0,0 +1,201 @@ +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 +from pythinker_code.wire.types import ProgressNote + + +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. + monkeypatch.setattr( + "pythinker_code.tools.workflow.AgentTool", lambda rt: FakeAgentTool(responder) + ) + # Suppress wire emission (no Wire ContextVar in a unit test). + monkeypatch.setattr("pythinker_code.tools.workflow.wire_send", lambda *a, **k: None) + tool = Workflow(cast(Runtime, 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 isinstance(res.output, str) + 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_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() + + +@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)) + assert not res.is_error + + +@pytest.mark.asyncio +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 any("custom checkpoint reached" in note.body for note in notes) + + +@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 isinstance(res.output, str) + 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", ]