From c9cd0b8e1d26c88256a3f9c8496bb4693d02e6a5 Mon Sep 17 00:00:00 2001 From: Kiran Gopinathan Date: Tue, 10 Mar 2026 23:34:35 -0400 Subject: [PATCH 1/4] implemented multi-agent support in effectful using a basic choreographic programming primitive --- docs/source/multi_agent_example.py | 310 +++++++++ effectful/handlers/llm/multi.py | 657 ++++++++++++++++++ tests/test_handlers_llm_provider.py | 118 ++++ tests/test_multi_agent_epp.py | 989 ++++++++++++++++++++++++++++ 4 files changed, 2074 insertions(+) create mode 100644 docs/source/multi_agent_example.py create mode 100644 effectful/handlers/llm/multi.py create mode 100644 tests/test_multi_agent_epp.py diff --git a/docs/source/multi_agent_example.py b/docs/source/multi_agent_example.py new file mode 100644 index 000000000..7c0304654 --- /dev/null +++ b/docs/source/multi_agent_example.py @@ -0,0 +1,310 @@ +"""Multi-agent system using choreographic endpoint projection. + +Demonstrates: +- Choreographic programming: one function describes the entire workflow +- Automatic endpoint projection: each agent gets its own thread +- Crash tolerance: Ctrl-C and restart, agents resume where they left off +- Scatter: two coder agents share the implementation work via claim-based pull +- PersistentAgent for automatic checkpointing and context compaction + +The scenario: a team of agents collaboratively builds a small Python library. +An architect agent breaks the project into module specs, two coder agents +implement the modules in parallel (via scatter), and two reviewer agents +review modules in parallel and request fixes if needed. + +Usage:: + + # First run — agents start working + python docs/source/multi_agent_example.py + + # Ctrl-C mid-run, then restart — agents pick up where they left off + python docs/source/multi_agent_example.py + +Requirements: + pip install effectful[llm] + export OPENAI_API_KEY=... # or any LiteLLM-supported provider + +""" + +import json +import logging +from pathlib import Path +from typing import Literal, TypedDict + +from effectful.handlers.llm import Template, Tool +from effectful.handlers.llm.completions import LiteLLMProvider, RetryLLMHandler +from effectful.handlers.llm.multi import Choreography, ChoreographyError, scatter +from effectful.handlers.llm.persistence import PersistenceHandler, PersistentAgent +from effectful.ops.types import NotHandled + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(threadName)s] %(message)s", + datefmt="%H:%M:%S", +) +log = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +WORKSPACE = Path("./multi_agent_workspace") +STATE_DIR = WORKSPACE / ".state" +OUTPUT_DIR = WORKSPACE / "output" +MODEL = "gpt-4o-mini" + +# The project to build +PROJECT_SPEC = """\ +Build a small Python utility library called 'textkit' with these modules: +1. textkit/slugify.py — convert strings to URL-safe slugs +2. textkit/wrap.py — word-wrap text to a given width +3. textkit/redact.py — redact email addresses and phone numbers from text +Each module should have a clear public API, docstrings, and at least 3 +test cases written as a separate test_.py file. +""" + + +# --------------------------------------------------------------------------- +# Structured types — constrained decoding for LLM output +# --------------------------------------------------------------------------- + + +class ModuleSpec(TypedDict): + """Schema for architect planning output — constrained decoding ensures valid shape.""" + + module_path: str + description: str + public_api: str + test_path: str + + +class PlanResult(TypedDict): + """Wrapper for list output — LiteLLM requires a root object, not bare array.""" + + modules: list[ModuleSpec] + + +class ReviewResult(TypedDict): + """Schema for reviewer output — verdict constrained to PASS or NEEDS_FIXES.""" + + verdict: Literal["PASS", "NEEDS_FIXES"] + feedback: str + + +# --------------------------------------------------------------------------- +# Agents +# --------------------------------------------------------------------------- + + +class ArchitectAgent(PersistentAgent): + """You are a software architect. Given a project specification, you break + it into individual module implementation tasks. Each task should specify + the module filename, its public API, and what tests to write. + Be concrete and specific — the coder will follow your spec exactly. + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._output_dir = OUTPUT_DIR + + @Tool.define + def read_existing_files(self) -> str: + """List files already written to the output directory.""" + if not self._output_dir.exists(): + return "No files yet." + files = sorted(self._output_dir.rglob("*.py")) + if not files: + return "No Python files yet." + return "\n".join(str(f.relative_to(self._output_dir)) for f in files) + + @Template.define + def plan_modules(self, project_spec: str) -> PlanResult: + """Given this project specification, output a plan with a "modules" list. + Each module spec has: module_path, description, public_api, test_path. + + Use `read_existing_files` to check what's already been written + and skip those. + + Project spec: + {project_spec}""" + raise NotHandled + + +class CoderAgent(PersistentAgent): + """You are an expert Python developer. Given a module specification, + you write clean, well-documented Python code. You also write thorough + test files. Output ONLY the Python source code, no markdown fences. + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._output_dir = OUTPUT_DIR + + @Tool.define + def read_file(self, path: str) -> str: + """Read a file from the output directory.""" + full = self._output_dir / path + if full.exists(): + return full.read_text() + return f"File not found: {path}" + + @Tool.define + def write_file(self, path: str, content: str) -> str: + """Write a file to the output directory.""" + full = self._output_dir / path + full.parent.mkdir(parents=True, exist_ok=True) + full.write_text(content) + return f"Wrote {len(content)} chars to {path}" + + @Template.define + def implement_module(self, module_spec: str) -> str: + """Implement the following module specification. Use `write_file` + to write both the module and its test file. Use `read_file` to + check existing code if needed. + + Specification: + {module_spec}""" + raise NotHandled + + +class ReviewerAgent(PersistentAgent): + """You are a senior code reviewer. You review Python modules for + correctness, style, edge cases, and test coverage. Be specific + about issues and provide actionable feedback. + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._output_dir = OUTPUT_DIR + + @Tool.define + def read_file(self, path: str) -> str: + """Read a file from the output directory.""" + full = self._output_dir / path + if full.exists(): + return full.read_text() + return f"File not found: {path}" + + @Template.define + def review_module(self, module_path: str, test_path: str) -> ReviewResult: + """Review the module at {module_path} and its tests at {test_path}. + Use `read_file` to read them. Return verdict "PASS" or "NEEDS_FIXES" + and feedback. If NEEDS_FIXES, explain exactly what to change.""" + raise NotHandled + + +# --------------------------------------------------------------------------- +# Choreographic program — the entire multi-agent workflow in one function +# --------------------------------------------------------------------------- + + +def build_project( + project_spec: str, + architect: ArchitectAgent, + coder: CoderAgent, + reviewer: ReviewerAgent, +) -> list[ReviewResult]: + """Choreographic program describing the full build workflow. + + 1. Architect breaks the project into module specs. + 2. Coders implement modules in parallel (scatter distributes via claim-based pull). + 3. Reviewers review modules in parallel; coders fix in parallel until all pass. + """ + # Step 1: Architect plans modules + plan = architect.plan_modules(project_spec) + + # Step 2: Scatter implementation across coders + # Each module becomes a task in the queue; coders claim until none remain. + scatter( + plan["modules"], + coder, + lambda c, mod: c.implement_module(json.dumps(mod, indent=2)), + ) + + # Step 3: Review loop — keep fixing until reviewers accept all modules + while True: + reviews: list[ReviewResult] = scatter( + plan["modules"], + reviewer, + lambda r, mod: r.review_module(mod["module_path"], mod["test_path"]), + ) + + needs_fixes = [ + (mod, review) + for mod, review in zip(plan["modules"], reviews) + if review["verdict"] == "NEEDS_FIXES" + ] + + if not needs_fixes: + return reviews + + # Scatter fixes across coders, then re-review + scatter( + needs_fixes, + coder, + lambda c, pair: c.implement_module( + json.dumps( + {**pair[0], "fix_feedback": pair[1]["feedback"]}, + indent=2, + ) + ), + ) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + WORKSPACE.mkdir(parents=True, exist_ok=True) + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + + # Create agents + architect = ArchitectAgent(agent_id="architect") + coder1 = CoderAgent(agent_id="coder-1") + coder2 = CoderAgent(agent_id="coder-2") + reviewer1 = ReviewerAgent(agent_id="reviewer-1") + reviewer2 = ReviewerAgent(agent_id="reviewer-2") + + # Build the choreography — all boilerplate (threads, queues, signal + # handling, crash recovery) is handled automatically. + choreo = Choreography( + build_project, + agents=[architect, coder1, coder2, reviewer1, reviewer2], + state_dir=STATE_DIR, + handlers=[ + LiteLLMProvider(model=MODEL), + RetryLLMHandler(), + PersistenceHandler(STATE_DIR), + ], + ) + + log.info("Starting multi-agent build (Ctrl-C to pause, re-run to resume)") + + try: + reviews = choreo.run( + project_spec=PROJECT_SPEC, + architect=architect, + coder=[coder1, coder2], + reviewer=[reviewer1, reviewer2], + ) + except ChoreographyError as e: + log.error("Choreography failed: %s", e) + return + + # Summary + output_files = list(OUTPUT_DIR.rglob("*.py")) + passed = sum(1 for r in reviews if r["verdict"] == "PASS") + log.info( + "Done: %d modules reviewed (%d passed), %d output files", + len(reviews), + passed, + len(output_files), + ) + for f in output_files: + log.info(" %s", f.relative_to(WORKSPACE)) + + +if __name__ == "__main__": + main() diff --git a/effectful/handlers/llm/multi.py b/effectful/handlers/llm/multi.py new file mode 100644 index 000000000..541de5c1b --- /dev/null +++ b/effectful/handlers/llm/multi.py @@ -0,0 +1,657 @@ +"""Choreographic programming for multi-agent LLM systems. + +Write a single function describing how agents interact from a global +perspective, then run it with automatic endpoint projection (EPP). +Each agent gets its own thread, inter-agent communication is +handled automatically via a persistent :class:`TaskQueue`, and the +entire process is crash-tolerant and restartable. + +**How it works.** The choreographic program is a plain Python function +whose arguments are agent instances. All agent threads run this same +function. The :class:`EndpointProjection` handler intercepts +:attr:`~effectful.handlers.llm.template.Template.__apply__`: + +- When it is the current agent's template: claim a task in the + queue, execute via ``fwd``, and store the result. +- When it is another agent's template: poll the queue until + the result appears. + +Each statement in the choreography is assigned an incrementing step ID. +Completed steps are persisted to disk. On restart, the program re-runs +from the start; completed steps return their cached results instantly, +and execution resumes from the first incomplete step. + +Example — sequential choreography with a review loop:: + + from pathlib import Path + from typing import Literal, TypedDict + + from effectful.handlers.llm import Template + from effectful.handlers.llm.completions import LiteLLMProvider, RetryLLMHandler + from effectful.handlers.llm.multi import Choreography + from effectful.handlers.llm.persistence import PersistenceHandler, PersistentAgent + from effectful.ops.types import NotHandled + + class ModuleSpec(TypedDict): + module_path: str + description: str + + class PlanResult(TypedDict): + modules: list[ModuleSpec] + + class ReviewResult(TypedDict): + verdict: Literal["PASS", "NEEDS_FIXES"] + feedback: str + + class Architect(PersistentAgent): + \"\"\"You are a software architect.\"\"\" + + @Template.define + def plan_modules(self, project_spec: str) -> PlanResult: + \"\"\"Break this project into modules: {project_spec}\"\"\" + raise NotHandled + + class Coder(PersistentAgent): + \"\"\"You are a Python developer.\"\"\" + + @Template.define + def implement_module(self, spec: str) -> str: + \"\"\"Implement the module: {spec}\"\"\" + raise NotHandled + + class Reviewer(PersistentAgent): + \"\"\"You are a code reviewer.\"\"\" + + @Template.define + def review_code(self, code: str) -> ReviewResult: + \"\"\"Review this code: {code}\"\"\" + raise NotHandled + + def build_codebase( + project_spec: str, + architect: Architect, + coder: Coder, + reviewer: Reviewer, + ) -> str: + plan = architect.plan_modules(project_spec) + code = coder.implement_module(str(plan)) + while True: + result = reviewer.review_code(code) + if result["verdict"] == "PASS": + return code + code = coder.implement_module(result["feedback"]) + + architect = Architect(agent_id="architect") + coder = Coder(agent_id="coder") + reviewer = Reviewer(agent_id="reviewer") + + choreo = Choreography( + build_codebase, + agents=[architect, coder, reviewer], + state_dir=Path("./state"), + handlers=[ + LiteLLMProvider(model="gpt-4o-mini"), + RetryLLMHandler(), + PersistenceHandler(Path("./state")), + ], + ) + # Kill at any point, restart, and it resumes where it left off. + result = choreo.run( + project_spec="Build a URL slugify library", + architect=architect, + coder=coder, + reviewer=reviewer, + ) + +Example — parallel scatter across multiple coders:: + + from effectful.handlers.llm.multi import Choreography, scatter + + def build_parallel( + project_spec: str, + architect: Architect, + coder: Coder, + reviewer: Reviewer, + ) -> list[ReviewResult]: + plan = architect.plan_modules(project_spec) + # Each module becomes a task; coders claim from the queue + # until none remain — natural load balancing. + codes = scatter( + plan["modules"], coder, + lambda coder, mod: coder.implement_module(str(mod)), + ) + return [reviewer.review_code(code) for code in codes] + + coder1 = Coder(agent_id="coder-1") + coder2 = Coder(agent_id="coder-2") + coder3 = Coder(agent_id="coder-3") + + choreo = Choreography( + build_parallel, + agents=[architect, coder1, coder2, coder3, reviewer], + state_dir=Path("./state"), + handlers=[LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler()], + ) + # Pass coder as a list — scatter distributes across all three + reviews = choreo.run( + project_spec="Build textkit with slugify, wrap, and redact modules", + architect=architect, + coder=[coder1, coder2, coder3], + reviewer=reviewer, + ) + +""" + +import contextlib +import json +import threading +import time +import uuid +from collections.abc import Callable, Sequence +from enum import StrEnum +from pathlib import Path +from typing import Any + +from effectful.handlers.llm.template import Agent, Template, get_bound_agent +from effectful.ops.semantics import fwd, handler +from effectful.ops.syntax import ObjectInterpretation, implements +from effectful.ops.types import Interpretation, Operation + +# ── TaskQueue ────────────────────────────────────────────────────── + + +class TaskStatus(StrEnum): + PENDING = "pending" + CLAIMED = "claimed" + DONE = "done" + FAILED = "failed" + + +class TaskQueue: + """File-based task queue with claim-based ownership. + + Each task is a JSON file in *queue_dir*. Claiming a task + atomically renames it from ``.pending.json`` to + ``.claimed..json``, preventing double-claiming even + across process restarts. + + The queue is fully crash-tolerant: call + :meth:`release_stale_claims` on restart to reclaim work from a + prior crashed session. + """ + + def __init__(self, queue_dir: Path): + self.queue_dir = queue_dir + self.queue_dir.mkdir(parents=True, exist_ok=True) + self._lock = threading.Lock() + + def _task_path(self, task_id: str, status: str, owner: str = "") -> Path: + if owner: + return self.queue_dir / f"{task_id}.{status}.{owner}.json" + return self.queue_dir / f"{task_id}.{status}.json" + + def submit( + self, + task_type: str, + payload: dict, + task_id: str | None = None, + ) -> str: + """Add a new task. Returns the task ID. + + Idempotent when *task_id* is specified: if a task with that ID + already exists (in any state), the call is a no-op. + """ + if task_id is None: + task_id = str(uuid.uuid4())[:8] + with self._lock: + if list(self.queue_dir.glob(f"{task_id}.*")): + return task_id + task = { + "id": task_id, + "type": task_type, + "payload": payload, + "status": TaskStatus.PENDING, + "owner": "", + "result": None, + } + path = self._task_path(task_id, TaskStatus.PENDING) + path.write_text(json.dumps(task, indent=2, default=str)) + return task_id + + def claim(self, task_type: str, owner: str) -> dict | None: + """Atomically claim the next pending task of the given type. + + Returns the task dict if one was claimed, or ``None``. + """ + with self._lock: + for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): + task = json.loads(path.read_text()) + if task["type"] != task_type: + continue + task["status"] = TaskStatus.CLAIMED + task["owner"] = owner + claimed = self._task_path(task["id"], TaskStatus.CLAIMED, owner) + try: + path.rename(claimed) + except FileNotFoundError: + continue # another thread claimed it + claimed.write_text(json.dumps(task, indent=2, default=str)) + return task + return None + + def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: + """Claim any pending task whose ID starts with *prefix*.""" + with self._lock: + for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): + fname = path.name.split(".")[0] + if not fname.startswith(prefix): + continue + task = json.loads(path.read_text()) + task["status"] = TaskStatus.CLAIMED + task["owner"] = owner + claimed = self._task_path(task["id"], TaskStatus.CLAIMED, owner) + try: + path.rename(claimed) + except FileNotFoundError: + continue + claimed.write_text(json.dumps(task, indent=2, default=str)) + return task + return None + + def complete(self, task_id: str, owner: str, result: Any = None) -> None: + """Mark a claimed task as done with *result*.""" + claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) + if not claimed.exists(): + return + task = json.loads(claimed.read_text()) + task["status"] = TaskStatus.DONE + task["result"] = result + done = self._task_path(task_id, TaskStatus.DONE) + tmp = done.with_suffix(".tmp") + tmp.write_text(json.dumps(task, indent=2, default=str)) + tmp.replace(done) + try: + claimed.unlink() + except FileNotFoundError: + pass + + def fail(self, task_id: str, owner: str, error: str) -> None: + """Mark a claimed task as failed.""" + claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) + if not claimed.exists(): + return + task = json.loads(claimed.read_text()) + task["status"] = TaskStatus.FAILED + task["result"] = {"error": error} + failed = self._task_path(task_id, TaskStatus.FAILED) + claimed.rename(failed) + failed.write_text(json.dumps(task, indent=2, default=str)) + + def get_result(self, task_id: str) -> Any | None: + """Return the result of a completed task, or ``None``.""" + done = self._task_path(task_id, TaskStatus.DONE) + if done.exists(): + task = json.loads(done.read_text()) + return task.get("result") + return None + + def release_stale_claims(self, owner: str) -> int: + """Release tasks claimed by *owner* back to pending. + + Call on startup to reclaim work from a prior crashed session. + """ + count = 0 + with self._lock: + for path in self.queue_dir.glob(f"*.{TaskStatus.CLAIMED}.{owner}.json"): + task = json.loads(path.read_text()) + task["status"] = TaskStatus.PENDING + task["owner"] = "" + pending = self._task_path(task["id"], TaskStatus.PENDING) + path.rename(pending) + pending.write_text(json.dumps(task, indent=2, default=str)) + count += 1 + return count + + def pending_count(self, task_type: str | None = None) -> int: + """Count pending tasks, optionally filtered by type.""" + count = 0 + for path in self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json"): + if task_type is None: + count += 1 + else: + task = json.loads(path.read_text()) + if task["type"] == task_type: + count += 1 + return count + + def all_done(self) -> bool: + """``True`` if no pending or claimed tasks remain.""" + for status in (TaskStatus.PENDING, TaskStatus.CLAIMED): + if list(self.queue_dir.glob(f"*.{status}*")): + return False + return True + + +# ── scatter ──────────────────────────────────────────────────────── + + +@Operation.define +def scatter(items: list, agent: Agent, fn: Callable) -> list: + """Distribute *items* by calling ``fn(agent, item)`` for each item. + + **Default** (no EPP handler): sequential + ``[fn(agent, item) for item in items]``. + + **Under** :class:`EndpointProjection`: each item becomes a task in + the queue. When a list of agents is passed for the same role + (e.g. ``coder=[coder1, coder2]``), agents claim tasks until none + remain — providing natural load balancing with crash recovery. + On restart, completed items are returned from cache; only + remaining items are re-executed. + + .. warning:: + + ``fn`` should only call templates on the assigned agent. + Cross-agent template calls inside scatter are not supported. + """ + return [fn(agent, item) for item in items] + + +# ── Endpoint Projection ─────────────────────────────────────────── + + +class ChoreographyError(Exception): + """Raised when a choreography fails due to an agent error.""" + + +class CancelledError(Exception): + """Raised inside an agent thread when the choreography is cancelled.""" + + +class EndpointProjection(ObjectInterpretation): + """Handler that projects a choreographic program onto a single agent. + + Each template call in the choreography is assigned a step ID + (incrementing counter). Steps become tasks in the + :class:`TaskQueue`. + + - **Own agent's templates**: check if the step is already done + (cached); if not, claim the task, execute, and store the result. + - **Other agent's templates**: poll the queue until the result + appears. + - **Unbound templates**: execute directly on all threads. + + Also handles :func:`scatter` for data-parallel distribution via + claim-based pull. + """ + + def __init__( + self, + agent: Agent, + queue: TaskQueue, + agent_ids: frozenset[str], + poll_interval: float = 0.1, + cancel_event: threading.Event | None = None, + ) -> None: + self._agent = agent + self._agent_id = agent.__agent_id__ + self._queue = queue + self._agent_ids = agent_ids + self._poll = poll_interval + self._step = 0 + self._in_scatter = False + self._cancel = cancel_event + + def _next_step(self) -> str: + step_id = f"step-{self._step:04d}" + self._step += 1 + return step_id + + def _check_cancelled(self) -> None: + if self._cancel is not None and self._cancel.is_set(): + raise CancelledError("Choreography cancelled") + + def _wait_result(self, step_id: str) -> Any: + """Poll queue until task result is available.""" + while True: + self._check_cancelled() + r = self._queue.get_result(step_id) + if r is not None: + return r + time.sleep(self._poll) + + @implements(Template.__apply__) + def _call[**P, T]( + self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs + ) -> T: + bound = get_bound_agent(template) + + # Inside scatter: execute directly, no task management + if self._in_scatter: + if bound and bound.__agent_id__ == self._agent_id: + return fwd(template, *args, **kwargs) + raise RuntimeError( + f"Cross-agent call in scatter: {self._agent_id} -> " + f"{bound.__agent_id__ if bound else '?'}" + ) + + step_id = self._next_step() + + if bound is not None and bound.__agent_id__ == self._agent_id: + # My template — check done cache, or claim and execute + cached = self._queue.get_result(step_id) + if cached is not None: + return cached + + self._queue.submit( + task_type=template.__name__, + payload={"agent": self._agent_id}, + task_id=step_id, + ) + task = self._queue.claim(template.__name__, self._agent_id) + if task is None: + # Already claimed (e.g. restarted while another thread + # is executing) — poll for result + return self._wait_result(step_id) + + try: + result = fwd(template, *args, **kwargs) + self._queue.complete(step_id, self._agent_id, result) + return result + except Exception as e: + self._queue.fail(step_id, self._agent_id, str(e)) + raise + + elif bound is not None: + # Another agent's template — poll for result + return self._wait_result(step_id) + + else: + # Unbound template — execute directly + return fwd(template, *args, **kwargs) + + @implements(scatter) + def _scatter(self, items: list, agent: Agent, fn: Callable) -> list: + step_id = self._next_step() + + # agent may be a single Agent or a list of Agents (passed + # transparently from choreo.run kwargs). Normalize to a list. + agents = agent if isinstance(agent, list) else [agent] + scatter_ids = {a.__agent_id__ for a in agents} + + # Submit one task per item + for i in range(len(items)): + self._queue.submit( + task_type=f"scatter-{step_id}", + payload={"item_index": i}, + task_id=f"{step_id}:{i:04d}", + ) + + # If I'm a scatter agent, claim and execute until none left + if self._agent_id in scatter_ids: + while True: + task = self._queue.claim_by_prefix(f"{step_id}:", self._agent_id) + if task is None: + break + idx = task["payload"]["item_index"] + self._in_scatter = True + try: + result = fn(self._agent, items[idx]) + self._queue.complete(task["id"], self._agent_id, result) + except Exception as e: + self._queue.fail(task["id"], self._agent_id, str(e)) + raise + finally: + self._in_scatter = False + + # Gather all results (blocking until done) + return [self._wait_result(f"{step_id}:{i:04d}") for i in range(len(items))] + + +# ── Choreography runner ─────────────────────────────────────────── + + +class Choreography: + """Run a choreographic program with crash-tolerant endpoint projection. + + Each agent gets its own thread. Template calls are routed via + the :class:`TaskQueue`: the owning agent claims and executes, + others poll for results. On restart, completed steps are + returned from cache. + + Args: + program: The choreographic function. All agent threads run + this same function; EPP makes each thread behave + differently. + agents: The agents participating in the choreography. + state_dir: Directory for the persistent task queue. + handlers: Handler instances to install per-thread beneath + the EPP handler (e.g. LLM provider, retry handler, + persistence handler). + poll_interval: Seconds between polling for task results + (default 0.1). + + Example:: + + choreo = Choreography( + build_codebase, + agents=[architect, coder, reviewer], + state_dir=Path("./state"), + handlers=[ + LiteLLMProvider(model="gpt-4o-mini"), + RetryLLMHandler(), + PersistenceHandler(Path("./state")), + ], + ) + result = choreo.run( + project_spec="Build a library...", + architect=architect, + coder=coder, + reviewer=reviewer, + ) + """ + + def __init__( + self, + program: Callable[..., Any], + agents: Sequence[Agent], + state_dir: Path, + handlers: Sequence[Interpretation | ObjectInterpretation] | None = None, + poll_interval: float = 0.1, + ) -> None: + self.program = program + self.agents = list(agents) + self.state_dir = Path(state_dir) + self.handlers = list(handlers or []) + self.poll_interval = poll_interval + self._queue = TaskQueue(self.state_dir / "choreo_queue") + + @property + def queue(self) -> TaskQueue: + """The underlying task queue (for inspection or manual ops).""" + return self._queue + + def project( + self, + agent: Agent, + cancel_event: threading.Event | None = None, + ) -> EndpointProjection: + """Return the EPP handler for a specific agent. + + Useful for manual thread management:: + + proj = choreo.project(agent) + with handler(provider), handler(proj): + result = choreo.program(**kwargs) + """ + return EndpointProjection( + agent, + self._queue, + frozenset(a.__agent_id__ for a in self.agents), + self.poll_interval, + cancel_event=cancel_event, + ) + + def run(self, **kwargs: Any) -> Any: + """Run the choreography to completion. + + Keyword arguments are forwarded to the choreographic function. + Returns the result (identical across all agent threads). + + On restart after a crash, completed steps return cached + results; stale claims are released and re-executed. + + Raises: + ChoreographyError: If any agent thread fails. + """ + # Release stale claims from prior crashed run + for agent in self.agents: + self._queue.release_stale_claims(agent.__agent_id__) + + cancel = threading.Event() + results: dict[str, Any] = {} + errors: list[tuple[str, BaseException]] = [] + lock = threading.Lock() + + def agent_main(agent: Agent) -> None: + try: + proj = self.project(agent, cancel_event=cancel) + result = self._run_with_handlers(proj, **kwargs) + with lock: + results[agent.__agent_id__] = result + except CancelledError: + pass # another agent failed; this thread was cancelled + except BaseException as e: + cancel.set() # signal other threads to stop + with lock: + errors.append((agent.__agent_id__, e)) + + threads = [] + for agent in self.agents: + t = threading.Thread( + target=agent_main, + args=(agent,), + name=f"choreo-{agent.__agent_id__}", + daemon=True, + ) + t.start() + threads.append(t) + + for t in threads: + t.join() + + if errors: + agent_id, exc = errors[0] + raise ChoreographyError(f"Agent '{agent_id}' failed: {exc}") from exc + + # All agents compute the same result; return any. + return next(iter(results.values())) + + def _run_with_handlers(self, proj: EndpointProjection, **kwargs: Any) -> Any: + """Install handlers and EPP, then run the program.""" + with contextlib.ExitStack() as stack: + for h in self.handlers: + stack.enter_context(handler(h)) + # EPP outermost — intercepts before providers + stack.enter_context(handler(proj)) + return self.program(**kwargs) diff --git a/tests/test_handlers_llm_provider.py b/tests/test_handlers_llm_provider.py index f3d2af9ca..5899cb82b 100644 --- a/tests/test_handlers_llm_provider.py +++ b/tests/test_handlers_llm_provider.py @@ -2604,3 +2604,121 @@ def ask(self, q: str) -> str: history = json.loads(row[0]) first_msg = history[0] assert "CONTEXT SUMMARY" in first_msg["content"] + + +# --------------------------------------------------------------------------- +# Integration tests: Multi-agent choreography (based on multi_agent_example.py) +# --------------------------------------------------------------------------- + + +@requires_openai +def test_multi_agent_choreography_integration(tmp_path): + """Multi-agent choreography: architect plans, coder implements, reviewer reviews. + + Based on docs/source/multi_agent_example.py but with constrained scope + (one module, capped LLM calls) so the test completes quickly. + """ + from typing import Literal, TypedDict + + from effectful.handlers.llm.multi import Choreography, scatter + from effectful.handlers.llm.persistence import PersistenceHandler, PersistentAgent + + class ModuleSpec(TypedDict): + module_path: str + description: str + + class PlanResult(TypedDict): + modules: list[ModuleSpec] + + class ReviewResult(TypedDict): + verdict: Literal["PASS", "NEEDS_FIXES"] + feedback: str + + class Architect(PersistentAgent): + """You are a software architect. Given a project spec, output a plan + with a 'modules' list. Each module has module_path and description. + Output exactly ONE module.""" + + @Template.define + def plan_modules(self, project_spec: str) -> PlanResult: + """Plan modules for: {project_spec}. Return exactly one module.""" + raise NotHandled + + class Coder(PersistentAgent): + """You are a Python developer. Write clean code. Reply with ONLY + the Python source code, no markdown.""" + + @Template.define + def implement_module(self, module_spec: str) -> str: + """Implement: {module_spec}""" + raise NotHandled + + class Reviewer(PersistentAgent): + """You are a code reviewer. Review for correctness and style.""" + + @Template.define + def review_code(self, code: str) -> ReviewResult: + """Review this code. Return verdict PASS or NEEDS_FIXES: {code}""" + raise NotHandled + + MAX_REVIEW_ROUNDS = 3 + + def build_project( + project_spec: str, + architect: Architect, + coders: list[Coder], + reviewer: Reviewer, + ) -> list[ReviewResult]: + # Step 1: Architect plans modules + plan = architect.plan_modules(project_spec) + + # Step 2: Scatter implementation across coders + codes = scatter( + plan["modules"], + coders, + lambda coder, mod: coder.implement_module(str(mod)), + ) + + # Step 3: Review each module; if NEEDS_FIXES, re-implement with feedback + reviews: list[ReviewResult] = [] + for mod, code in zip(plan["modules"], codes): + for _round in range(MAX_REVIEW_ROUNDS): + review = reviewer.review_code(code) + if review["verdict"] == "PASS": + break + # Re-implement with reviewer feedback + fix_spec = f"{mod}\nFix feedback: {review['feedback']}" + code = coders[0].implement_module(fix_spec) + reviews.append(review) + + return reviews + + architect = Architect(agent_id="architect") + coder1 = Coder(agent_id="coder-1") + reviewer = Reviewer(agent_id="reviewer") + + state_dir = tmp_path / "state" + choreo = Choreography( + build_project, + agents=[architect, coder1, reviewer], + state_dir=state_dir, + handlers=[ + LiteLLMProvider(model="gpt-4o-mini", max_tokens=300), + LimitLLMCallsHandler(max_calls=15), + PersistenceHandler(state_dir), + ], + poll_interval=0.05, + ) + + reviews = choreo.run( + project_spec="Build a single-function module: slugify(text) -> str", + architect=architect, + coders=[coder1], + reviewer=reviewer, + ) + + assert isinstance(reviews, list) + assert len(reviews) >= 1 + for r in reviews: + assert r["verdict"] in ("PASS", "NEEDS_FIXES") + assert isinstance(r["feedback"], str) diff --git a/tests/test_multi_agent_epp.py b/tests/test_multi_agent_epp.py new file mode 100644 index 000000000..1d50d3926 --- /dev/null +++ b/tests/test_multi_agent_epp.py @@ -0,0 +1,989 @@ +"""Tests for effectful.handlers.llm.multi — choreographic EPP with TaskQueue.""" + +import itertools +import json +import shutil +import threading +import time +from pathlib import Path +from typing import Any + +import pytest + +from effectful.handlers.llm import Template +from effectful.handlers.llm.multi import ( + Choreography, + ChoreographyError, + EndpointProjection, + TaskQueue, + TaskStatus, + scatter, +) +from effectful.handlers.llm.persistence import PersistentAgent +from effectful.handlers.llm.template import get_bound_agent +from effectful.ops.semantics import handler +from effectful.ops.syntax import ObjectInterpretation, implements +from effectful.ops.types import NotHandled + +# ── Fixtures and helpers ────────────────────────────────────────── + +STATE_DIR = Path("/tmp/test_multi_epp") + +# Default timeout for all concurrent tests (seconds). +# Concurrency bugs often manifest as hangs — this catches them. +THREAD_TIMEOUT = 10 + + +@pytest.fixture(autouse=True) +def clean_state(): + if STATE_DIR.exists(): + shutil.rmtree(STATE_DIR) + STATE_DIR.mkdir(parents=True) + yield + if STATE_DIR.exists(): + shutil.rmtree(STATE_DIR) + + +class MockLLM(ObjectInterpretation): + """Mock LLM handler that returns canned responses.""" + + def __init__(self, responses: dict[str, Any]): + self._responses = responses + self.calls: list[str] = [] + + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + key = ( + f"{bound.__agent_id__}.{template.__name__}" if bound else template.__name__ + ) + self.calls.append(key) + return self._responses.get( + key, self._responses.get(template.__name__, f"mock-{template.__name__}") + ) + + +class DelayedMockLLM(ObjectInterpretation): + """Mock LLM that introduces per-agent delays to force scheduling orderings. + + ``delays`` maps agent_id to a sleep duration (seconds) applied before + each template call for that agent. This lets tests deterministically + force one thread to run before another. + """ + + def __init__(self, responses: dict[str, Any], delays: dict[str, float]): + self._responses = responses + self._delays = delays + self.calls: list[str] = [] + self._lock = threading.Lock() + + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + agent_id = bound.__agent_id__ if bound else None + if agent_id and agent_id in self._delays: + time.sleep(self._delays[agent_id]) + key = ( + f"{bound.__agent_id__}.{template.__name__}" if bound else template.__name__ + ) + with self._lock: + self.calls.append(key) + return self._responses.get( + key, self._responses.get(template.__name__, f"mock-{template.__name__}") + ) + + +class FailingMockLLM(ObjectInterpretation): + """Mock LLM that raises on specific agent.template keys.""" + + def __init__(self, responses: dict[str, Any], fail_on: set[str]): + self._responses = responses + self._fail_on = fail_on + + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + key = ( + f"{bound.__agent_id__}.{template.__name__}" if bound else template.__name__ + ) + if key in self._fail_on: + raise RuntimeError(f"Simulated failure on {key}") + return self._responses.get( + key, self._responses.get(template.__name__, f"mock-{template.__name__}") + ) + + +def _run_threads_with_timeout(targets, timeout=THREAD_TIMEOUT): + """Start threads and join with timeout. Raises if any thread hangs.""" + threads = [threading.Thread(target=t, daemon=True) for t in targets] + for t in threads: + t.start() + for t in threads: + t.join(timeout=timeout) + if t.is_alive(): + raise TimeoutError( + f"Thread {t.name} did not finish within {timeout}s — " + "possible deadlock or infinite poll" + ) + + +# ── Agent definitions ───────────────────────────────────────────── + + +class Architect(PersistentAgent): + """Plans modules.""" + + @Template.define + def plan(self, spec: str) -> str: + """Plan modules for: {spec}""" + raise NotHandled + + +class Coder(PersistentAgent): + """Writes code.""" + + @Template.define + def implement(self, spec: str) -> str: + """Implement: {spec}""" + raise NotHandled + + +class Reviewer(PersistentAgent): + """Reviews code.""" + + @Template.define + def review(self, code: str) -> str: + """Review: {code}""" + raise NotHandled + + +# ── TaskQueue tests ─────────────────────────────────────────────── + + +class TestTaskQueue: + def test_submit_and_claim(self): + tq = TaskQueue(STATE_DIR / "q") + tid = tq.submit("code", {"file": "main.py"}, task_id="t1") + assert tid == "t1" + + task = tq.claim("code", "worker1") + assert task is not None + assert task["id"] == "t1" + assert task["status"] == TaskStatus.CLAIMED + + # Can't claim again + assert tq.claim("code", "worker2") is None + + def test_idempotent_submit(self): + tq = TaskQueue(STATE_DIR / "q") + tq.submit("code", {}, task_id="t1") + tq.submit("code", {}, task_id="t1") # no-op + assert tq.pending_count() == 1 + + def test_complete_and_get_result(self): + tq = TaskQueue(STATE_DIR / "q") + tq.submit("code", {}, task_id="t1") + tq.claim("code", "w1") + tq.complete("t1", "w1", {"output": "hello"}) + assert tq.get_result("t1") == {"output": "hello"} + + def test_release_stale_claims(self): + tq = TaskQueue(STATE_DIR / "q") + tq.submit("code", {}, task_id="t1") + tq.claim("code", "crashed_worker") + assert tq.pending_count() == 0 + + released = tq.release_stale_claims("crashed_worker") + assert released == 1 + assert tq.pending_count() == 1 + + # Can re-claim + task = tq.claim("code", "new_worker") + assert task is not None + + def test_claim_by_prefix(self): + tq = TaskQueue(STATE_DIR / "q") + tq.submit("scatter", {}, task_id="step-0001:0000") + tq.submit("scatter", {}, task_id="step-0001:0001") + tq.submit("other", {}, task_id="step-0002") + + task = tq.claim_by_prefix("step-0001:", "w1") + assert task is not None + assert task["id"].startswith("step-0001:") + + task2 = tq.claim_by_prefix("step-0001:", "w1") + assert task2 is not None + assert task2["id"] != task["id"] + + assert tq.claim_by_prefix("step-0001:", "w1") is None + + +# ── EPP tests ───────────────────────────────────────────────────── + + +class TestEndpointProjection: + def _run_choreo( + self, + agents, + choreo_fn, + responses, + *, + mock_cls=None, + timeout=THREAD_TIMEOUT, + **kwargs, + ): + """Helper: run a choreography with mock LLM. + + *mock_cls* can be a callable ``(responses) -> ObjectInterpretation`` + to inject custom mock behaviour (e.g. delays). + """ + tq = TaskQueue(STATE_DIR / "q") + ids = frozenset(a.__agent_id__ for a in agents) + results: dict[str, Any] = {} + llm_calls: dict[str, list[str]] = {} + errors: list[tuple[str, Exception]] = [] + lock = threading.Lock() + + def run_agent(agent): + try: + mock = mock_cls(responses) if mock_cls else MockLLM(responses) + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo_fn(**kwargs) + with lock: + results[agent.__agent_id__] = r + llm_calls[agent.__agent_id__] = list( + mock.calls if hasattr(mock, "calls") else [] + ) + except Exception as e: + with lock: + errors.append((agent.__agent_id__, e)) + + _run_threads_with_timeout( + [lambda a=a: run_agent(a) for a in agents], + timeout=timeout, + ) + + return results, llm_calls, errors + + def test_basic_sequential(self): + """Two agents: planner plans, worker executes.""" + planner = Architect(agent_id="arch") + worker = Coder(agent_id="coder") + + def choreo(spec, arch, coder): + plan = arch.plan(spec) + return coder.implement(plan) + + results, calls, errors = self._run_choreo( + [planner, worker], + choreo, + {"plan": "the plan", "implement": "code"}, + spec="build it", + arch=planner, + coder=worker, + ) + + assert not errors, errors + assert results["arch"] == "code" + assert results["coder"] == "code" + # Planner executed plan, worker executed implement + assert "arch.plan" in calls["arch"] + assert "coder.implement" in calls["coder"] + + def test_all_agents_same_result(self): + """All agent threads produce the same result.""" + arch = Architect(agent_id="a") + coder = Coder(agent_id="c") + reviewer = Reviewer(agent_id="r") + + def choreo(arch, coder, reviewer): + plan = arch.plan("spec") + code = coder.implement(plan) + return reviewer.review(code) + + results, _, errors = self._run_choreo( + [arch, coder, reviewer], + choreo, + {"plan": "plan", "implement": "code", "review": "PASS"}, + arch=arch, + coder=coder, + reviewer=reviewer, + ) + + assert not errors + assert results["a"] == results["c"] == results["r"] == "PASS" + + def test_while_loop(self): + """Control flow: reviewer retries, then passes.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + reviewer = Reviewer(agent_id="rev") + + review_count = {"n": 0} + + class LoopMock(ObjectInterpretation): + def __init__(self): + self.calls: list[str] = [] + + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + self.calls.append(template.__name__) + if template.__name__ == "plan": + return "plan" + elif template.__name__ == "implement": + return f"code-v{len(self.calls)}" + elif template.__name__ == "review": + review_count["n"] += 1 + return "RETRY" if review_count["n"] <= 1 else "PASS" + return "?" + + def choreo(arch, coder, reviewer): + plan = arch.plan("spec") + code = coder.implement(plan) + while True: + verdict = reviewer.review(code) + if verdict == "PASS": + return code + code = coder.implement(verdict) + + tq = TaskQueue(STATE_DIR / "q") + ids = frozenset(["arch", "coder", "rev"]) + results: dict[str, Any] = {} + errors: list = [] + lock = threading.Lock() + + def run(agent): + try: + mock = LoopMock() + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(arch=arch, coder=coder, reviewer=reviewer) + with lock: + results[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [arch, coder, reviewer]]) + + assert not errors, errors + assert all(r == results["arch"] for r in results.values()) + # Should have 5 persisted steps: plan, implement, review, implement, review + done_files = list(tq.queue_dir.glob(f"*.{TaskStatus.DONE}.json")) + assert len(done_files) == 5 + + def test_crash_recovery(self): + """Pre-cache step 0, restart: step 0 from cache, step 1 fresh.""" + tq = TaskQueue(STATE_DIR / "q") + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + # Simulate prior run: step 0 done + tq.submit("plan", {"agent": "arch"}, task_id="step-0000") + tq.claim("plan", "arch") + tq.complete("step-0000", "arch", "cached-plan") + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + ids = frozenset(["arch", "coder"]) + results: dict[str, Any] = {} + llm_calls: dict[str, list[str]] = {} + lock = threading.Lock() + errors: list = [] + + def run(agent): + try: + mock = MockLLM({"plan": "SHOULD NOT RUN", "implement": "fresh-code"}) + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(arch=arch, coder=coder) + with lock: + results[agent.__agent_id__] = r + llm_calls[agent.__agent_id__] = list(mock.calls) + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [arch, coder]]) + + assert not errors, errors + assert results["arch"] == "fresh-code" + assert results["coder"] == "fresh-code" + # arch should NOT have called LLM for plan + assert "arch.plan" not in llm_calls.get("arch", []) + # coder should have called implement + assert "coder.implement" in llm_calls.get("coder", []) + + +# ── Ordering permutation tests ──────────────────────────────────── + + +class TestOrderingPermutations: + """Run choreographies under every possible thread-scheduling order. + + Uses controlled delays to deterministically force one agent to + execute before another. For N agents there are N! orderings; + all must produce the same result. + """ + + def _run_with_ordering(self, agents, choreo_fn, responses, ordering, **kwargs): + """Run *choreo_fn* with agents delayed so they execute in *ordering*.""" + # Give each agent a staggered delay: first in ordering gets 0, + # second gets a small delay, etc. + delays = {agent.__agent_id__: i * 0.03 for i, agent in enumerate(ordering)} + tq = TaskQueue(STATE_DIR / "q") + ids = frozenset(a.__agent_id__ for a in agents) + results: dict[str, Any] = {} + errors: list[tuple[str, Exception]] = [] + lock = threading.Lock() + + def run_agent(agent): + try: + mock = DelayedMockLLM(responses, delays) + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo_fn(**kwargs) + with lock: + results[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append((agent.__agent_id__, e)) + + _run_threads_with_timeout([lambda a=a: run_agent(a) for a in agents]) + return results, errors + + def test_two_agent_all_orderings(self): + """Two agents, two orderings — both must agree on the result.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + agents = [arch, coder] + responses = {"plan": "the-plan", "implement": "the-code"} + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + all_results = [] + for perm in itertools.permutations(agents): + # Each permutation needs a fresh queue directory + q_dir = STATE_DIR / "q" + if q_dir.exists(): + shutil.rmtree(q_dir) + results, errors = self._run_with_ordering( + agents, + choreo, + responses, + perm, + arch=arch, + coder=coder, + ) + assert not errors, f"Ordering {[a.__agent_id__ for a in perm]}: {errors}" + all_results.append(results) + + # All orderings must produce the same result for every agent + for r in all_results: + assert r["arch"] == r["coder"] == "the-code" + + def test_three_agent_all_orderings(self): + """Three agents, six orderings — all must agree.""" + arch = Architect(agent_id="a") + coder = Coder(agent_id="c") + reviewer = Reviewer(agent_id="r") + agents = [arch, coder, reviewer] + responses = {"plan": "plan", "implement": "code", "review": "PASS"} + + def choreo(arch, coder, reviewer): + plan = arch.plan("spec") + code = coder.implement(plan) + return reviewer.review(code) + + expected = "PASS" + for perm in itertools.permutations(agents): + q_dir = STATE_DIR / "q" + if q_dir.exists(): + shutil.rmtree(q_dir) + results, errors = self._run_with_ordering( + agents, + choreo, + responses, + perm, + arch=arch, + coder=coder, + reviewer=reviewer, + ) + assert not errors, f"Ordering {[a.__agent_id__ for a in perm]}: {errors}" + for aid, r in results.items(): + assert r == expected, ( + f"Agent {aid} got {r!r} != {expected!r} " + f"with ordering {[a.__agent_id__ for a in perm]}" + ) + + def test_scatter_all_orderings(self): + """Scatter with two coders + reviewer, all orderings agree.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + rev = Reviewer(agent_id="rev") + agents = [c1, c2, rev] + responses = {"implement": "code", "review": "ok"} + + def choreo(items, coders, reviewer): + codes = scatter(items, coders, lambda c, m: c.implement(m)) + return [reviewer.review(c) for c in codes] + + for perm in itertools.permutations(agents): + q_dir = STATE_DIR / "q" + if q_dir.exists(): + shutil.rmtree(q_dir) + results, errors = self._run_with_ordering( + agents, + choreo, + responses, + perm, + items=["A", "B"], + coders=[c1, c2], + reviewer=rev, + ) + assert not errors, f"Ordering {[a.__agent_id__ for a in perm]}: {errors}" + for aid, r in results.items(): + assert r == ["ok", "ok"], ( + f"Agent {aid} got {r!r} with ordering " + f"{[a.__agent_id__ for a in perm]}" + ) + + +# ── Race condition tests ────────────────────────────────────────── + + +class TestRaceConditions: + def test_concurrent_claims_no_double_execution(self): + """Multiple workers racing to claim the same task type — + exactly one wins, no double execution.""" + tq = TaskQueue(STATE_DIR / "q") + tq.submit("work", {"data": "x"}, task_id="t1") + + claimed_by: list[str] = [] + lock = threading.Lock() + barrier = threading.Barrier(5) + + def try_claim(worker_id): + barrier.wait() # all threads start claiming simultaneously + task = tq.claim("work", worker_id) + if task is not None: + with lock: + claimed_by.append(worker_id) + + _run_threads_with_timeout([lambda w=f"w{i}": try_claim(w) for i in range(5)]) + assert len(claimed_by) == 1, f"Double-claim: {claimed_by}" + + def test_concurrent_submit_idempotent(self): + """Multiple threads submitting the same task_id — only one task created.""" + tq = TaskQueue(STATE_DIR / "q") + barrier = threading.Barrier(5) + + def submit(i): + barrier.wait() + tq.submit("work", {"thread": i}, task_id="same-id") + + _run_threads_with_timeout([lambda i=i: submit(i) for i in range(5)]) + assert tq.pending_count() == 1 + + def test_step_counters_stay_in_sync(self): + """All agent threads must assign the same step ID to each + choreographic statement, even under different scheduling.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + step_ids_seen: dict[str, list[str]] = {"arch": [], "coder": []} + lock = threading.Lock() + + class StepTrackingMock(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + return "result" + + tq = TaskQueue(STATE_DIR / "q") + ids = frozenset(["arch", "coder"]) + + def run_agent(agent): + mock = StepTrackingMock() + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + arch.plan("spec") + coder.implement("plan") + # After execution, check the step counter + with lock: + step_ids_seen[agent.__agent_id__].append(epp._step) + + _run_threads_with_timeout([lambda a=a: run_agent(a) for a in [arch, coder]]) + # Both agents should have advanced through the same number of steps + assert step_ids_seen["arch"] == step_ids_seen["coder"] + + def test_many_agents_many_steps(self): + """Stress test: 5 coders scatter over 20 items.""" + coders = [Coder(agent_id=f"c{i}") for i in range(5)] + items = list(range(20)) + responses = {"implement": "done"} + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(str(m))) + + mock = MockLLM(responses) + c = Choreography( + choreo, + agents=coders, + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + result = c.run(items=items, coders=coders) + assert result == ["done"] * 20 + + +# ── Edge case tests ─────────────────────────────────────────────── + + +class TestEdgeCases: + def test_single_agent_choreography(self): + """A choreography with only one agent works without deadlock.""" + coder = Coder(agent_id="solo") + + def choreo(coder): + return coder.implement("just me") + + mock = MockLLM({"implement": "solo-code"}) + c = Choreography( + choreo, + agents=[coder], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + result = c.run(coder=coder) + assert result == "solo-code" + + def test_empty_scatter(self): + """Scatter over an empty list returns [] without hanging.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + mock = MockLLM({"implement": "code"}) + c = Choreography( + choreo, + agents=[c1, c2], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + result = c.run(items=[], coders=[c1, c2]) + assert result == [] + + def test_agent_error_propagates(self): + """An exception in one agent's template propagates as ChoreographyError.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + mock = FailingMockLLM( + responses={"implement": "code"}, + fail_on={"arch.plan"}, + ) + c = Choreography( + choreo, + agents=[arch, coder], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + with pytest.raises(ChoreographyError, match="arch"): + c.run(arch=arch, coder=coder) + + def test_scatter_single_worker(self): + """Scatter with one worker still completes all items.""" + coder = Coder(agent_id="c1") + + call_count = {"n": 0} + call_lock = threading.Lock() + + class CountingMock(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + with call_lock: + call_count["n"] += 1 + return f"result-{call_count['n']}" + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + c = Choreography( + choreo, + agents=[coder], + state_dir=STATE_DIR, + handlers=[CountingMock()], + poll_interval=0.02, + ) + result = c.run(items=["a", "b", "c"], coders=[coder]) + assert len(result) == 3 + # All results should be non-None + assert all(r is not None for r in result) + + def test_scatter_error_propagates(self): + """An error inside a scatter item propagates as ChoreographyError.""" + c1 = Coder(agent_id="c1") + + class ScatterFailMock(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + raise RuntimeError("scatter item failed") + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + c = Choreography( + choreo, + agents=[c1], + state_dir=STATE_DIR, + handlers=[ScatterFailMock()], + poll_interval=0.02, + ) + with pytest.raises(ChoreographyError): + c.run(items=["x"], coders=[c1]) + + def test_result_none_vs_not_done(self): + """get_result returns None for non-existent tasks, not for tasks + whose result *is* None — ensuring poll loops don't confuse the two.""" + tq = TaskQueue(STATE_DIR / "q") + # Non-existent task + assert tq.get_result("nonexistent") is None + + # Task with an actual result of a falsy value + tq.submit("work", {}, task_id="t1") + tq.claim("work", "w1") + tq.complete("t1", "w1", 0) # result is 0 (falsy but not None) + assert tq.get_result("t1") == 0 + + def test_repeated_runs_deterministic(self): + """Running the same choreography 5 times gives identical results.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + results = [] + for i in range(5): + q_dir = STATE_DIR / "q" + if q_dir.exists(): + shutil.rmtree(q_dir) + mock = MockLLM({"plan": "plan", "implement": "code"}) + c = Choreography( + choreo, + agents=[arch, coder], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + results.append(c.run(arch=arch, coder=coder)) + + assert all(r == results[0] for r in results) + + +# ── Scatter tests ───────────────────────────────────────────────── + + +class TestScatter: + def test_scatter_distributes_work(self): + """Scatter distributes items across agents via claim mechanism.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + rev = Reviewer(agent_id="rev") + + tq = TaskQueue(STATE_DIR / "q") + ids = frozenset(["c1", "c2", "rev"]) + items = ["A", "B", "C", "D"] + + def choreo(items, coders, reviewer): + codes = scatter(items, coders, lambda c, m: c.implement(m)) + return [reviewer.review(code) for code in codes] + + results: dict[str, Any] = {} + llm_calls: dict[str, list[str]] = {} + errors: list = [] + lock = threading.Lock() + + def run(agent): + try: + mock = MockLLM({"implement": "code", "review": "PASS"}) + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(items, [c1, c2], rev) + with lock: + results[agent.__agent_id__] = r + llm_calls[agent.__agent_id__] = list(mock.calls) + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [c1, c2, rev]]) + + assert not errors, errors + assert results["c1"] == results["c2"] == results["rev"] + # Total execute calls across coders should be 4 + c1_impl = [c for c in llm_calls.get("c1", []) if "implement" in c] + c2_impl = [c for c in llm_calls.get("c2", []) if "implement" in c] + assert len(c1_impl) + len(c2_impl) == 4 + # Reviewer did all 4 reviews + rev_reviews = [c for c in llm_calls.get("rev", []) if "review" in c] + assert len(rev_reviews) == 4 + + def test_scatter_crash_recovery(self): + """Scatter with some items cached: only remaining items executed.""" + tq = TaskQueue(STATE_DIR / "q") + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + + items = ["A", "B", "C", "D"] + + # Pre-cache items 0 and 1 + for i in range(2): + step = f"step-0000:{i:04d}" + done = tq._task_path(step, TaskStatus.DONE) + done.write_text( + json.dumps( + { + "id": step, + "type": "scatter-step-0000", + "status": "done", + "result": f"cached-{items[i]}", + "payload": {"item_index": i}, + } + ) + ) + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + ids = frozenset(["c1", "c2"]) + results: dict[str, Any] = {} + llm_calls: dict[str, list[str]] = {} + errors: list = [] + lock = threading.Lock() + + def run(agent): + try: + mock = MockLLM({"implement": "fresh"}) + tq.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(items, [c1, c2]) + with lock: + results[agent.__agent_id__] = r + llm_calls[agent.__agent_id__] = list(mock.calls) + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [c1, c2]]) + + assert not errors, errors + r = results["c1"] + assert r[0] == "cached-A" + assert r[1] == "cached-B" + assert r[2] == "fresh" + assert r[3] == "fresh" + + # Only 2 LLM calls total (items 2 and 3) + total = sum(len(calls) for calls in llm_calls.values()) + assert total == 2 + + +# ── Choreography runner tests ───────────────────────────────────── + + +class TestChoreography: + def test_run(self): + """High-level Choreography.run() orchestrates everything.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + mock = MockLLM({"plan": "plan", "implement": "code"}) + c = Choreography( + choreo, + agents=[arch, coder], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + result = c.run(arch=arch, coder=coder) + assert result == "code" + + def test_run_restart(self): + """Run, restart, verify cached results are used.""" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + mock1 = MockLLM({"plan": "plan-v1", "implement": "code-v1"}) + c1 = Choreography( + choreo, + agents=[arch, coder], + state_dir=STATE_DIR, + handlers=[mock1], + poll_interval=0.02, + ) + result1 = c1.run(arch=arch, coder=coder) + assert result1 == "code-v1" + + # "Restart" — new Choreography, same state_dir + # Even with different responses, should use cached + mock2 = MockLLM({"plan": "plan-v2", "implement": "code-v2"}) + c2 = Choreography( + choreo, + agents=[arch, coder], + state_dir=STATE_DIR, + handlers=[mock2], + poll_interval=0.02, + ) + result2 = c2.run(arch=arch, coder=coder) + # Both steps were cached from first run + assert result2 == "code-v1" + + def test_scatter_with_choreography(self): + """Choreography with scatter.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + mock = MockLLM({"implement": "code"}) + c = Choreography( + choreo, + agents=[c1, c2], + state_dir=STATE_DIR, + handlers=[mock], + poll_interval=0.02, + ) + result = c.run(items=["A", "B", "C"], coders=[c1, c2]) + assert result == ["code", "code", "code"] From 167e763dab5b086169be51c4049f8e5d74441404 Mon Sep 17 00:00:00 2001 From: Kiran Gopinathan Date: Tue, 10 Mar 2026 23:58:17 -0400 Subject: [PATCH 2/4] made taskqueue persistence optional --- docs/source/multi_agent_example.py | 9 +- effectful/handlers/llm/multi.py | 207 ++++++++++++++++++++++++---- tests/test_handlers_llm_provider.py | 4 +- tests/test_multi_agent_epp.py | 171 +++++++++++++++-------- 4 files changed, 298 insertions(+), 93 deletions(-) diff --git a/docs/source/multi_agent_example.py b/docs/source/multi_agent_example.py index 7c0304654..1bce13956 100644 --- a/docs/source/multi_agent_example.py +++ b/docs/source/multi_agent_example.py @@ -33,7 +33,12 @@ from effectful.handlers.llm import Template, Tool from effectful.handlers.llm.completions import LiteLLMProvider, RetryLLMHandler -from effectful.handlers.llm.multi import Choreography, ChoreographyError, scatter +from effectful.handlers.llm.multi import ( + Choreography, + ChoreographyError, + PersistentTaskQueue, + scatter, +) from effectful.handlers.llm.persistence import PersistenceHandler, PersistentAgent from effectful.ops.types import NotHandled @@ -272,7 +277,7 @@ def main() -> None: choreo = Choreography( build_project, agents=[architect, coder1, coder2, reviewer1, reviewer2], - state_dir=STATE_DIR, + queue=PersistentTaskQueue(STATE_DIR / "choreo_queue"), handlers=[ LiteLLMProvider(model=MODEL), RetryLLMHandler(), diff --git a/effectful/handlers/llm/multi.py b/effectful/handlers/llm/multi.py index 541de5c1b..096423464 100644 --- a/effectful/handlers/llm/multi.py +++ b/effectful/handlers/llm/multi.py @@ -88,7 +88,7 @@ def build_codebase( choreo = Choreography( build_codebase, agents=[architect, coder, reviewer], - state_dir=Path("./state"), + queue=PersistentTaskQueue(Path("./state/choreo_queue")), handlers=[ LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler(), @@ -105,7 +105,7 @@ def build_codebase( Example — parallel scatter across multiple coders:: - from effectful.handlers.llm.multi import Choreography, scatter + from effectful.handlers.llm.multi import Choreography, PersistentTaskQueue, scatter def build_parallel( project_spec: str, @@ -129,7 +129,7 @@ def build_parallel( choreo = Choreography( build_parallel, agents=[architect, coder1, coder2, coder3, reviewer], - state_dir=Path("./state"), + queue=PersistentTaskQueue(Path("./state/choreo_queue")), handlers=[LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler()], ) # Pass coder as a list — scatter distributes across all three @@ -142,6 +142,7 @@ def build_parallel( """ +import abc import contextlib import json import threading @@ -167,7 +168,168 @@ class TaskStatus(StrEnum): FAILED = "failed" -class TaskQueue: +class TaskQueue(abc.ABC): + """Abstract task queue with claim-based ownership. + + Subclasses implement persistent (file-based) or in-memory storage. + All methods are thread-safe. + """ + + @abc.abstractmethod + def submit( + self, + task_type: str, + payload: dict, + task_id: str | None = None, + ) -> str: + """Add a new task. Returns the task ID. + + Idempotent when *task_id* is specified: if a task with that ID + already exists (in any state), the call is a no-op. + """ + + @abc.abstractmethod + def claim(self, task_type: str, owner: str) -> dict | None: + """Atomically claim the next pending task of the given type. + + Returns the task dict if one was claimed, or ``None``. + """ + + @abc.abstractmethod + def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: + """Claim any pending task whose ID starts with *prefix*.""" + + @abc.abstractmethod + def complete(self, task_id: str, owner: str, result: Any = None) -> None: + """Mark a claimed task as done with *result*.""" + + @abc.abstractmethod + def fail(self, task_id: str, owner: str, error: str) -> None: + """Mark a claimed task as failed.""" + + @abc.abstractmethod + def get_result(self, task_id: str) -> Any | None: + """Return the result of a completed task, or ``None``.""" + + @abc.abstractmethod + def release_stale_claims(self, owner: str) -> int: + """Release tasks claimed by *owner* back to pending. + + Call on startup to reclaim work from a prior crashed session. + """ + + @abc.abstractmethod + def pending_count(self, task_type: str | None = None) -> int: + """Count pending tasks, optionally filtered by type.""" + + @abc.abstractmethod + def all_done(self) -> bool: + """``True`` if no pending or claimed tasks remain.""" + + +class InMemoryTaskQueue(TaskQueue): + """In-memory task queue for testing or ephemeral workflows. + + Not crash-tolerant — all state is lost when the process exits. + Thread-safe via a single lock. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._tasks: dict[str, dict] = {} # task_id -> task dict + + def submit( + self, + task_type: str, + payload: dict, + task_id: str | None = None, + ) -> str: + if task_id is None: + task_id = str(uuid.uuid4())[:8] + with self._lock: + if task_id in self._tasks: + return task_id + self._tasks[task_id] = { + "id": task_id, + "type": task_type, + "payload": payload, + "status": TaskStatus.PENDING, + "owner": "", + "result": None, + } + return task_id + + def claim(self, task_type: str, owner: str) -> dict | None: + with self._lock: + for task_id in sorted(self._tasks): + task = self._tasks[task_id] + if task["status"] == TaskStatus.PENDING and task["type"] == task_type: + task["status"] = TaskStatus.CLAIMED + task["owner"] = owner + return dict(task) + return None + + def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: + with self._lock: + for task_id in sorted(self._tasks): + task = self._tasks[task_id] + if task["status"] == TaskStatus.PENDING and task_id.startswith(prefix): + task["status"] = TaskStatus.CLAIMED + task["owner"] = owner + return dict(task) + return None + + def complete(self, task_id: str, owner: str, result: Any = None) -> None: + with self._lock: + task = self._tasks.get(task_id) + if task is None or task["status"] != TaskStatus.CLAIMED: + return + task["status"] = TaskStatus.DONE + task["result"] = result + + def fail(self, task_id: str, owner: str, error: str) -> None: + with self._lock: + task = self._tasks.get(task_id) + if task is None or task["status"] != TaskStatus.CLAIMED: + return + task["status"] = TaskStatus.FAILED + task["result"] = {"error": error} + + def get_result(self, task_id: str) -> Any | None: + with self._lock: + task = self._tasks.get(task_id) + if task is not None and task["status"] == TaskStatus.DONE: + return task["result"] + return None + + def release_stale_claims(self, owner: str) -> int: + count = 0 + with self._lock: + for task in self._tasks.values(): + if task["status"] == TaskStatus.CLAIMED and task["owner"] == owner: + task["status"] = TaskStatus.PENDING + task["owner"] = "" + count += 1 + return count + + def pending_count(self, task_type: str | None = None) -> int: + with self._lock: + return sum( + 1 + for t in self._tasks.values() + if t["status"] == TaskStatus.PENDING + and (task_type is None or t["type"] == task_type) + ) + + def all_done(self) -> bool: + with self._lock: + return not any( + t["status"] in (TaskStatus.PENDING, TaskStatus.CLAIMED) + for t in self._tasks.values() + ) + + +class PersistentTaskQueue(TaskQueue): """File-based task queue with claim-based ownership. Each task is a JSON file in *queue_dir*. Claiming a task @@ -196,11 +358,6 @@ def submit( payload: dict, task_id: str | None = None, ) -> str: - """Add a new task. Returns the task ID. - - Idempotent when *task_id* is specified: if a task with that ID - already exists (in any state), the call is a no-op. - """ if task_id is None: task_id = str(uuid.uuid4())[:8] with self._lock: @@ -219,10 +376,6 @@ def submit( return task_id def claim(self, task_type: str, owner: str) -> dict | None: - """Atomically claim the next pending task of the given type. - - Returns the task dict if one was claimed, or ``None``. - """ with self._lock: for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): task = json.loads(path.read_text()) @@ -240,7 +393,6 @@ def claim(self, task_type: str, owner: str) -> dict | None: return None def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: - """Claim any pending task whose ID starts with *prefix*.""" with self._lock: for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): fname = path.name.split(".")[0] @@ -259,7 +411,6 @@ def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: return None def complete(self, task_id: str, owner: str, result: Any = None) -> None: - """Mark a claimed task as done with *result*.""" claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) if not claimed.exists(): return @@ -276,7 +427,6 @@ def complete(self, task_id: str, owner: str, result: Any = None) -> None: pass def fail(self, task_id: str, owner: str, error: str) -> None: - """Mark a claimed task as failed.""" claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) if not claimed.exists(): return @@ -288,7 +438,6 @@ def fail(self, task_id: str, owner: str, error: str) -> None: failed.write_text(json.dumps(task, indent=2, default=str)) def get_result(self, task_id: str) -> Any | None: - """Return the result of a completed task, or ``None``.""" done = self._task_path(task_id, TaskStatus.DONE) if done.exists(): task = json.loads(done.read_text()) @@ -296,10 +445,6 @@ def get_result(self, task_id: str) -> Any | None: return None def release_stale_claims(self, owner: str) -> int: - """Release tasks claimed by *owner* back to pending. - - Call on startup to reclaim work from a prior crashed session. - """ count = 0 with self._lock: for path in self.queue_dir.glob(f"*.{TaskStatus.CLAIMED}.{owner}.json"): @@ -313,7 +458,6 @@ def release_stale_claims(self, owner: str) -> int: return count def pending_count(self, task_type: str | None = None) -> int: - """Count pending tasks, optionally filtered by type.""" count = 0 for path in self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json"): if task_type is None: @@ -325,7 +469,6 @@ def pending_count(self, task_type: str | None = None) -> int: return count def all_done(self) -> bool: - """``True`` if no pending or claimed tasks remain.""" for status in (TaskStatus.PENDING, TaskStatus.CLAIMED): if list(self.queue_dir.glob(f"*.{status}*")): return False @@ -479,7 +622,10 @@ def _scatter(self, items: list, agent: Agent, fn: Callable) -> list: agents = agent if isinstance(agent, list) else [agent] scatter_ids = {a.__agent_id__ for a in agents} - # Submit one task per item + # Submit one task per item. All agent threads execute this + # loop, but submit() is idempotent on task_id — the + # deterministic ID (step_id:index) ensures each task is + # created exactly once regardless of how many threads call it. for i in range(len(items)): self._queue.submit( task_type=f"scatter-{step_id}", @@ -512,7 +658,7 @@ def _scatter(self, items: list, agent: Agent, fn: Callable) -> list: class Choreography: - """Run a choreographic program with crash-tolerant endpoint projection. + """Run a choreographic program with endpoint projection. Each agent gets its own thread. Template calls are routed via the :class:`TaskQueue`: the owning agent claims and executes, @@ -524,7 +670,9 @@ class Choreography: this same function; EPP makes each thread behave differently. agents: The agents participating in the choreography. - state_dir: Directory for the persistent task queue. + queue: The task queue to use. Defaults to + :class:`InMemoryTaskQueue` if not provided. Pass a + :class:`PersistentTaskQueue` for crash tolerance. handlers: Handler instances to install per-thread beneath the EPP handler (e.g. LLM provider, retry handler, persistence handler). @@ -536,7 +684,7 @@ class Choreography: choreo = Choreography( build_codebase, agents=[architect, coder, reviewer], - state_dir=Path("./state"), + queue=PersistentTaskQueue(Path("./state/choreo_queue")), handlers=[ LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler(), @@ -555,16 +703,15 @@ def __init__( self, program: Callable[..., Any], agents: Sequence[Agent], - state_dir: Path, + queue: TaskQueue | None = None, handlers: Sequence[Interpretation | ObjectInterpretation] | None = None, poll_interval: float = 0.1, ) -> None: self.program = program self.agents = list(agents) - self.state_dir = Path(state_dir) self.handlers = list(handlers or []) self.poll_interval = poll_interval - self._queue = TaskQueue(self.state_dir / "choreo_queue") + self._queue = queue if queue is not None else InMemoryTaskQueue() @property def queue(self) -> TaskQueue: diff --git a/tests/test_handlers_llm_provider.py b/tests/test_handlers_llm_provider.py index 5899cb82b..20f43d2eb 100644 --- a/tests/test_handlers_llm_provider.py +++ b/tests/test_handlers_llm_provider.py @@ -2697,11 +2697,13 @@ def build_project( coder1 = Coder(agent_id="coder-1") reviewer = Reviewer(agent_id="reviewer") + from effectful.handlers.llm.multi import PersistentTaskQueue + state_dir = tmp_path / "state" choreo = Choreography( build_project, agents=[architect, coder1, reviewer], - state_dir=state_dir, + queue=PersistentTaskQueue(state_dir / "choreo_queue"), handlers=[ LiteLLMProvider(model="gpt-4o-mini", max_tokens=300), LimitLLMCallsHandler(max_calls=15), diff --git a/tests/test_multi_agent_epp.py b/tests/test_multi_agent_epp.py index 1d50d3926..2c0a5a143 100644 --- a/tests/test_multi_agent_epp.py +++ b/tests/test_multi_agent_epp.py @@ -1,7 +1,6 @@ """Tests for effectful.handlers.llm.multi — choreographic EPP with TaskQueue.""" import itertools -import json import shutil import threading import time @@ -15,7 +14,8 @@ Choreography, ChoreographyError, EndpointProjection, - TaskQueue, + InMemoryTaskQueue, + PersistentTaskQueue, TaskStatus, scatter, ) @@ -159,10 +159,28 @@ def review(self, code: str) -> str: # ── TaskQueue tests ─────────────────────────────────────────────── +# Counter to avoid directory collisions between parametrized persistent tests. +_ptq_counter = 0 + + +def _make_persistent_queue(): + global _ptq_counter + _ptq_counter += 1 + return PersistentTaskQueue(STATE_DIR / f"q-{_ptq_counter}") + + +@pytest.fixture(params=["persistent", "in_memory"]) +def make_queue(request): + """Parametrized fixture — runs each test against both queue backends.""" + if request.param == "persistent": + return _make_persistent_queue + else: + return InMemoryTaskQueue + class TestTaskQueue: - def test_submit_and_claim(self): - tq = TaskQueue(STATE_DIR / "q") + def test_submit_and_claim(self, make_queue): + tq = make_queue() tid = tq.submit("code", {"file": "main.py"}, task_id="t1") assert tid == "t1" @@ -174,21 +192,21 @@ def test_submit_and_claim(self): # Can't claim again assert tq.claim("code", "worker2") is None - def test_idempotent_submit(self): - tq = TaskQueue(STATE_DIR / "q") + def test_idempotent_submit(self, make_queue): + tq = make_queue() tq.submit("code", {}, task_id="t1") tq.submit("code", {}, task_id="t1") # no-op assert tq.pending_count() == 1 - def test_complete_and_get_result(self): - tq = TaskQueue(STATE_DIR / "q") + def test_complete_and_get_result(self, make_queue): + tq = make_queue() tq.submit("code", {}, task_id="t1") tq.claim("code", "w1") tq.complete("t1", "w1", {"output": "hello"}) assert tq.get_result("t1") == {"output": "hello"} - def test_release_stale_claims(self): - tq = TaskQueue(STATE_DIR / "q") + def test_release_stale_claims(self, make_queue): + tq = make_queue() tq.submit("code", {}, task_id="t1") tq.claim("code", "crashed_worker") assert tq.pending_count() == 0 @@ -201,8 +219,8 @@ def test_release_stale_claims(self): task = tq.claim("code", "new_worker") assert task is not None - def test_claim_by_prefix(self): - tq = TaskQueue(STATE_DIR / "q") + def test_claim_by_prefix(self, make_queue): + tq = make_queue() tq.submit("scatter", {}, task_id="step-0001:0000") tq.submit("scatter", {}, task_id="step-0001:0001") tq.submit("other", {}, task_id="step-0002") @@ -217,6 +235,58 @@ def test_claim_by_prefix(self): assert tq.claim_by_prefix("step-0001:", "w1") is None + def test_all_done(self, make_queue): + tq = make_queue() + assert tq.all_done() + + tq.submit("code", {}, task_id="t1") + assert not tq.all_done() + + tq.claim("code", "w1") + assert not tq.all_done() # claimed but not done + + tq.complete("t1", "w1", "result") + assert tq.all_done() + + def test_fail(self, make_queue): + tq = make_queue() + tq.submit("code", {}, task_id="t1") + tq.claim("code", "w1") + tq.fail("t1", "w1", "boom") + # Failed tasks are not pending/claimed, so all_done is True + assert tq.all_done() + # get_result returns None for failed tasks + assert tq.get_result("t1") is None + + def test_concurrent_claims(self, make_queue): + """Multiple threads claiming — no double claims.""" + tq = make_queue() + n_tasks = 20 + for i in range(n_tasks): + tq.submit("work", {"i": i}, task_id=f"t{i:03d}") + + claimed: list[dict] = [] + lock = threading.Lock() + + def claimer(owner): + while True: + task = tq.claim("work", owner) + if task is None: + break + with lock: + claimed.append(task) + + threads = [threading.Thread(target=claimer, args=(f"w{i}",)) for i in range(5)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=THREAD_TIMEOUT) + + # Each task claimed exactly once + ids = [t["id"] for t in claimed] + assert len(ids) == n_tasks + assert len(set(ids)) == n_tasks + # ── EPP tests ───────────────────────────────────────────────────── @@ -237,7 +307,7 @@ def _run_choreo( *mock_cls* can be a callable ``(responses) -> ObjectInterpretation`` to inject custom mock behaviour (e.g. delays). """ - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() ids = frozenset(a.__agent_id__ for a in agents) results: dict[str, Any] = {} llm_calls: dict[str, list[str]] = {} @@ -348,7 +418,7 @@ def choreo(arch, coder, reviewer): return code code = coder.implement(verdict) - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() ids = frozenset(["arch", "coder", "rev"]) results: dict[str, Any] = {} errors: list = [] @@ -370,13 +440,10 @@ def run(agent): assert not errors, errors assert all(r == results["arch"] for r in results.values()) - # Should have 5 persisted steps: plan, implement, review, implement, review - done_files = list(tq.queue_dir.glob(f"*.{TaskStatus.DONE}.json")) - assert len(done_files) == 5 def test_crash_recovery(self): """Pre-cache step 0, restart: step 0 from cache, step 1 fresh.""" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() arch = Architect(agent_id="arch") coder = Coder(agent_id="coder") @@ -436,7 +503,7 @@ def _run_with_ordering(self, agents, choreo_fn, responses, ordering, **kwargs): # Give each agent a staggered delay: first in ordering gets 0, # second gets a small delay, etc. delays = {agent.__agent_id__: i * 0.03 for i, agent in enumerate(ordering)} - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() ids = frozenset(a.__agent_id__ for a in agents) results: dict[str, Any] = {} errors: list[tuple[str, Exception]] = [] @@ -471,10 +538,6 @@ def choreo(arch, coder): all_results = [] for perm in itertools.permutations(agents): - # Each permutation needs a fresh queue directory - q_dir = STATE_DIR / "q" - if q_dir.exists(): - shutil.rmtree(q_dir) results, errors = self._run_with_ordering( agents, choreo, @@ -505,9 +568,6 @@ def choreo(arch, coder, reviewer): expected = "PASS" for perm in itertools.permutations(agents): - q_dir = STATE_DIR / "q" - if q_dir.exists(): - shutil.rmtree(q_dir) results, errors = self._run_with_ordering( agents, choreo, @@ -537,9 +597,6 @@ def choreo(items, coders, reviewer): return [reviewer.review(c) for c in codes] for perm in itertools.permutations(agents): - q_dir = STATE_DIR / "q" - if q_dir.exists(): - shutil.rmtree(q_dir) results, errors = self._run_with_ordering( agents, choreo, @@ -564,7 +621,7 @@ class TestRaceConditions: def test_concurrent_claims_no_double_execution(self): """Multiple workers racing to claim the same task type — exactly one wins, no double execution.""" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() tq.submit("work", {"data": "x"}, task_id="t1") claimed_by: list[str] = [] @@ -583,7 +640,7 @@ def try_claim(worker_id): def test_concurrent_submit_idempotent(self): """Multiple threads submitting the same task_id — only one task created.""" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() barrier = threading.Barrier(5) def submit(i): @@ -607,7 +664,7 @@ class StepTrackingMock(ObjectInterpretation): def _call(self, template, *args, **kwargs): return "result" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() ids = frozenset(["arch", "coder"]) def run_agent(agent): @@ -638,7 +695,7 @@ def choreo(items, coders): c = Choreography( choreo, agents=coders, - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -661,7 +718,7 @@ def choreo(coder): c = Choreography( choreo, agents=[coder], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -680,7 +737,7 @@ def choreo(items, coders): c = Choreography( choreo, agents=[c1, c2], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -703,7 +760,7 @@ def choreo(arch, coder): c = Choreography( choreo, agents=[arch, coder], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -730,7 +787,7 @@ def choreo(items, coders): c = Choreography( choreo, agents=[coder], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[CountingMock()], poll_interval=0.02, ) @@ -754,7 +811,7 @@ def choreo(items, coders): c = Choreography( choreo, agents=[c1], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[ScatterFailMock()], poll_interval=0.02, ) @@ -764,7 +821,7 @@ def choreo(items, coders): def test_result_none_vs_not_done(self): """get_result returns None for non-existent tasks, not for tasks whose result *is* None — ensuring poll loops don't confuse the two.""" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() # Non-existent task assert tq.get_result("nonexistent") is None @@ -792,7 +849,7 @@ def choreo(arch, coder): c = Choreography( choreo, agents=[arch, coder], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -811,7 +868,7 @@ def test_scatter_distributes_work(self): c2 = Coder(agent_id="c2") rev = Reviewer(agent_id="rev") - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() ids = frozenset(["c1", "c2", "rev"]) items = ["A", "B", "C", "D"] @@ -852,27 +909,18 @@ def run(agent): def test_scatter_crash_recovery(self): """Scatter with some items cached: only remaining items executed.""" - tq = TaskQueue(STATE_DIR / "q") + tq = InMemoryTaskQueue() c1 = Coder(agent_id="c1") c2 = Coder(agent_id="c2") items = ["A", "B", "C", "D"] - # Pre-cache items 0 and 1 + # Pre-cache items 0 and 1 via submit/claim/complete for i in range(2): step = f"step-0000:{i:04d}" - done = tq._task_path(step, TaskStatus.DONE) - done.write_text( - json.dumps( - { - "id": step, - "type": "scatter-step-0000", - "status": "done", - "result": f"cached-{items[i]}", - "payload": {"item_index": i}, - } - ) - ) + tq.submit("scatter-step-0000", {"item_index": i}, task_id=step) + tq.claim("scatter-step-0000", "prior-worker") + tq.complete(step, "prior-worker", f"cached-{items[i]}") def choreo(items, coders): return scatter(items, coders, lambda c, m: c.implement(m)) @@ -928,7 +976,7 @@ def choreo(arch, coder): c = Choreography( choreo, agents=[arch, coder], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) @@ -944,24 +992,27 @@ def choreo(arch, coder): plan = arch.plan("spec") return coder.implement(plan) + # Shared persistent queue survives across Choreography instances + shared_queue = PersistentTaskQueue(STATE_DIR / "restart_q") + mock1 = MockLLM({"plan": "plan-v1", "implement": "code-v1"}) c1 = Choreography( choreo, agents=[arch, coder], - state_dir=STATE_DIR, + queue=shared_queue, handlers=[mock1], poll_interval=0.02, ) result1 = c1.run(arch=arch, coder=coder) assert result1 == "code-v1" - # "Restart" — new Choreography, same state_dir + # "Restart" — new Choreography, same persistent queue # Even with different responses, should use cached mock2 = MockLLM({"plan": "plan-v2", "implement": "code-v2"}) c2 = Choreography( choreo, agents=[arch, coder], - state_dir=STATE_DIR, + queue=shared_queue, handlers=[mock2], poll_interval=0.02, ) @@ -981,7 +1032,7 @@ def choreo(items, coders): c = Choreography( choreo, agents=[c1, c2], - state_dir=STATE_DIR, + queue=InMemoryTaskQueue(), handlers=[mock], poll_interval=0.02, ) From 7c98cb793d0b790a19487f59560f23ea2c4007d6 Mon Sep 17 00:00:00 2001 From: Kiran Gopinathan Date: Wed, 11 Mar 2026 15:20:49 -0400 Subject: [PATCH 3/4] updated to use sqlite db --- docs/source/multi_agent_example.py | 4 +- effectful/handlers/llm/multi.py | 308 ++++++++++++++++++---------- tests/test_handlers_llm_provider.py | 85 +++++++- tests/test_multi_agent_epp.py | 216 ++++++++++++++++++- 4 files changed, 498 insertions(+), 115 deletions(-) diff --git a/docs/source/multi_agent_example.py b/docs/source/multi_agent_example.py index 1bce13956..74533d22b 100644 --- a/docs/source/multi_agent_example.py +++ b/docs/source/multi_agent_example.py @@ -277,11 +277,11 @@ def main() -> None: choreo = Choreography( build_project, agents=[architect, coder1, coder2, reviewer1, reviewer2], - queue=PersistentTaskQueue(STATE_DIR / "choreo_queue"), + queue=PersistentTaskQueue(STATE_DIR / "task_queue.db"), handlers=[ LiteLLMProvider(model=MODEL), RetryLLMHandler(), - PersistenceHandler(STATE_DIR), + PersistenceHandler(STATE_DIR / "checkpoints.db"), ], ) diff --git a/effectful/handlers/llm/multi.py b/effectful/handlers/llm/multi.py index 096423464..74fb0e6e3 100644 --- a/effectful/handlers/llm/multi.py +++ b/effectful/handlers/llm/multi.py @@ -88,11 +88,11 @@ def build_codebase( choreo = Choreography( build_codebase, agents=[architect, coder, reviewer], - queue=PersistentTaskQueue(Path("./state/choreo_queue")), + queue=PersistentTaskQueue(Path("./state/task_queue.db")), handlers=[ LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler(), - PersistenceHandler(Path("./state")), + PersistenceHandler(Path("./state/checkpoints.db")), ], ) # Kill at any point, restart, and it resumes where it left off. @@ -129,7 +129,7 @@ def build_parallel( choreo = Choreography( build_parallel, agents=[architect, coder1, coder2, coder3, reviewer], - queue=PersistentTaskQueue(Path("./state/choreo_queue")), + queue=PersistentTaskQueue(Path("./state/task_queue.db")), handlers=[LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler()], ) # Pass coder as a list — scatter distributes across all three @@ -145,6 +145,7 @@ def build_parallel( import abc import contextlib import json +import sqlite3 import threading import time import uuid @@ -329,28 +330,69 @@ def all_done(self) -> bool: ) +def _init_queue_db(conn: sqlite3.Connection) -> None: + """Create the tasks table and configure WAL mode for crash tolerance.""" + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS tasks ( + id TEXT PRIMARY KEY, + type TEXT NOT NULL, + payload TEXT NOT NULL DEFAULT '{}', + status TEXT NOT NULL DEFAULT 'pending', + owner TEXT NOT NULL DEFAULT '', + result TEXT + ) + """ + ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)") + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_tasks_status_type ON tasks(status, type)" + ) + conn.commit() + + class PersistentTaskQueue(TaskQueue): - """File-based task queue with claim-based ownership. + """SQLite-backed task queue with claim-based ownership. + + All task state is stored in a single SQLite database using WAL + journal mode for crash tolerance. If the process is killed + mid-transaction, SQLite's journal-based recovery ensures the + database remains consistent. - Each task is a JSON file in *queue_dir*. Claiming a task - atomically renames it from ``.pending.json`` to - ``.claimed..json``, preventing double-claiming even + Claiming a task atomically updates its status from ``pending`` to + ``claimed`` inside a transaction, preventing double-claiming even across process restarts. The queue is fully crash-tolerant: call :meth:`release_stale_claims` on restart to reclaim work from a prior crashed session. + + Args: + db_path: Path to the SQLite database file. """ - def __init__(self, queue_dir: Path): - self.queue_dir = queue_dir - self.queue_dir.mkdir(parents=True, exist_ok=True) + def __init__(self, db_path: Path): + self._db_path = Path(db_path) self._lock = threading.Lock() + self._db_initialized = False + self._init_lock = threading.Lock() - def _task_path(self, task_id: str, status: str, owner: str = "") -> Path: - if owner: - return self.queue_dir / f"{task_id}.{status}.{owner}.json" - return self.queue_dir / f"{task_id}.{status}.json" + @property + def db_path(self) -> Path: + """Path to the SQLite database file.""" + return self._db_path + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(str(self._db_path), timeout=10) + conn.execute("PRAGMA busy_timeout=5000") + if not self._db_initialized: + with self._init_lock: + if not self._db_initialized: + _init_queue_db(conn) + self._db_initialized = True + return conn def submit( self, @@ -360,119 +402,171 @@ def submit( ) -> str: if task_id is None: task_id = str(uuid.uuid4())[:8] - with self._lock: - if list(self.queue_dir.glob(f"{task_id}.*")): - return task_id - task = { - "id": task_id, - "type": task_type, - "payload": payload, - "status": TaskStatus.PENDING, - "owner": "", - "result": None, - } - path = self._task_path(task_id, TaskStatus.PENDING) - path.write_text(json.dumps(task, indent=2, default=str)) - return task_id + payload_json = json.dumps(payload, default=str) + conn = self._connect() + try: + conn.execute( + """ + INSERT OR IGNORE INTO tasks (id, type, payload, status, owner, result) + VALUES (?, ?, ?, ?, '', NULL) + """, + (task_id, task_type, payload_json, TaskStatus.PENDING), + ) + conn.commit() + finally: + conn.close() + return task_id def claim(self, task_type: str, owner: str) -> dict | None: with self._lock: - for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): - task = json.loads(path.read_text()) - if task["type"] != task_type: - continue - task["status"] = TaskStatus.CLAIMED - task["owner"] = owner - claimed = self._task_path(task["id"], TaskStatus.CLAIMED, owner) - try: - path.rename(claimed) - except FileNotFoundError: - continue # another thread claimed it - claimed.write_text(json.dumps(task, indent=2, default=str)) - return task - return None + conn = self._connect() + try: + row = conn.execute( + """ + SELECT id, type, payload, status, owner, result + FROM tasks + WHERE status = ? AND type = ? + ORDER BY id LIMIT 1 + """, + (TaskStatus.PENDING, task_type), + ).fetchone() + if row is None: + return None + task_id = row[0] + conn.execute( + "UPDATE tasks SET status = ?, owner = ? WHERE id = ?", + (TaskStatus.CLAIMED, owner, task_id), + ) + conn.commit() + return { + "id": task_id, + "type": row[1], + "payload": json.loads(row[2]), + "status": TaskStatus.CLAIMED, + "owner": owner, + "result": json.loads(row[5]) if row[5] is not None else None, + } + finally: + conn.close() def claim_by_prefix(self, prefix: str, owner: str) -> dict | None: with self._lock: - for path in sorted(self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json")): - fname = path.name.split(".")[0] - if not fname.startswith(prefix): - continue - task = json.loads(path.read_text()) - task["status"] = TaskStatus.CLAIMED - task["owner"] = owner - claimed = self._task_path(task["id"], TaskStatus.CLAIMED, owner) - try: - path.rename(claimed) - except FileNotFoundError: - continue - claimed.write_text(json.dumps(task, indent=2, default=str)) - return task - return None + conn = self._connect() + try: + row = conn.execute( + """ + SELECT id, type, payload, status, owner, result + FROM tasks + WHERE status = ? AND id LIKE ? + ORDER BY id LIMIT 1 + """, + (TaskStatus.PENDING, prefix + "%"), + ).fetchone() + if row is None: + return None + task_id = row[0] + conn.execute( + "UPDATE tasks SET status = ?, owner = ? WHERE id = ?", + (TaskStatus.CLAIMED, owner, task_id), + ) + conn.commit() + return { + "id": task_id, + "type": row[1], + "payload": json.loads(row[2]), + "status": TaskStatus.CLAIMED, + "owner": owner, + "result": json.loads(row[5]) if row[5] is not None else None, + } + finally: + conn.close() def complete(self, task_id: str, owner: str, result: Any = None) -> None: - claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) - if not claimed.exists(): - return - task = json.loads(claimed.read_text()) - task["status"] = TaskStatus.DONE - task["result"] = result - done = self._task_path(task_id, TaskStatus.DONE) - tmp = done.with_suffix(".tmp") - tmp.write_text(json.dumps(task, indent=2, default=str)) - tmp.replace(done) + result_json = json.dumps(result, default=str) + conn = self._connect() try: - claimed.unlink() - except FileNotFoundError: - pass + conn.execute( + """ + UPDATE tasks SET status = ?, result = ? + WHERE id = ? AND status = ? + """, + (TaskStatus.DONE, result_json, task_id, TaskStatus.CLAIMED), + ) + conn.commit() + finally: + conn.close() def fail(self, task_id: str, owner: str, error: str) -> None: - claimed = self._task_path(task_id, TaskStatus.CLAIMED, owner) - if not claimed.exists(): - return - task = json.loads(claimed.read_text()) - task["status"] = TaskStatus.FAILED - task["result"] = {"error": error} - failed = self._task_path(task_id, TaskStatus.FAILED) - claimed.rename(failed) - failed.write_text(json.dumps(task, indent=2, default=str)) + error_json = json.dumps({"error": error}, default=str) + conn = self._connect() + try: + conn.execute( + """ + UPDATE tasks SET status = ?, result = ? + WHERE id = ? AND status = ? + """, + (TaskStatus.FAILED, error_json, task_id, TaskStatus.CLAIMED), + ) + conn.commit() + finally: + conn.close() def get_result(self, task_id: str) -> Any | None: - done = self._task_path(task_id, TaskStatus.DONE) - if done.exists(): - task = json.loads(done.read_text()) - return task.get("result") - return None + conn = self._connect() + try: + row = conn.execute( + "SELECT result FROM tasks WHERE id = ? AND status = ?", + (task_id, TaskStatus.DONE), + ).fetchone() + if row is None: + return None + return json.loads(row[0]) if row[0] is not None else None + finally: + conn.close() def release_stale_claims(self, owner: str) -> int: - count = 0 with self._lock: - for path in self.queue_dir.glob(f"*.{TaskStatus.CLAIMED}.{owner}.json"): - task = json.loads(path.read_text()) - task["status"] = TaskStatus.PENDING - task["owner"] = "" - pending = self._task_path(task["id"], TaskStatus.PENDING) - path.rename(pending) - pending.write_text(json.dumps(task, indent=2, default=str)) - count += 1 - return count + conn = self._connect() + try: + cursor = conn.execute( + """ + UPDATE tasks SET status = ?, owner = '' + WHERE status = ? AND owner = ? + """, + (TaskStatus.PENDING, TaskStatus.CLAIMED, owner), + ) + conn.commit() + return cursor.rowcount + finally: + conn.close() def pending_count(self, task_type: str | None = None) -> int: - count = 0 - for path in self.queue_dir.glob(f"*.{TaskStatus.PENDING}.json"): + conn = self._connect() + try: if task_type is None: - count += 1 + row = conn.execute( + "SELECT COUNT(*) FROM tasks WHERE status = ?", + (TaskStatus.PENDING,), + ).fetchone() else: - task = json.loads(path.read_text()) - if task["type"] == task_type: - count += 1 - return count + row = conn.execute( + "SELECT COUNT(*) FROM tasks WHERE status = ? AND type = ?", + (TaskStatus.PENDING, task_type), + ).fetchone() + return row[0] if row else 0 + finally: + conn.close() def all_done(self) -> bool: - for status in (TaskStatus.PENDING, TaskStatus.CLAIMED): - if list(self.queue_dir.glob(f"*.{status}*")): - return False - return True + conn = self._connect() + try: + row = conn.execute( + "SELECT COUNT(*) FROM tasks WHERE status IN (?, ?)", + (TaskStatus.PENDING, TaskStatus.CLAIMED), + ).fetchone() + return row[0] == 0 if row else True + finally: + conn.close() # ── scatter ──────────────────────────────────────────────────────── @@ -684,11 +778,11 @@ class Choreography: choreo = Choreography( build_codebase, agents=[architect, coder, reviewer], - queue=PersistentTaskQueue(Path("./state/choreo_queue")), + queue=PersistentTaskQueue(Path("./state/task_queue.db")), handlers=[ LiteLLMProvider(model="gpt-4o-mini"), RetryLLMHandler(), - PersistenceHandler(Path("./state")), + PersistenceHandler(Path("./state/checkpoints.db")), ], ) result = choreo.run( diff --git a/tests/test_handlers_llm_provider.py b/tests/test_handlers_llm_provider.py index 20f43d2eb..528f2853e 100644 --- a/tests/test_handlers_llm_provider.py +++ b/tests/test_handlers_llm_provider.py @@ -2700,14 +2700,15 @@ def build_project( from effectful.handlers.llm.multi import PersistentTaskQueue state_dir = tmp_path / "state" + state_dir.mkdir() choreo = Choreography( build_project, agents=[architect, coder1, reviewer], - queue=PersistentTaskQueue(state_dir / "choreo_queue"), + queue=PersistentTaskQueue(state_dir / "task_queue.db"), handlers=[ - LiteLLMProvider(model="gpt-4o-mini", max_tokens=300), - LimitLLMCallsHandler(max_calls=15), - PersistenceHandler(state_dir), + LiteLLMProvider(model="gpt-4o-mini", max_tokens=200), + LimitLLMCallsHandler(max_calls=8), + PersistenceHandler(state_dir / "checkpoints.db"), ], poll_interval=0.05, ) @@ -2724,3 +2725,79 @@ def build_project( for r in reviews: assert r["verdict"] in ("PASS", "NEEDS_FIXES") assert isinstance(r["feedback"], str) + + +# --------------------------------------------------------------------------- +# Integration tests: SQLite task queue crash tolerance +# --------------------------------------------------------------------------- + + +@requires_openai +def test_sqlite_task_queue_crash_recovery_integration(tmp_path): + """Choreography with SQLite queue: completed steps survive restart.""" + from effectful.handlers.llm.multi import Choreography, PersistentTaskQueue + from effectful.handlers.llm.persistence import PersistenceHandler, PersistentAgent + + class Bot(PersistentAgent): + """You are a concise assistant. Reply in at most 10 words.""" + + @Template.define + def step1(self, q: str) -> str: + """Answer briefly: {q}""" + raise NotHandled + + @Template.define + def step2(self, context: str) -> str: + """Given context, reply in one word: {context}""" + raise NotHandled + + bot = Bot(agent_id="crash-q-bot") + + db_path = tmp_path / "task_queue.db" + + # Run 1: complete full choreography + def two_step(bot): + r1 = bot.step1("What is 2+2?") + return bot.step2(r1) + + choreo1 = Choreography( + two_step, + agents=[bot], + queue=PersistentTaskQueue(db_path), + handlers=[ + LiteLLMProvider(model="gpt-4o-mini", max_tokens=30), + RetryLLMHandler(), + LimitLLMCallsHandler(max_calls=3), + PersistenceHandler(tmp_path / "checkpoints.db"), + ], + poll_interval=0.05, + ) + result1 = choreo1.run(bot=bot) + assert isinstance(result1, str) + + # Verify SQLite DB exists and has task data + conn = sqlite3.connect(str(db_path)) + integrity = conn.execute("PRAGMA integrity_check").fetchone()[0] + assert integrity == "ok" + done_count = conn.execute( + "SELECT COUNT(*) FROM tasks WHERE status = 'done'" + ).fetchone()[0] + conn.close() + assert done_count == 2 # step-0000 and step-0001 + + # Run 2: new Choreography on same DB — steps return cached results + bot2 = Bot(agent_id="crash-q-bot") + choreo2 = Choreography( + two_step, + agents=[bot2], + queue=PersistentTaskQueue(db_path), + handlers=[ + LiteLLMProvider(model="gpt-4o-mini", max_tokens=30), + RetryLLMHandler(), + LimitLLMCallsHandler(max_calls=0), # no LLM calls allowed — must use cache + PersistenceHandler(tmp_path / "checkpoints.db"), + ], + poll_interval=0.05, + ) + result2 = choreo2.run(bot=bot2) + assert result2 == result1 # exact same cached result diff --git a/tests/test_multi_agent_epp.py b/tests/test_multi_agent_epp.py index 2c0a5a143..b9895f659 100644 --- a/tests/test_multi_agent_epp.py +++ b/tests/test_multi_agent_epp.py @@ -166,7 +166,7 @@ def review(self, code: str) -> str: def _make_persistent_queue(): global _ptq_counter _ptq_counter += 1 - return PersistentTaskQueue(STATE_DIR / f"q-{_ptq_counter}") + return PersistentTaskQueue(STATE_DIR / f"q-{_ptq_counter}.db") @pytest.fixture(params=["persistent", "in_memory"]) @@ -993,7 +993,7 @@ def choreo(arch, coder): return coder.implement(plan) # Shared persistent queue survives across Choreography instances - shared_queue = PersistentTaskQueue(STATE_DIR / "restart_q") + shared_queue = PersistentTaskQueue(STATE_DIR / "restart_q.db") mock1 = MockLLM({"plan": "plan-v1", "implement": "code-v1"}) c1 = Choreography( @@ -1038,3 +1038,215 @@ def choreo(items, coders): ) result = c.run(items=["A", "B", "C"], coders=[c1, c2]) assert result == ["code", "code", "code"] + + +# ── SQLite crash tolerance tests ───────────────────────────────── + + +class TestSQLiteTaskQueueCrashTolerance: + """Tests that verify SQLite-specific crash tolerance properties.""" + + def test_wal_mode_enabled(self): + """WAL journal mode is enabled for crash tolerance.""" + import sqlite3 + + tq = PersistentTaskQueue(STATE_DIR / "wal_test.db") + tq.submit("work", {}, task_id="t1") + + conn = sqlite3.connect(str(tq.db_path)) + mode = conn.execute("PRAGMA journal_mode").fetchone()[0] + conn.close() + assert mode == "wal" + + def test_db_integrity_after_operations(self): + """Database passes integrity check after various operations.""" + import sqlite3 + + tq = PersistentTaskQueue(STATE_DIR / "integrity_test.db") + + # Submit, claim, complete, fail + tq.submit("work", {"data": "a"}, task_id="t1") + tq.submit("work", {"data": "b"}, task_id="t2") + tq.claim("work", "w1") + tq.complete("t1", "w1", {"result": "done"}) + tq.claim("work", "w1") + tq.fail("t2", "w1", "boom") + + conn = sqlite3.connect(str(tq.db_path)) + result = conn.execute("PRAGMA integrity_check").fetchone()[0] + conn.close() + assert result == "ok" + + def test_persistence_across_instances(self): + """Data survives creating a new PersistentTaskQueue on same db.""" + db_path = STATE_DIR / "persist_test.db" + + # Instance 1: submit and complete + tq1 = PersistentTaskQueue(db_path) + tq1.submit("work", {"key": "value"}, task_id="t1") + tq1.claim("work", "w1") + tq1.complete("t1", "w1", "result-1") + + # Instance 2: verify result survives + tq2 = PersistentTaskQueue(db_path) + assert tq2.get_result("t1") == "result-1" + assert tq2.all_done() + + def test_stale_claims_survive_restart(self): + """Claimed-but-not-completed tasks are recoverable after restart.""" + db_path = STATE_DIR / "stale_test.db" + + # Instance 1: submit and claim (simulating crash before completion) + tq1 = PersistentTaskQueue(db_path) + tq1.submit("work", {"data": "x"}, task_id="t1") + tq1.submit("work", {"data": "y"}, task_id="t2") + tq1.claim("work", "crashed_worker") + tq1.claim("work", "crashed_worker") + + # Instance 2: restart, release stale claims, re-claim + tq2 = PersistentTaskQueue(db_path) + assert tq2.pending_count() == 0 # both are claimed + released = tq2.release_stale_claims("crashed_worker") + assert released == 2 + assert tq2.pending_count() == 2 + + # Can re-claim and complete + task = tq2.claim("work", "new_worker") + assert task is not None + tq2.complete(task["id"], "new_worker", "recovered") + assert tq2.get_result(task["id"]) == "recovered" + + def test_partial_choreography_restart(self): + """Choreography can resume from a partially completed state.""" + db_path = STATE_DIR / "partial_choreo.db" + arch = Architect(agent_id="arch") + coder = Coder(agent_id="coder") + + # "Run 1": complete step 0 (plan), simulate crash before step 1 + tq1 = PersistentTaskQueue(db_path) + tq1.submit("plan", {"agent": "arch"}, task_id="step-0000") + tq1.claim("plan", "arch") + tq1.complete("step-0000", "arch", "the-plan") + # step 1 never submitted (crash) + + # "Run 2": restart with same db — step 0 cached, step 1 fresh + tq2 = PersistentTaskQueue(db_path) + + def choreo(arch, coder): + plan = arch.plan("spec") + return coder.implement(plan) + + ids = frozenset(["arch", "coder"]) + results: dict[str, Any] = {} + errors: list = [] + lock = threading.Lock() + + def run(agent): + try: + mock = MockLLM({"plan": "SHOULD-NOT-RUN", "implement": "fresh-code"}) + tq2.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq2, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(arch=arch, coder=coder) + with lock: + results[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [arch, coder]]) + + assert not errors, errors + # Step 0 was cached as "the-plan", step 1 ran fresh + assert results["arch"] == "fresh-code" + assert results["coder"] == "fresh-code" + + def test_concurrent_claims_across_connections(self): + """Multiple threads with separate connections cannot double-claim.""" + db_path = STATE_DIR / "concurrent_conn.db" + tq = PersistentTaskQueue(db_path) + + n_tasks = 20 + for i in range(n_tasks): + tq.submit("work", {"i": i}, task_id=f"t{i:03d}") + + claimed: list[dict] = [] + lock = threading.Lock() + + def claimer(owner): + while True: + task = tq.claim("work", owner) + if task is None: + break + with lock: + claimed.append(task) + + _run_threads_with_timeout([lambda w=f"w{i}": claimer(w) for i in range(5)]) + + ids = [t["id"] for t in claimed] + assert len(ids) == n_tasks + assert len(set(ids)) == n_tasks + + def test_scatter_crash_recovery_sqlite(self): + """Scatter with SQLite queue: cached items survive restart.""" + db_path = STATE_DIR / "scatter_crash.db" + + # "Run 1": complete scatter items 0 and 1 + tq1 = PersistentTaskQueue(db_path) + for i in range(2): + step = f"step-0000:{i:04d}" + tq1.submit("scatter-step-0000", {"item_index": i}, task_id=step) + tq1.claim("scatter-step-0000", "prior-worker") + tq1.complete(step, "prior-worker", f"cached-{i}") + + # "Run 2": scatter should pick up cached items 0,1 and run 2,3 fresh + tq2 = PersistentTaskQueue(db_path) + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + items = ["A", "B", "C", "D"] + + def choreo(items, coders): + return scatter(items, coders, lambda c, m: c.implement(m)) + + ids = frozenset(["c1", "c2"]) + results: dict[str, Any] = {} + errors: list = [] + lock = threading.Lock() + + def run(agent): + try: + mock = MockLLM({"implement": "fresh"}) + tq2.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq2, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(items, [c1, c2]) + with lock: + results[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in [c1, c2]]) + + assert not errors, errors + r = results["c1"] + assert r[0] == "cached-0" + assert r[1] == "cached-1" + assert r[2] == "fresh" + assert r[3] == "fresh" + + def test_idempotent_submit_across_restarts(self): + """Re-submitting existing task_id after restart is a no-op.""" + db_path = STATE_DIR / "idempotent_restart.db" + + tq1 = PersistentTaskQueue(db_path) + tq1.submit("work", {"original": True}, task_id="t1") + + tq2 = PersistentTaskQueue(db_path) + tq2.submit("work", {"duplicate": True}, task_id="t1") # should be ignored + assert tq2.pending_count() == 1 + + # Verify original payload preserved + task = tq2.claim("work", "w1") + assert task is not None + assert task["payload"] == {"original": True} From 468203d75ce0b2d1f2f12a2abe376cd1b00b0b99 Mon Sep 17 00:00:00 2001 From: Kiran Gopinathan Date: Wed, 11 Mar 2026 15:37:49 -0400 Subject: [PATCH 4/4] primitive for more parallelism --- effectful/handlers/llm/multi.py | 95 ++++++ tests/test_multi_agent_epp.py | 539 ++++++++++++++++++++++++++++++++ 2 files changed, 634 insertions(+) diff --git a/effectful/handlers/llm/multi.py b/effectful/handlers/llm/multi.py index 74fb0e6e3..019c78fb5 100644 --- a/effectful/handlers/llm/multi.py +++ b/effectful/handlers/llm/multi.py @@ -594,6 +594,46 @@ def scatter(items: list, agent: Agent, fn: Callable) -> list: return [fn(agent, item) for item in items] +@Operation.define +def fan_out(groups: list[tuple[list, Agent, Callable]]) -> list[list]: + """Run multiple scatter-like operations concurrently. + + Each element of *groups* is a ``(items, agent, fn)`` triple — the + same arguments you would pass to :func:`scatter`. Returns a list + of result lists, one per group, in the same order as *groups*. + + **Default** (no EPP handler): sequential execution of each group:: + + [ + [fn(agent, item) for item in items] + for items, agent, fn in groups + ] + + **Under** :class:`EndpointProjection`: all groups' items are + submitted as tasks under a single step ID. Agents from *every* + group claim and execute work concurrently, so a spec-writer, + tester, and prover can all be working at the same time rather + than waiting for the previous scatter to finish. + + Example:: + + spec_results, test_results, proof_results = fan_out([ + (spec_tasks, spec_writer, + lambda w, b: w.write_spec(json.dumps(b, indent=2))), + (test_tasks, tester, + lambda t, b: t.write_tests_and_validate(json.dumps(b, indent=2))), + (proof_tasks, prover, + lambda p, b: p.prove_theorem(json.dumps(b, indent=2))), + ]) + + .. warning:: + + ``fn`` should only call templates on the assigned agent. + Cross-agent template calls inside fan_out are not supported. + """ + return [[fn(agent, item) for item in items] for items, agent, fn in groups] + + # ── Endpoint Projection ─────────────────────────────────────────── @@ -747,6 +787,61 @@ def _scatter(self, items: list, agent: Agent, fn: Callable) -> list: # Gather all results (blocking until done) return [self._wait_result(f"{step_id}:{i:04d}") for i in range(len(items))] + @implements(fan_out) + def _fan_out(self, groups: list[tuple[list, Agent, Callable]]) -> list[list]: + step_id = self._next_step() + + # For each group, normalize agents and build a mapping from + # agent_id → list of (group_index, items, fn) so each agent + # knows which groups it participates in. + group_agents: list[set[str]] = [] + group_fns: list[Callable] = [] + group_items: list[list] = [] + + for g, (items, agent, fn) in enumerate(groups): + agents = agent if isinstance(agent, list) else [agent] + group_agents.append({a.__agent_id__ for a in agents}) + group_fns.append(fn) + group_items.append(items) + + # Submit all tasks across all groups. Deterministic IDs: + # {step_id}:g{group}:{item_index} + for g in range(len(groups)): + for i in range(len(group_items[g])): + self._queue.submit( + task_type=f"fan-{step_id}:g{g}", + payload={"group": g, "item_index": i}, + task_id=f"{step_id}:g{g}:{i:04d}", + ) + + # Claim and execute tasks from my groups + my_groups = [g for g in range(len(groups)) if self._agent_id in group_agents[g]] + for g in my_groups: + prefix = f"{step_id}:g{g}:" + while True: + task = self._queue.claim_by_prefix(prefix, self._agent_id) + if task is None: + break + idx = task["payload"]["item_index"] + self._in_scatter = True + try: + result = group_fns[g](self._agent, group_items[g][idx]) + self._queue.complete(task["id"], self._agent_id, result) + except Exception as e: + self._queue.fail(task["id"], self._agent_id, str(e)) + raise + finally: + self._in_scatter = False + + # Gather results per group (blocking) + return [ + [ + self._wait_result(f"{step_id}:g{g}:{i:04d}") + for i in range(len(group_items[g])) + ] + for g in range(len(groups)) + ] + # ── Choreography runner ─────────────────────────────────────────── diff --git a/tests/test_multi_agent_epp.py b/tests/test_multi_agent_epp.py index b9895f659..459271e9b 100644 --- a/tests/test_multi_agent_epp.py +++ b/tests/test_multi_agent_epp.py @@ -17,6 +17,7 @@ InMemoryTaskQueue, PersistentTaskQueue, TaskStatus, + fan_out, scatter, ) from effectful.handlers.llm.persistence import PersistentAgent @@ -157,6 +158,24 @@ def review(self, code: str) -> str: raise NotHandled +class TesterAgent(PersistentAgent): + """Writes tests.""" + + @Template.define + def write_tests(self, spec: str) -> str: + """Write tests for: {spec}""" + raise NotHandled + + +class Prover(PersistentAgent): + """Proves theorems.""" + + @Template.define + def prove(self, spec: str) -> str: + """Prove: {spec}""" + raise NotHandled + + # ── TaskQueue tests ─────────────────────────────────────────────── # Counter to avoid directory collisions between parametrized persistent tests. @@ -1250,3 +1269,523 @@ def test_idempotent_submit_across_restarts(self): task = tq2.claim("work", "w1") assert task is not None assert task["payload"] == {"original": True} + + +# ── fan_out tests ───────────────────────────────────────────────── + + +class TestFanOutDefault: + """fan_out without EPP handler — sequential fallback.""" + + def test_default_sequential(self): + """Default fan_out runs groups sequentially.""" + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + + mock = MockLLM( + { + "c1.implement": "code-result", + "t1.write_tests": "test-result", + } + ) + with handler(mock): + results = fan_out( + [ + (["a", "b"], coder, lambda c, m: c.implement(m)), + (["x"], tester, lambda t, m: t.write_tests(m)), + ] + ) + assert len(results) == 2 + assert results[0] == ["code-result", "code-result"] + assert results[1] == ["test-result"] + + def test_default_empty_groups(self): + """fan_out with empty item lists returns empty lists.""" + coder = Coder(agent_id="c1") + results = fan_out( + [ + ([], coder, lambda c, m: c.implement(m)), + ] + ) + assert results == [[]] + + def test_default_no_groups(self): + """fan_out with no groups returns empty list.""" + results = fan_out([]) + assert results == [] + + +class TestFanOutEPP: + """fan_out under EndpointProjection — concurrent execution.""" + + def test_concurrent_different_agent_types(self): + """Three different agent types work concurrently via fan_out.""" + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + prover = Prover(agent_id="p1") + + mock = MockLLM( + { + "c1.implement": "impl-result", + "t1.write_tests": "test-result", + "p1.prove": "proof-result", + } + ) + queue = InMemoryTaskQueue() + agents = [coder, tester, prover] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def choreo(c, t, p): + return fan_out( + [ + (["a", "b"], c, lambda c, m: c.implement(m)), + (["x", "y", "z"], t, lambda t, m: t.write_tests(m)), + (["th1"], p, lambda p, m: p.prove(m)), + ] + ) + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(coder, tester, prover) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + # All agents compute same result + for aid in ["c1", "t1", "p1"]: + r = results_map[aid] + assert r[0] == ["impl-result", "impl-result"] + assert r[1] == ["test-result", "test-result", "test-result"] + assert r[2] == ["proof-result"] + + def test_fan_out_with_multiple_workers_per_group(self): + """fan_out with multiple coders in one group, single tester in another.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + t1 = TesterAgent(agent_id="t1") + + execution_log: list[str] = [] + log_lock = threading.Lock() + + class TrackingMockLLM(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + key = f"{bound.__agent_id__}.{template.__name__}" if bound else "" + with log_lock: + execution_log.append(key) + return f"result-{key}" + + mock = TrackingMockLLM() + queue = InMemoryTaskQueue() + agents = [c1, c2, t1] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def choreo(coders, tester): + return fan_out( + [ + (["m1", "m2", "m3", "m4"], coders, lambda c, m: c.implement(m)), + (["t1", "t2"], tester, lambda t, m: t.write_tests(m)), + ] + ) + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo([c1, c2], t1) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + # Coders split the 4 items; tester handles 2 items + r = results_map["c1"] + assert len(r[0]) == 4 # 4 code results + assert len(r[1]) == 2 # 2 test results + # Both coders should have executed some implement calls + coder_calls = [c for c in execution_log if c.endswith(".implement")] + tester_calls = [c for c in execution_log if c.endswith(".write_tests")] + assert len(coder_calls) == 4 + assert len(tester_calls) == 2 + + def test_fan_out_empty_group_no_hang(self): + """Empty items in one group don't cause hangs.""" + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + + mock = MockLLM({"c1.implement": "code", "t1.write_tests": "test"}) + queue = InMemoryTaskQueue() + agents = [coder, tester] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def choreo(c, t): + return fan_out( + [ + ([], c, lambda c, m: c.implement(m)), # empty! + (["x"], t, lambda t, m: t.write_tests(m)), + ] + ) + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(coder, tester) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + r = results_map["c1"] + assert r[0] == [] + assert r[1] == ["test"] + + def test_fan_out_step_counter_sync(self): + """fan_out consumes exactly one step ID across all agent threads.""" + coder = Coder(agent_id="c1") + reviewer = Reviewer(agent_id="r1") + + mock = MockLLM( + { + "c1.implement": "code", + "r1.review": "lgtm", + } + ) + queue = InMemoryTaskQueue() + agents = [coder, reviewer] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def choreo(c, r): + # fan_out uses 1 step, then review uses 1 step + fan_results = fan_out( + [ + (["a"], c, lambda c, m: c.implement(m)), + ] + ) + verdict = r.review(str(fan_results)) + return verdict + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(coder, reviewer) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + # Both agents compute the same final result + assert results_map["c1"] == "lgtm" + assert results_map["r1"] == "lgtm" + + def test_fan_out_error_in_one_group(self): + """Error in one group's fn propagates as ChoreographyError.""" + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + + class ErrorMockLLM(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + key = f"{bound.__agent_id__}.{template.__name__}" if bound else "" + if key == "t1.write_tests": + raise RuntimeError("test failure!") + return "ok" + + mock = ErrorMockLLM() + + def program(c, t): + return fan_out( + [ + (["a"], c, lambda c, m: c.implement(m)), + (["x"], t, lambda t, m: t.write_tests(m)), + ] + ) + + choreo = Choreography( + program, + agents=[coder, tester], + handlers=[mock], + ) + with pytest.raises(ChoreographyError, match="test failure"): + choreo.run(c=coder, t=tester) + + def test_fan_out_all_orderings_agree(self): + """All thread scheduling orderings produce the same result.""" + for delays in [ + {"c1": 0.0, "t1": 0.05}, + {"c1": 0.05, "t1": 0.0}, + {"c1": 0.0, "t1": 0.0}, + ]: + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + + mock = DelayedMockLLM( + {"c1.implement": "code", "t1.write_tests": "test"}, + delays, + ) + queue = InMemoryTaskQueue() + agents = [coder, tester] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def choreo(c, t): + return fan_out( + [ + (["a", "b"], c, lambda c, m: c.implement(m)), + (["x"], t, lambda t, m: t.write_tests(m)), + ] + ) + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(coder, tester) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + assert results_map["c1"] == results_map["t1"] + + +class TestFanOutChoreography: + """fan_out via the high-level Choreography runner.""" + + def test_choreography_fan_out(self): + """fan_out works through Choreography.run().""" + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + + mock = MockLLM({"c1.implement": "code-out", "t1.write_tests": "test-out"}) + + def program(c, t): + return fan_out( + [ + (["a", "b"], c, lambda c, m: c.implement(m)), + (["x"], t, lambda t, m: t.write_tests(m)), + ] + ) + + choreo = Choreography( + program, + agents=[coder, tester], + handlers=[mock], + ) + result = choreo.run(c=coder, t=tester) + assert result == [["code-out", "code-out"], ["test-out"]] + + def test_choreography_fan_out_with_scatter(self): + """fan_out and scatter can be mixed in the same choreography.""" + architect = Architect(agent_id="arch") + coder = Coder(agent_id="c1") + tester = TesterAgent(agent_id="t1") + reviewer = Reviewer(agent_id="r1") + + mock = MockLLM( + { + "arch.plan": "plan-result", + "c1.implement": "code", + "t1.write_tests": "test", + "r1.review": "lgtm", + } + ) + + def program(arch, c, t, r): + plan = arch.plan("project") + # fan_out: code and test in parallel + code_results, test_results = fan_out( + [ + (["m1", "m2"], c, lambda c, m: c.implement(m)), + (["t1"], t, lambda t, m: t.write_tests(m)), + ] + ) + # Then scatter reviews sequentially + reviews = scatter(code_results, r, lambda r, code: r.review(code)) + return {"plan": plan, "reviews": reviews, "tests": test_results} + + choreo = Choreography( + program, + agents=[architect, coder, tester, reviewer], + handlers=[mock], + ) + result = choreo.run(arch=architect, c=coder, t=tester, r=reviewer) + assert result["plan"] == "plan-result" + assert result["reviews"] == ["lgtm", "lgtm"] + assert result["tests"] == ["test"] + + +class TestFanOutCrashTolerance: + """fan_out crash recovery with PersistentTaskQueue.""" + + def test_fan_out_crash_recovery_with_cached_results(self): + """Pre-cached fan_out items are reused on restart.""" + db_path = STATE_DIR / "fan_out_crash.db" + + # "Run 1": simulate partial completion + tq1 = PersistentTaskQueue(db_path) + # Pre-cache group 0 items + tq1.submit( + "fan-step-0000:g0", + {"group": 0, "item_index": 0}, + task_id="step-0000:g0:0000", + ) + tq1.claim_by_prefix("step-0000:g0:0000", "prior-worker") + tq1.complete("step-0000:g0:0000", "prior-worker", "cached-code") + + tq1.submit( + "fan-step-0000:g1", + {"group": 1, "item_index": 0}, + task_id="step-0000:g1:0000", + ) + tq1.claim_by_prefix("step-0000:g1:0000", "prior-worker") + tq1.complete("step-0000:g1:0000", "prior-worker", "cached-test") + + # "Run 2": restart with one uncached item per group + tq2 = PersistentTaskQueue(db_path) + c1 = Coder(agent_id="c1") + t1 = TesterAgent(agent_id="t1") + + mock = MockLLM({"c1.implement": "fresh-code", "t1.write_tests": "fresh-test"}) + + def choreo(c, t): + return fan_out( + [ + (["A", "B"], c, lambda c, m: c.implement(m)), + (["X", "Y"], t, lambda t, m: t.write_tests(m)), + ] + ) + + agents = [c1, t1] + ids = frozenset(a.__agent_id__ for a in agents) + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + def run(agent): + try: + tq2.release_stale_claims(agent.__agent_id__) + epp = EndpointProjection(agent, tq2, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo(c1, t1) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + r = results_map["c1"] + assert r[0][0] == "cached-code" + assert r[0][1] == "fresh-code" + assert r[1][0] == "cached-test" + assert r[1][1] == "fresh-test" + + def test_fan_out_concurrent_no_double_execution(self): + """No item is executed twice even under concurrent claiming.""" + c1 = Coder(agent_id="c1") + c2 = Coder(agent_id="c2") + t1 = TesterAgent(agent_id="t1") + t2 = TesterAgent(agent_id="t2") + + executed: list[str] = [] + exec_lock = threading.Lock() + + class TrackingMock(ObjectInterpretation): + @implements(Template.__apply__) + def _call(self, template, *args, **kwargs): + bound = get_bound_agent(template) + key = f"{bound.__agent_id__}.{template.__name__}" if bound else "" + with exec_lock: + executed.append(key) + return f"result-{key}" + + mock = TrackingMock() + queue = InMemoryTaskQueue() + agents = [c1, c2, t1, t2] + ids = frozenset(a.__agent_id__ for a in agents) + + results_map: dict[str, Any] = {} + errors: list[Exception] = [] + lock = threading.Lock() + + items_code = [f"mod-{i}" for i in range(10)] + items_test = [f"test-{i}" for i in range(8)] + + def choreo(coders, testers): + return fan_out( + [ + (items_code, coders, lambda c, m: c.implement(m)), + (items_test, testers, lambda t, m: t.write_tests(m)), + ] + ) + + def run(agent): + try: + epp = EndpointProjection(agent, queue, ids, poll_interval=0.02) + with handler(mock), handler(epp): + r = choreo([c1, c2], [t1, t2]) + with lock: + results_map[agent.__agent_id__] = r + except Exception as e: + with lock: + errors.append(e) + + _run_threads_with_timeout([lambda a=a: run(a) for a in agents]) + + assert not errors, errors + # Exactly 10 implement + 8 write_tests calls + impl_calls = [c for c in executed if ".implement" in c] + test_calls = [c for c in executed if ".write_tests" in c] + assert len(impl_calls) == 10 + assert len(test_calls) == 8 + # All agents see the same result + for aid in ["c1", "c2", "t1", "t2"]: + assert results_map[aid] == results_map["c1"]