pycastle 0.1.3.7.dev0__tar.gz → 0.1.3.9.dev0__tar.gz

{pycastle-0.1.3.7.dev0 → pycastle-0.1.3.9.dev0}/CONTEXT.md

@@ -109,14 +109,14 @@
 
 | Term | Definition | Aliases to avoid |
 | --- | --- | --- |
-| **agent output protocol** | The contract between prompts and the orchestrator: the set of XML tags agents emit to signal structured output (`<plan>`, `<issue>`, `<promise>`), plus the module that owns parsing and extraction of those tags | output format, agent tags, agent signals |
+| **agent output protocol** | The contract between prompts and the orchestrator: the set of XML tags agents emit to signal structured output (`<plan>`, `<issue>`, `<promise>`), plus the module that owns the complete NDJSON stream → typed output pipeline | output format, agent tags, agent signals |
 | **`<plan>` tag** | XML tag emitted by the Planner containing a JSON payload listing unblocked issues for the current iteration; extracted by the agent output protocol module | plan output, plan block |
 | **`<issue>` tag** | XML tag emitted by the preflight-issue agent containing the GitHub issue number it filed; extracted by the agent output protocol module | issue output, issue number tag |
 | **`<promise>COMPLETE</promise>`** | XML tag emitted by Implementers, Reviewers, and the Merger to declare that their work phase is complete; detected by the agent output protocol module | done signal, completion tag |
 | **`AgentOutputProtocolError`** | Base exception raised by the agent output protocol module when a required tag is missing or malformed; subclassed by `PlanParseError`, `IssueParseError`, and `PromiseParseError` | parse error, protocol error |
-| **`parse()`** | Entry point for data-bearing roles (Planner, preflight-issue agent); returns a typed output (`PlannerOutput` or `IssueOutput`) resolved statically by role; also checks for `<promise>COMPLETE</promise>` | protocol parser, output extractor |
-| **`assert_complete()`** | Entry point for completion-only roles (Implementer, Reviewer, Merger); verifies the `<promise>COMPLETE</promise>` tag is present and raises `PromiseParseError` if not; returns `None` on success — never a bool | parse_completion, check_promise |
-| **Claude streaming envelope** | The NDJSON format Claude Code uses for structured output; lines are JSON objects and the agent's final result is carried in the `{"type": "result", "result": "..."}` line; unwrapped internally by the agent output protocol module before tag extraction | streaming format, NDJSON output |
+| **`process_stream()`** | Single entry point in the agent output protocol module; accepts an iterable of decoded NDJSON lines, an `on_turn` callback, an `AgentRole`, and `usage_limit_patterns`; drives the per-line loop, emits complete assistant turns via the callback, raises `UsageLimitError` immediately on detection, unwraps the result envelope, and returns a typed `AgentOutput`; the container runner is the only caller — phases never call it directly | protocol entry point, stream processor |
+| **`on_turn` callback** | A `Callable[[str], None]` passed to `process_stream` by the container runner; invoked once per complete assistant turn during the Work phase; constructed by the container runner as a lambda over `StatusDisplay.print` so the agent output protocol module has no dependency on `StatusDisplay` | turn callback, display hook |
+| **Claude streaming envelope** | The NDJSON format Claude Code uses for structured output; lines are JSON objects and the agent's final result is carried in the `{"type": "result", "result": "..."}` line; unwrapped internally by `process_stream` before tag extraction | streaming format, NDJSON output |
 
 ## Agent Lifecycle
 
@@ -143,7 +143,7 @@
 | Term | Definition | Aliases to avoid |
 | --- | --- | --- |
 | **Dockerfile** | File in the pycastle directory defining the Docker image for agent containers — ships without baked-in credentials and without baked-in dev tools; system utilities (git, gh), Claude Code CLI, and the Python runtime are the only baked-in contents; all dev tools (e.g. ruff, mypy, pytest) must be declared in the consuming project's dependency file and are installed at runtime during the Setup phase | image definition |
-| **container runner** | Package module that manages Docker container lifecycle, injects runtime secrets, and drives the four agent lifecycle phases (Setup, Pre-flight, Prepare, Work) via instance methods; holds `status_display` at construction time so phase methods can update terminal state without caller involvement | docker wrapper |
+| **container runner** | Package module that manages Docker container lifecycle, injects runtime secrets, and drives the four agent lifecycle phases (Setup, Pre-flight, Prepare, Work) via instance methods; holds `status_display` at construction time so phase methods can update terminal state without caller involvement; during the Work phase owns Docker byte chunking, byte-to-line splitting, log writing, and idle timeout detection, then delegates the line stream to `process_stream` | docker wrapper |
 | **host repo** | The git repository on the developer's machine that is mounted into each agent container | project repo, local repo |
 | **volume mount** | A Docker bind mount attaching a host filesystem path to a container-internal path, with an explicit read/write mode | bind mount, volume |
 | **RO mount** | A volume mount with `mode: "ro"` — the container cannot write to it; used for the host repo | read-only mount |
@@ -158,10 +158,11 @@
 | **worktree contents check** | Guard step run after `git worktree add` that verifies `pyproject.toml` or `requirements.txt` is present; fails with the worktree path and directory listing if absent | checkout guard, file check |
 | **`detached_worktree`** | Async context manager in `worktree.py` that creates a detached checkout at a given SHA, yields the path, and guarantees removal in `__aexit__` regardless of outcome; used by `planning_phase` and `preflight_phase` for their sandbox worktrees | managed_worktree |
 | **`branch_worktree`** | Async context manager in `worktree.py` that creates a named-branch worktree at a given SHA, yields the path, and on exit removes the worktree and optionally deletes the branch; used by `merge_phase` for the merge-sandbox worktree | managed_worktree |
+| **`_agent_worktree`** | Async context manager in `implement.py` that owns the full Implementer and Reviewer worktree lifecycle; accepts a branch name, SHA, `CancellationToken`, and `Deps`; on entry creates the worktree and gitdir overlay; on exit conditionally removes the worktree based on `token.wants_worktree_preserved` and working-tree cleanliness, and always removes the gitdir overlay; used by `run_issue` twice per issue — once for the Implementer (new-branch path) and once for the Reviewer (existing-branch path); defined in `implement.py` not `worktree.py` because its cleanup policy depends on agent-lifecycle state (`CancellationToken`) rather than being unconditional | managed_worktree |
 | **`worktree_name_for_branch`** | Function in `worktree.py` that derives a short directory name from a branch string: extracts `issue-N` from `pycastle/issue-N-slug` or falls back to a sanitised slug; single authoritative definition replacing duplicated regex in `agent_runner` and `merge_phase` | — |
 | **`worktree_path`** | Function in `worktree.py` that constructs the host filesystem path for a named worktree at `<repo_root>/<pycastle_dir>/.worktrees/<name>`; single authoritative path expression replacing duplication across all phase modules | — |
 | **runtime injection** | The act of reading `~/.claude.json` from the host and writing it to `/home/agent/.claude.json` inside a container before the agent runs | baking in, build-time config |
-| **StreamParser** | Stateful class in `stream_parser.py` that processes stream-json lines one at a time; `feed(line: str) -> Optional[str]` returns the full text of a complete assistant turn (all text content blocks joined with `\n\n`) when the turn boundary is reached, or `None` to suppress; tool-use and tool-result content blocks are always dropped; one `StreamParser` instance per agent run | stream processor, message parser |
+| **StreamParser** | Retired — its assistant-turn assembly logic is now a private implementation detail of `process_stream` in the agent output protocol module; `stream_parser.py` no longer exists as a public module | stream processor, message parser |
 | **agent message** | The text content emitted by an agent during a single assistant turn; excludes tool-use and tool-result blocks; during the Work phase, printed to the console prefixed with the agent name and followed by a blank line; not shown in the status panel | assistant message, agent output |
 | **PycastleError** | Base exception class for all pycastle domain errors | — |
 | **DockerError** | Error raised when a Docker operation (container start, stop, remove) fails | container error |
@@ -186,11 +187,12 @@
 | **GithubService** | Service that encapsulates `gh` CLI calls for GitHub issue operations: closing issues, querying parent issues, listing open sub-issues, and reading issue labels | GitHub wrapper, gh provider |
 | **Logger** | Injectable abstraction that owns all structured log output for one iteration; exposes named channels (`log_error`, `log_agent_output`) each writing to a dedicated file under `logs/`; injected via `Deps` so tests never touch the filesystem | log writer, output handler |
 | **RecordingLogger** | Test double for `Logger` that records every call in memory; tests assert on recorded calls rather than capturing stderr or reading log files | mock logger, spy logger |
-| **StatusDisplay** | Injectable abstraction that owns the live terminal status line; exposes `add_agent(name, phase, work_body="")`, `update_phase`, `reset_idle_timer`, `remove_agent`, and `print` methods; backed by a `rich` `Live` display in production and a `NullStatusDisplay` no-op in tests; injected via `Deps` as a separate concern from `Logger` | terminal display, status bar |
-| **work_body** | The caller-constructed string passed as the third argument to `add_agent`; displayed in the body column of the agent status row during the Work phase; empty string for agents that do not reach Work | — |
-| **NullStatusDisplay** | Test/no-op implementation of `StatusDisplay`; all methods are no-ops; `print` falls back to `builtins.print` | — |
-| **agent status row** | One headerless line in the `StatusDisplay` live panel representing one active agent; format: `elapsed \| RoleName \| idle \| body`; elapsed is dim and right-justified so all role names align across rows; role name is bold, with any numeric part (e.g. `5` in `Implementer #5`) styled bold cyan; idle is dim; body column shows the current lifecycle phase name (e.g. `Setup`, `Running ruff Checks`, `Prepare`) for all non-Work states, or the `work_body` string during Work (caller-constructed: `"implementing \"{issue_title}\""` for Implementers, `"reviewing \"{issue_title}\""` for Reviewers, `"Creating Plan from {n} issues"` for the Planner, `"Merging {n} Branches"` for the Merger, `"Checking"` for the preflight-checks agent, `"reporting {check_name} issue"` for the preflight-issue agent); elapsed counts up from agent start and never resets; idle resets on each Docker stream chunk; the live panel is preceded by one blank line to visually separate it from scrollback output; ordered by orchestration phase (plan → implement → review → merge) then by issue number; disappears when the agent finishes | status entry, agent row |
-| **IterationOutcome** | Sealed return type of `run_iteration()`; one of four variants: `Continue` (iteration completed, keep looping), `Done` (no issues found, stop cleanly), `AbortedHITL` (HITL verdict — carries `issue_number`; orchestrator exits non-zero), `AbortedUsageLimit` (token ceiling hit — worktrees preserved, safe to retry; orchestrator exits non-zero) | iteration result, loop result |
+| **StatusDisplay** | Injectable abstraction that owns the live terminal status panel and all formatted terminal output; exposes `register(caller, startup_message="started", work_body="")`, `update_phase`, `reset_idle_timer`, `remove(caller, shutdown_message="finished", shutdown_style="success")`, and `print(caller, message, style=None)` methods; backed by a `rich` `Live` display in production and a `PlainStatusDisplay` in tests; injected via `Deps` as a separate concern from `Logger`; defined in `status_display` module | terminal display, status bar |
+| **caller** | The identity string passed as the first argument to `StatusDisplay.register`, `remove`, and `print`; rendered as a `[Caller]` prefix on every terminal output line; empty string `""` is the anonymous caller — no brackets are printed and the message is output as-is; a blank line is inserted before any output call (`register`, `remove`, or `print`) when the caller differs from the previous one, or unconditionally when the caller is `""` (anonymous outputs always stand alone); canonical callers — phase rows: `"Preflight"`, `"Plan"`, `"Implement"`, `"Merge"`; agents: `"Preflight Agent"`, `"Plan Agent"`, `"Implement Agent #N"`, `"Review Agent #N"`, `"Merge Agent"` | source, label |
+| **work_body** | The caller-constructed string passed as the third argument to `register`; displayed in the body column of the status row during the Work phase; empty string for callers that do not reach Work | — |
+| **PlainStatusDisplay** | Plain-terminal adapter for `StatusDisplay` defined in `status_display` module; panel methods (`update_phase`, `reset_idle_timer`) are no-ops; `register` and `remove` print their startup/shutdown messages; `print(caller, message, style=None)` formats output as `[Caller] message` with no ANSI colour codes, no bold, and style ignored; used in tests so assertions can match the full formatted line | NullStatusDisplay |
+| **status row** | One headerless line in the `StatusDisplay` live panel; created by `register` and removed by `remove`; two kinds: **agent rows** (one per active agent — `"Preflight Agent"`, `"Plan Agent"`, `"Implement Agent #N"`, `"Review Agent #N"`, `"Merge Agent"`) and **phase rows** (one per active phase — `"Preflight"`, `"Plan"`, `"Implement"`, `"Merge"`); phase rows and agent rows within the same phase coexist; format: `elapsed \| Name \| idle \| body`; elapsed is dim and right-justified; name is bold with any numeric part styled bold cyan; idle is dim; body shows the current lifecycle phase name for all non-Work states, or the `work_body` string during Work; elapsed counts up from `register` and never resets; idle resets on each Docker stream chunk; the live panel is preceded by one blank line to visually separate it from scrollback; ordered by orchestration phase (plan → implement → review → merge) then by issue number | agent status row, status entry, agent row |
+| **IterationOutcome** | Sealed return type of `run_iteration()`; one of four variants: `Continue` (iteration completed, keep looping), `Done` (no issues found, stop cleanly), `AbortedHITL` (HITL verdict — carries `issue_number`; orchestrator exits non-zero), `AbortedUsageLimit` (token ceiling hit — worktrees preserved; orchestrator sleeps until 2 minutes past the next local-time full hour, then continues the loop to retry the current issue from scratch; repeats indefinitely on consecutive hits) | iteration result, loop result |
 
 ## Test Anti-Patterns (Red Flags)
 
@@ -216,14 +218,15 @@
 - The **Planner** and all **Implementer** worktrees are created from the pinned **safe SHA**, never from HEAD directly; this guarantees every agent sees the same verified-clean committed state regardless of external commits that land on main after preflight passes.
 - In **sequential mode** (`max_parallel = 1`), the iteration processes issues one by one: after each issue's merge the safe SHA is re-pinned to the new HEAD, and the next Implementer starts from that SHA; a failed issue is skipped (remains `ready-for-agent`) and the queue continues; the Merger remains available as a fallback for unexpected conflicts; no additional pre-flight checks run between issues.
 - The **Pre-flight phase** (agent lifecycle) runs quality checks inside the container and returns a list of failure tuples to the orchestrator; it never spawns agents internally.
-- An **orphan sweep** runs once at orchestrator startup; **collision detection** holds a per-branch lock for the full duration of each agent run.
-- **`detached_worktree`** is used by `planning_phase` (for the plan-sandbox worktree) and `preflight_phase` (for the pre-flight-sandbox worktree); **`branch_worktree`** is used by `merge_phase` (for the merge-sandbox worktree); Implementer and Reviewer worktrees are still managed directly inside `agent_runner` because their cleanup is conditional on cancellation state. **`worktree_path`** and **`worktree_name_for_branch`** are the single authoritative path and name expressions used by all of the above.
+- An **orphan sweep** runs once at orchestrator startup; **collision detection** uses a per-branch `asyncio.Lock` held in `implement_phase` for the full duration of each `run_issue` call — from first worktree creation to final worktree teardown.
+- **`detached_worktree`** is used by `planning_phase` (for the plan-sandbox worktree) and `preflight_phase` (for the pre-flight-sandbox worktree); **`branch_worktree`** is used by `merge_phase` (for the merge-sandbox worktree); **`_agent_worktree`** is used by `run_issue` in `implement.py` for Implementer and Reviewer worktrees — its cleanup is conditional on cancellation state, unlike the unconditional teardown in `detached_worktree` and `branch_worktree`. **`worktree_path`** and **`worktree_name_for_branch`** are the single authoritative path and name expressions used by all of the above.
 - Host mounts per container: host repo → RO at `/home/agent/repo`; worktree → RW at `/home/agent/workspace`; `<host-repo>/.git` → RW at `/.pycastle-parent-git`; on Windows, gitdir overlay → RO over `/home/agent/workspace/.git`.
 - A **Service** defines a Custom exception hierarchy so callers never handle raw subprocess exceptions; tests inject Default implementations from a test fixture and override per-test for error paths.
 - **StatusDisplay** is a separate injectable in `Deps` alongside `Logger`; `Logger` owns file I/O, `StatusDisplay` owns the live terminal UI — they never overlap.
-- The **agent status row** is created when the Setup phase begins; the row is removed when the agent finishes or errors; the `rich` `Live` display is started on the first `add_agent` call and stopped after the last `remove_agent` call.
+- Rich markup (e.g. `[red]...[/red]`) must never be embedded in a `StatusDisplay.print` message string; colouring is expressed exclusively via the `style` parameter (`"error"`, `"success"`).
+- A **status row** is created by `StatusDisplay.register` and removed by `StatusDisplay.remove`; phase rows are registered at the start of each orchestration phase and removed at its end; agent rows are registered at container Setup and removed when the agent finishes or errors; the `rich` `Live` display is started on the first `register` call and stopped after the last `remove` call.
 - All orchestrator-level terminal output (e.g. "Planning complete…") is routed through `StatusDisplay.print()` so `rich` can coordinate it with the live panel; bare `print()` calls are not used while a `StatusDisplay` is active.
-- The streaming loop calls `reset_idle_timer` on every Docker chunk (matching the kill-agent timeout behavior); each complete line is fed to `StreamParser.feed()` — during the Work phase, if a complete assistant turn is returned, it is printed to the console via `StatusDisplay.print()` with the agent-name prefix and a trailing blank line; tool-use blocks are silently dropped; Setup, Pre-flight, and Prepare phases produce no console output — their activity is reflected only in the body column of the agent status row.
+- During the Work phase the container runner owns byte chunking, byte-to-line splitting, log writing, and idle timeout detection; it passes the decoded NDJSON line stream and an **`on_turn` callback** to **`process_stream`**, which assembles assistant turns (invoking the callback for each), detects usage limit lines and raises `UsageLimitError` immediately, unwraps the result envelope, and returns a typed `AgentOutput`; phases receive `AgentOutput` directly from `AgentRunner.run()` — no phase calls `parse()` or `assert_complete()`. Setup, Pre-flight, and Prepare phases produce no console output — their activity is reflected only in the body column of the agent status row.
 
 ## Example dialogue
 

{pycastle-0.1.3.7.dev0 → pycastle-0.1.3.9.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pycastle
-Version: 0.1.3.7.dev0
+Version: 0.1.3.9.dev0
 Summary: Python orchestrator for autonomous Claude Code agents in Docker
 License: MIT License
         

{pycastle-0.1.3.7.dev0 → pycastle-0.1.3.9.dev0}/src/pycastle/agent_output_protocol.py

@@ -2,7 +2,10 @@ import dataclasses
 import enum
 import json
 import re
-from typing import Literal, TypeAlias, overload
+from collections.abc import Callable, Iterable
+from typing import TypeAlias
+
+from .errors import UsageLimitError
 
 
 class AgentRole(enum.Enum):
@@ -48,21 +51,6 @@ class PromiseParseError(AgentOutputProtocolError):
     pass
 
 
-def _unwrap(output: str) -> str:
-    for line in output.splitlines():
-        line = line.strip()
-        if not line:
-            continue
-        try:
-            obj = json.loads(line)
-        except json.JSONDecodeError:
-            continue
-        if isinstance(obj, dict) and obj.get("type") == "result":
-            result = obj.get("result")
-            return result if isinstance(result, str) else output
-    return output
-
-
 def _extract_planner_output(text: str) -> PlannerOutput:
     match = re.search(r"<plan>([\s\S]*?)</plan>", text)
     if not match:
@@ -111,20 +99,80 @@ def _extract_issue_output(text: str) -> IssueOutput:
     return IssueOutput(labels=labels, number=number)
 
 
-@overload
-def parse(output: str, role: Literal[AgentRole.PLANNER]) -> PlannerOutput: ...
-
-
-@overload
-def parse(output: str, role: Literal[AgentRole.PREFLIGHT_ISSUE]) -> IssueOutput: ...
-
-
-@overload
-def parse(output: str, role: AgentRole) -> AgentOutput: ...
-
-
-def parse(output: str, role: AgentRole) -> AgentOutput:
-    text = _unwrap(output)
+def _is_usage_limit_line(line: str, patterns: tuple[str, ...]) -> bool:
+    try:
+        obj = json.loads(line)
+        if isinstance(obj, dict):
+            if obj.get("type") == "result" and obj.get("is_error"):
+                if obj.get("api_error_status") == 429:
+                    return True
+                result_text = obj.get("result")
+                if isinstance(result_text, str) and any(
+                    p.lower() in result_text.lower() for p in patterns
+                ):
+                    return True
+            return False
+    except json.JSONDecodeError:
+        pass
+    line_lower = line.lower()
+    return any(p.lower() in line_lower for p in patterns)
+
+
+def _extract_turn(line: str) -> str | None:
+    try:
+        obj = json.loads(line)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(obj, dict) or obj.get("type") != "assistant":
+        return None
+    content = (obj.get("message") or {}).get("content") or []
+    parts: list[str] = []
+    for block in content:
+        if isinstance(block, dict) and block.get("type") == "text":
+            text = (block.get("text") or "").strip()
+            if text:
+                parts.append(text)
+    return "\n\n".join(parts) if parts else None
+
+
+def process_stream(
+    lines: Iterable[str],
+    on_turn: Callable[[str], None],
+    role: AgentRole,
+    usage_limit_patterns: tuple[str, ...],
+) -> AgentOutput:
+    collected: list[str] = []
+    result_text: str | None = None
+    for line in lines:
+        collected.append(line)
+        if _is_usage_limit_line(line, usage_limit_patterns):
+            raise UsageLimitError(line)
+        turn = _extract_turn(line)
+        if turn is not None:
+            on_turn(turn)
+            if role in (AgentRole.IMPLEMENTER, AgentRole.REVIEWER, AgentRole.MERGER):
+                if re.search(r"<promise>COMPLETE</promise>", turn):
+                    return CompletionOutput()
+            elif role == AgentRole.PLANNER:
+                try:
+                    return _extract_planner_output(turn)
+                except PlanParseError:
+                    pass
+            elif role == AgentRole.PREFLIGHT_ISSUE:
+                try:
+                    return _extract_issue_output(turn)
+                except IssueParseError:
+                    pass
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(obj, dict) and obj.get("type") == "result":
+            r = obj.get("result")
+            if isinstance(r, str):
+                result_text = r
+                break
+    text = result_text if result_text is not None else "\n".join(collected)
     tail = f"\nOutput tail: {text[-300:]!r}"
     if role == AgentRole.PREFLIGHT_ISSUE:
         try:
@@ -141,12 +189,3 @@ def parse(output: str, role: AgentRole) -> AgentOutput:
             f"Agent produced no <promise>COMPLETE</promise> tag.{tail}"
         )
     return CompletionOutput()
-
-
-def assert_complete(output: str) -> None:
-    text = _unwrap(output)
-    if not re.search(r"<promise>COMPLETE</promise>", text):
-        tail = text[-200:]
-        raise PromiseParseError(
-            f"Agent produced no <promise>COMPLETE</promise> tag. Output tail: {tail!r}"
-        )

pycastle-0.1.3.9.dev0/src/pycastle/agent_runner.py

@@ -0,0 +1,151 @@
+import dataclasses
+from pathlib import Path
+from typing import Any, Protocol
+
+from .agent_output_protocol import AgentOutput, AgentRole
+from .agent_result import CancellationToken, PreflightFailure
+from .config import Config
+from .container_runner import ContainerRunner
+from .errors import AgentTimeoutError, UsageLimitError
+from .services import GitService
+from .status_display import PlainStatusDisplay
+
+
+@dataclasses.dataclass
+class RunRequest:
+    name: str
+    prompt_file: Path
+    mount_path: Path
+    role: AgentRole = AgentRole.IMPLEMENTER
+    prompt_args: dict[str, str] | None = None
+    skip_preflight: bool = False
+    model: str = ""
+    effort: str = ""
+    stage: str = ""
+    token: CancellationToken | None = None
+    status_display: Any = None
+    issue_title: str = ""
+    work_body: str = ""
+
+
+class AgentRunnerProtocol(Protocol):
+    async def run(self, request: RunRequest) -> AgentOutput | PreflightFailure: ...
+
+    async def run_preflight(
+        self,
+        *,
+        name: str,
+        mount_path: Path,
+        stage: str = "",
+        status_display=None,
+        work_body: str = "",
+    ) -> list[tuple[str, str, str]]: ...
+
+
+class AgentRunner:
+    def __init__(
+        self,
+        env: dict[str, str],
+        cfg: Config,
+        git_service: GitService,
+        docker_client=None,
+    ) -> None:
+        self._env = env
+        self._cfg = cfg
+        self._git_service = git_service
+        self._docker_client = docker_client
+
+    async def run(self, request: RunRequest) -> AgentOutput | PreflightFailure:
+        name = request.name
+        prompt_file = request.prompt_file
+        mount_path = request.mount_path
+        prompt_args = request.prompt_args
+        skip_preflight = request.skip_preflight
+        model = request.model
+        effort = request.effort
+        token = request.token
+        status_display = request.status_display
+        work_body = request.work_body
+
+        if status_display is None:
+            status_display = PlainStatusDisplay()
+
+        _token = token if token is not None else CancellationToken()
+        if _token.is_cancelled:
+            raise UsageLimitError("Agent cancelled due to usage limit")
+
+        runner = ContainerRunner(
+            name,
+            mount_path,
+            self._env,
+            model=model,
+            effort=effort,
+            docker_client=self._docker_client,
+            status_display=status_display,
+            cfg=self._cfg,
+        )
+        try:
+            git_name = self._git_service.get_user_name()
+            git_email = self._git_service.get_user_email()
+            await runner.setup(git_name, git_email, work_body)
+            await runner.prepare(prompt_file, prompt_args or {})
+            if not skip_preflight:
+                failures = await runner.preflight(list(self._cfg.preflight_checks))
+                if failures:
+                    return PreflightFailure(failures=tuple(failures))
+            retries_left = self._cfg.timeout_retries
+            while True:
+                try:
+                    output = await runner.work(request.role)
+                    return output
+                except AgentTimeoutError:
+                    if retries_left <= 0:
+                        raise
+                    restart_num = self._cfg.timeout_retries - retries_left + 1
+                    status_display.print(
+                        name,
+                        f"Timeout — restarting"
+                        f" (attempt {restart_num}/{self._cfg.timeout_retries})",
+                    )
+                    retries_left -= 1
+                except UsageLimitError:
+                    _token.cancel(preserve_worktree=True)
+                    raise
+        finally:
+            status_display.remove(name)
+            try:
+                runner.__exit__(None, None, None)
+            except Exception:
+                pass
+
+    async def run_preflight(
+        self,
+        *,
+        name: str,
+        mount_path: Path,
+        stage: str = "",
+        status_display=None,
+        work_body: str = "",
+    ) -> list[tuple[str, str, str]]:
+        if status_display is None:
+            status_display = PlainStatusDisplay()
+
+        git_name = self._git_service.get_user_name()
+        git_email = self._git_service.get_user_email()
+        runner = ContainerRunner(
+            name,
+            mount_path,
+            self._env,
+            docker_client=self._docker_client,
+            status_display=status_display,
+            cfg=self._cfg,
+        )
+        try:
+            await runner.setup(git_name, git_email, work_body)
+            return await runner.preflight(list(self._cfg.preflight_checks))
+        finally:
+            status_display.remove(name)
+            try:
+                runner.__exit__(None, None, None)
+            except Exception:
+                pass

{pycastle-0.1.3.7.dev0 → pycastle-0.1.3.9.dev0}/src/pycastle/container_runner.py

@@ -1,6 +1,5 @@
 import asyncio
 import io
-import json
 import os
 import queue
 import re
@@ -8,12 +7,13 @@ import shlex
 import sys
 import tarfile
 import threading
+from collections.abc import Callable, Generator
 from pathlib import Path
 
 import docker
-from rich.text import Text
 from docker.models.containers import Container as DockerContainer
 
+from .agent_output_protocol import AgentOutput, AgentRole, process_stream
 from .config import Config
 from .errors import (
     AgentTimeoutError,
@@ -21,33 +21,13 @@ from .errors import (
     DockerTimeoutError,
     UsageLimitError,
 )
-from .stream_parser import StreamParser
+from .status_display import PlainStatusDisplay
 from .worktree import (
     CONTAINER_PARENT_GIT,
     patch_gitdir_for_container,
 )
 
 
-def _is_usage_limit_line(line: str, patterns: tuple[str, ...]) -> bool:
-    """Return True if line signals a usage limit — plain-text or a JSON result error."""
-    try:
-        obj = json.loads(line)
-        if isinstance(obj, dict):
-            if obj.get("type") == "result" and obj.get("is_error"):
-                if obj.get("api_error_status") == 429:
-                    return True
-                result_text = obj.get("result")
-                if isinstance(result_text, str) and any(
-                    p.lower() in result_text.lower() for p in patterns
-                ):
-                    return True
-            return False
-    except json.JSONDecodeError:
-        pass
-    line_lower = line.lower()
-    return any(p.lower() in line_lower for p in patterns)
-
-
 def _build_claude_command(model: str = "", effort: str = "") -> str:
     flags = "--verbose --dangerously-skip-permissions --output-format stream-json -p -"
     if model:
@@ -57,16 +37,6 @@ def _build_claude_command(model: str = "", effort: str = "") -> str:
     return f"claude {flags} < /tmp/.pycastle_prompt"
 
 
-def _build_agent_prefix(name: str) -> Text:
-    """Return a styled Text object for the agent output prefix, e.g. ``[Implementer #1] ``."""
-    msg = Text()
-    msg.append("[", style="bold")
-    for segment in re.split(r"(\d+)", name):
-        if segment:
-            msg.append(segment, style="bold cyan" if segment.isdigit() else "bold")
-    msg.append("] ", style="bold")
-    return msg
-
 
 class ContainerRunner:
     def __init__(
@@ -94,9 +64,7 @@ class ContainerRunner:
         self.effort = effort
         self._cfg = cfg
         if status_display is None:
-            from .iteration._deps import NullStatusDisplay
-
-            status_display = NullStatusDisplay()
+            status_display = PlainStatusDisplay()
         self._status_display = status_display
         self._owns_client = docker_client is None
         self._client = docker_client if docker_client is not None else docker.from_env()
@@ -268,7 +236,7 @@ class ContainerRunner:
     async def setup(self, git_name: str, git_email: str, work_body: str = "") -> None:
         loop = asyncio.get_running_loop()
         await loop.run_in_executor(None, self.__enter__)
-        self._status_display.add_agent(self.name, "Setup", work_body)
+        self._status_display.register(self.name, work_body=work_body)
         await loop.run_in_executor(
             None,
             self.exec_simple,
@@ -309,16 +277,19 @@ class ContainerRunner:
 
         self._prompt = await prepare_prompt(prompt_file, prompt_args, container_exec)
 
-    async def work(self) -> str:
+    async def work(self, role: AgentRole) -> AgentOutput:
         self._status_display.update_phase(self.name, "Work")
         loop = asyncio.get_running_loop()
+        on_turn: Callable[[str], None] = lambda turn: self._status_display.print(
+            self.name, turn
+        )
         return await loop.run_in_executor(
-            None, lambda: self.run_streaming(print_output=True)
+            None, lambda: self.run_streaming(role=role, on_turn=on_turn)
         )
 
-    def run_streaming(self, print_output: bool = False) -> str:
+    def run_streaming(self, role: AgentRole, on_turn: Callable[[str], None]) -> AgentOutput:
         self.write_file(self._prompt, "/tmp/.pycastle_prompt")
-        result = self._active_container.exec_run(
+        exec_result = self._active_container.exec_run(
             ["bash", "-c", _build_claude_command(model=self.model, effort=self.effort)],
             stream=True,
             workdir=self._worktree_path,
@@ -329,18 +300,17 @@ class ContainerRunner:
 
         def _feed():
             try:
-                for chunk in result.output:
+                for chunk in exec_result.output:
                     q.put(chunk)
             finally:
                 q.put(_sentinel)
 
         threading.Thread(target=_feed, daemon=True).start()
 
-        parts: list[str] = []
-        line_buf = ""
-        parser = StreamParser()
+        log = open(self._log_path, "wb")  # noqa: WPS515
         try:
-            with open(self._log_path, "wb") as log:
+            def _lines() -> Generator[str, None, None]:
+                line_buf = ""
                 while True:
                     try:
                         chunk = q.get(timeout=self._cfg.idle_timeout)
@@ -349,23 +319,19 @@ class ContainerRunner:
                             f"Agent idle for more than {self._cfg.idle_timeout}s"
                         )
                     if chunk is _sentinel:
-                        break
+                        return
                     log.write(chunk)
                     log.flush()
                     text = chunk.decode("utf-8", errors="replace")
-                    parts.append(text)
                     self._status_display.reset_idle_timer(self.name)
                     line_buf += text
                     while "\n" in line_buf:
                         line, line_buf = line_buf.split("\n", 1)
-                        if _is_usage_limit_line(line, self._cfg.usage_limit_patterns):
-                            raise UsageLimitError(line)
-                        turn = parser.feed(line)
-                        if print_output and turn is not None:
-                            msg = _build_agent_prefix(self.name)
-                            msg.append(f"{turn}\n")
-                            self._status_display.print(msg, source=self.name)
+                        yield line
+
+            return process_stream(_lines(), on_turn, role, self._cfg.usage_limit_patterns)
         finally:
+            log.close()
             try:
                 self._active_container.exec_run(
                     ["bash", "-c", "rm -f /tmp/.pycastle_prompt"],
@@ -373,5 +339,3 @@ class ContainerRunner:
                 )
             except Exception:
                 pass
-        return "".join(parts)
-