From e36a36f4846b9820b0cafa1ee0bdbb50c5f31ab8 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Mon, 14 Sep 2026 17:33:12 +0400 Subject: [PATCH 1/8] feat: tag subagent runs with identity and persist bounded activity Every child run now gets a stable run id plus its task description, forwarded with each progress event. Completed, failed, and cancelled runs append one bounded agent_run record (answer, prompt, per-tool summaries with caps and truncation markers) to the parent session, with a tombstone-aware reader. Direct @agent exchanges persist as activity without entering model history. --- src/lecode/agent/tools/task.py | 9 ++- src/lecode/extras/subagents.py | 130 +++++++++++++++++++++++++++++++-- src/lecode/session/storage.py | 67 +++++++++++++++++ src/lecode/tui/app.py | 20 +++-- tests/test_session_storage.py | 75 +++++++++++++++++++ tests/test_subagents.py | 130 ++++++++++++++++++++++++++++++++- 6 files changed, 417 insertions(+), 14 deletions(-) diff --git a/src/lecode/agent/tools/task.py b/src/lecode/agent/tools/task.py index ae6d788..294818d 100644 --- a/src/lecode/agent/tools/task.py +++ b/src/lecode/agent/tools/task.py @@ -84,6 +84,7 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: agents, name=str(args.get("agent") or DEFAULT_AGENT), prompt=prompt, + description=str(args.get("description") or ""), on_event=ctx.extras.get(SUBAGENT_EVENTS_EXTRA), ) except SubagentError as e: @@ -121,7 +122,13 @@ async def body(emit: Any) -> tuple[str, int | None]: # run_subagent builds fresh child extras itself — the parent's # "conversation" seam is never clobbered. outcome = await run_subagent( - ctx, registry, agents, name=agent_name, prompt=prompt, on_event=on_event + ctx, + registry, + agents, + name=agent_name, + prompt=prompt, + description=description, + on_event=on_event, ) except SubagentError as e: return f"error: {e}", 1 diff --git a/src/lecode/extras/subagents.py b/src/lecode/extras/subagents.py index 668b4d3..f7db03a 100644 --- a/src/lecode/extras/subagents.py +++ b/src/lecode/extras/subagents.py @@ -16,13 +16,16 @@ from __future__ import annotations import asyncio +import contextlib import dataclasses +import time +import uuid from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any from lecode.agent.prompts import build_system_prompt -from lecode.agent.runner import AgentRunner +from lecode.agent.runner import AgentRunner, ToolCall, ToolResult from lecode.agent.tools.base import ToolContext, ToolRegistry from lecode.hooks import SUBAGENT_END, SUBAGENT_START, build_envelope, dispatch_event from lecode.providers.openai_compat import ProviderError @@ -52,8 +55,23 @@ #: Clip for the SubagentEnd hook envelope's result content. _END_CONTENT_CLIP = 2000 -#: ``on_event(agent_name, event)`` — one call per child runner event. -OnSubagentEvent = Callable[[str, "AgentEvent"], Any] + +@dataclass(frozen=True) +class SubagentProgress: + """One child runner event, tagged with the run's stable identity. + + ``run_id`` distinguishes two concurrent runs of the same agent; + ``description`` is the human label the caller supplied (task description). + """ + + run_id: str + agent: str + description: str + event: AgentEvent + + +#: ``on_event(progress)`` — one call per child runner event. +OnSubagentEvent = Callable[[SubagentProgress], Any] class SubagentError(Exception): @@ -70,6 +88,8 @@ class SubagentOutcome: input_tokens: int output_tokens: int cost_usd: float + run_id: str = "" + description: str = "" def child_registry(parent: ToolRegistry) -> ToolRegistry: @@ -99,6 +119,52 @@ async def _fire_hook( await dispatch_event(event, envelope, handlers) +def _persist_agent_run( + ctx: ToolContext, + *, + run_id: str, + agent: str, + description: str, + prompt: str, + status: str, + answer: str, + error: str | None, + trail: dict[str, dict[str, Any]], + trail_order: list[str], + duration_s: float, + result: RunResult | None, +) -> None: + """Append the run's bounded activity trail to the parent session. + + Fail-open: activity is display-only, so a storage failure must not break + the run or mask its outcome. + """ + store = ctx.session_store + session = ctx.session + if store is None or session is None: + return + run: dict[str, Any] = { + "run_id": run_id, + "agent": agent, + "description": description, + "prompt": prompt, + "status": status, + "answer": answer or error or "", + "tool_calls": [trail[call_id] for call_id in trail_order], + "duration_s": duration_s, + } + if result is not None: + totals = result.usage_totals + run.update( + turns=result.turns, + input_tokens=totals.input_tokens, + output_tokens=totals.output_tokens, + cost_usd=totals.cost_usd, + ) + with contextlib.suppress(Exception): + store.record_agent_run(session, run) + + async def run_subagent( ctx: ToolContext, parent_registry: ToolRegistry, @@ -106,13 +172,16 @@ async def run_subagent( *, name: str, prompt: str, + description: str = "", on_event: OnSubagentEvent | None = None, ) -> SubagentOutcome: """Run subagent ``name`` on ``prompt``; returns its final text and usage. - Raises :class:`SubagentError` for unknown agents, a missing provider, - timeouts, and provider failures; ``asyncio.CancelledError`` propagates - so a cancelled parent turn takes the child down with it. + ``description`` labels the run for the UI and the persisted activity + record; empty falls back to a prompt preview. Raises :class:`SubagentError` + for unknown agents, a missing provider, timeouts, and provider failures; + ``asyncio.CancelledError`` propagates so a cancelled parent turn takes the + child down with it. """ available = [a.name for a in agents.subagents()] agent = agents.get(name) @@ -124,6 +193,9 @@ async def run_subagent( if provider is None: raise SubagentError("no provider available for subagents") + run_id = uuid.uuid4().hex[:8] + label = description or prompt[:60] + checker = ctx.permission_checker if agent.overlay is not None: checker = checker.for_agent(agent.overlay) @@ -152,17 +224,45 @@ async def run_subagent( child = AgentRunner(provider, child_registry(parent_registry), child_ctx, config=child_config) child.model = agent.model or ctx.config.agent.subagent_model or ctx.config.llm.model - forward = (lambda event: on_event(name, event)) if on_event is not None else None + # Bounded activity trail for the persisted record (tool calls paired by id). + trail: dict[str, dict[str, Any]] = {} + trail_order: list[str] = [] + + def forward(event: AgentEvent) -> Any: + if isinstance(event, ToolCall): + trail[event.id] = { + "name": event.name, + "args": event.arguments, + "result": "", + "is_error": False, + } + trail_order.append(event.id) + elif isinstance(event, ToolResult): + entry = trail.setdefault( + event.id, + {"name": event.name, "args": "", "result": "", "is_error": False}, + ) + entry["result"] = event.content + entry["is_error"] = event.is_error + if on_event is None: + return None + return on_event(SubagentProgress(run_id=run_id, agent=name, description=label, event=event)) + messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ] + started_at = time.monotonic() await _fire_hook(ctx, SUBAGENT_START, name, prompt=prompt) result: RunResult | None = None error: str | None = None + cancelled = False try: async with asyncio.timeout(SUBAGENT_TIMEOUT_S): result = await child.run(messages, on_event=forward) + except asyncio.CancelledError: + cancelled = True + raise except TimeoutError as e: error = f"subagent '{name}' timed out after {SUBAGENT_TIMEOUT_S:.0f}s" raise SubagentError(error) from e @@ -170,6 +270,20 @@ async def run_subagent( error = str(e) raise SubagentError(error) from e finally: + _persist_agent_run( + ctx, + run_id=run_id, + agent=name, + description=label, + prompt=prompt, + status="cancelled" if cancelled else ("ok" if result is not None else "error"), + answer=result.final_text if result is not None else "", + error=error, + trail=trail, + trail_order=trail_order, + duration_s=time.monotonic() - started_at, + result=result, + ) end_content = result.final_text[:_END_CONTENT_CLIP] if result is not None else error await _fire_hook( ctx, @@ -189,4 +303,6 @@ async def run_subagent( input_tokens=totals.input_tokens, output_tokens=totals.output_tokens, cost_usd=totals.cost_usd, + run_id=run_id, + description=label, ) diff --git a/src/lecode/session/storage.py b/src/lecode/session/storage.py index b47a562..4ae942b 100644 --- a/src/lecode/session/storage.py +++ b/src/lecode/session/storage.py @@ -39,6 +39,57 @@ except ImportError: # Windows: best-effort, attach locking disabled fcntl = None # type: ignore[assignment] +#: Caps for one persisted agent-run activity record (bounded trails). +AGENT_RUN_ANSWER_CAP = 32 * 1024 +AGENT_RUN_PROMPT_CAP = 2000 +AGENT_RUN_RESULT_CAP = 2000 +AGENT_RUN_ARGS_CAP = 500 +AGENT_RUN_MAX_TOOLS = 100 +_TRUNCATED_MARK = "… (truncated)" + + +def _clip_activity(text: str, cap: int) -> tuple[str, bool]: + """Clip ``text`` to ``cap`` chars; the flag reports a truncation.""" + if len(text) <= cap: + return text, False + return text[:cap] + f"\n{_TRUNCATED_MARK}", True + + +def _bounded_agent_run(run: dict[str, Any]) -> dict[str, Any]: + """Apply the activity caps so one run can never bloat the session file.""" + truncated = False + data: dict[str, Any] = { + "run_id": str(run.get("run_id") or ""), + "agent": str(run.get("agent") or ""), + "description": str(run.get("description") or ""), + "status": str(run.get("status") or ""), + } + for key, cap in (("prompt", AGENT_RUN_PROMPT_CAP), ("answer", AGENT_RUN_ANSWER_CAP)): + data[key], clipped = _clip_activity(str(run.get(key) or ""), cap) + truncated |= clipped + tools: list[dict[str, Any]] = [] + raw_tools = run.get("tool_calls") + raw_tools = raw_tools if isinstance(raw_tools, list) else [] + for raw in raw_tools[:AGENT_RUN_MAX_TOOLS]: + if not isinstance(raw, dict): + continue + entry: dict[str, Any] = { + "name": str(raw.get("name") or ""), + "is_error": bool(raw.get("is_error")), + } + for key, cap in (("args", AGENT_RUN_ARGS_CAP), ("result", AGENT_RUN_RESULT_CAP)): + entry[key], clipped = _clip_activity(str(raw.get(key) or ""), cap) + truncated |= clipped + tools.append(entry) + truncated |= len(raw_tools) > AGENT_RUN_MAX_TOOLS + data["tool_calls"] = tools + for key in ("turns", "input_tokens", "output_tokens"): + data[key] = int(run.get(key) or 0) + data["cost_usd"] = float(run.get("cost_usd") or 0.0) + data["duration_s"] = float(run.get("duration_s") or 0.0) + data["truncated"] = truncated + return data + class SessionNotFoundError(KeyError): """No session matched the reference.""" @@ -254,6 +305,10 @@ def append_tombstone(self, session: Session, up_to_seq: int) -> TombstoneRecord: self._append(session, record) return record + def record_agent_run(self, session: Session, run: dict[str, Any]) -> EventRecord: + """Append one bounded agent-run activity record (kind ``agent_run``).""" + return self.append_event(session, "agent_run", _bounded_agent_run(run)) + # -- reading ------------------------------------------------------------ def _read_records_at(self, path: Path) -> list[Record]: @@ -385,6 +440,18 @@ def load_messages(self, session: Session) -> list[MessageRecord]: if isinstance(r, MessageRecord) and not self._is_hidden(r.seq, tombstones) ] + def load_agent_runs(self, session: Session) -> list[dict[str, Any]]: + """Agent-run activity records with tombstones applied, in append order.""" + records = self.read_records(session) + tombstones = self._active_tombstones(records) + return [ + dict(r.data) + for r in records + if isinstance(r, EventRecord) + and r.kind == "agent_run" + and not self._is_hidden(r.seq, tombstones) + ] + def undo(self, session: Session) -> TombstoneRecord | None: """Hide the last user turn (the user message and everything after it).""" visible = self.load_messages(session) diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index 712f5b2..527aee1 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -65,7 +65,12 @@ from lecode.extras.mcp_client import MCP_EXTRA, attach_mcp from lecode.extras.proc import run_proc from lecode.extras.status_signals import START, STOP, StatusEmitter -from lecode.extras.subagents import SubagentError, SubagentOutcome, run_subagent +from lecode.extras.subagents import ( + SubagentError, + SubagentOutcome, + SubagentProgress, + run_subagent, +) from lecode.hooks import ( INTERRUPT, NOTIFICATION, @@ -1525,8 +1530,11 @@ async def _run_turn(self, text: MessageContent, *, overlay: str | None = None) - self._turn_task = asyncio.ensure_future(self._run_turn(follow_up)) async def _run_subagent_turn(self, name: str, prompt: str) -> None: - """A direct ``@agent`` submission: the subagent answers as a side - query — the exchange is not persisted to the session.""" + """A direct ``@agent`` submission: the subagent answers as a side query. + + The exchange never enters the message history the model replays, but + the run's activity record is persisted for resume/inspection. + """ self._status.state = StatusLineState.RUNNING self._activity(f"@{name} working") outcome: SubagentOutcome | None = None @@ -1705,12 +1713,14 @@ def _on_event(self, event: Any) -> None: # Rendered after the stats line in _run_turn, not mid-stream. self._pending_review = event - def _on_child_event(self, agent: str, event: Any) -> None: - """Subagent runner event → feed rendering, prefixed with the agent. + def _on_child_event(self, progress: SubagentProgress) -> None: + """Subagent progress → feed rendering, prefixed with the agent. Token/Reasoning are skipped: they would interleave with the parent's own stream. """ + agent = progress.agent + event = progress.event if isinstance(event, ToolCall): self._feed.tool_call(f"{agent}/{event.name}", " ".join(event.arguments.split())) self._activity(f"@{agent} running {event.name}") diff --git a/tests/test_session_storage.py b/tests/test_session_storage.py index 211a8d3..659358c 100644 --- a/tests/test_session_storage.py +++ b/tests/test_session_storage.py @@ -258,6 +258,81 @@ def test_import_rejects_non_session_file(store, tmp_path): store.import_session(bad) +# -- agent-run activity --------------------------------------------------------- + + +def test_agent_run_round_trip(store, session): + store.record_agent_run( + session, + { + "run_id": "ab12cd34", + "agent": "explore", + "description": "Explore src", + "prompt": "scan src", + "status": "ok", + "answer": "found 3 files", + "turns": 2, + "input_tokens": 10, + "output_tokens": 4, + "cost_usd": 0.001, + "duration_s": 1.5, + "tool_calls": [ + {"name": "read", "args": '{"path": "x"}', "result": "contents", "is_error": False} + ], + }, + ) + runs = store.load_agent_runs(session) + assert len(runs) == 1 + run = runs[0] + assert run["run_id"] == "ab12cd34" + assert run["agent"] == "explore" + assert run["description"] == "Explore src" + assert run["prompt"] == "scan src" + assert run["status"] == "ok" + assert run["answer"] == "found 3 files" + assert run["turns"] == 2 + assert run["tool_calls"] == [ + {"name": "read", "args": '{"path": "x"}', "result": "contents", "is_error": False} + ] + + +def test_agent_runs_keep_append_order(store, session): + store.record_agent_run(session, {"run_id": "one", "agent": "explore", "status": "ok"}) + store.record_agent_run(session, {"run_id": "two", "agent": "explore", "status": "error"}) + assert [run["run_id"] for run in store.load_agent_runs(session)] == ["one", "two"] + + +def test_agent_run_bounds_large_payloads(store, session): + store.record_agent_run( + session, + { + "run_id": "big", + "agent": "explore", + "status": "ok", + "prompt": "p" * 5000, + "answer": "a" * 40000, + "tool_calls": [ + {"name": "read", "args": "g" * 5000, "result": "r" * 10000, "is_error": False} + ], + }, + ) + run = store.load_agent_runs(session)[0] + assert run["truncated"] is True + assert run["answer"].endswith("… (truncated)") + assert len(run["answer"]) <= 32 * 1024 + len("\n… (truncated)") + assert len(run["prompt"]) <= 2000 + len("\n… (truncated)") + tool = run["tool_calls"][0] + assert tool["result"].endswith("… (truncated)") + assert len(tool["result"]) <= 2000 + len("\n… (truncated)") + assert len(tool["args"]) <= 500 + len("\n… (truncated)") + + +def test_agent_runs_respect_tombstones(store, session): + record = store.record_agent_run(session, {"run_id": "gone", "agent": "explore", "status": "ok"}) + store.append_tombstone(session, up_to_seq=record.seq - 1) + assert store.load_agent_runs(session) == [] + + # -- attach locking ------------------------------------------------------------- diff --git a/tests/test_subagents.py b/tests/test_subagents.py index 3ba872d..26fbc3e 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -12,6 +12,7 @@ from lecode.agent.builder import build_runtime from lecode.agent.runner import AgentRunner +from lecode.agent.runner import Done as RunDone from lecode.agent.tools import ToolRegistry from lecode.agent.tools.task import make_tool from lecode.config.models import Config @@ -22,6 +23,7 @@ run_subagent, ) from lecode.providers.types import Done, TokenDelta, ToolCallDelta +from lecode.session import SessionStore def make_runtime(tmp_path, monkeypatch, provider, config=None): @@ -133,6 +135,126 @@ async def test_child_uses_agent_prompt_and_lean_registry(tmp_path, monkeypatch): assert {"read", "grep", "list_dir"} <= tool_names +async def test_subagent_progress_carries_run_identity(tmp_path, monkeypatch): + """Every forwarded child event is tagged with one run's id, agent, description.""" + provider = FakeProvider([{"text": "done"}]) + runtime = make_runtime(tmp_path, monkeypatch, provider) + seen: list[Any] = [] + outcome = await run_subagent( + runtime.ctx, + runtime.registry, + runtime.agents, + name="explore", + prompt="scan the repo", + description="Explore src", + on_event=seen.append, + ) + assert outcome.text == "done" + assert seen + run_ids = {progress.run_id for progress in seen} + assert len(run_ids) == 1 + assert next(iter(run_ids)) + assert {progress.agent for progress in seen} == {"explore"} + assert {progress.description for progress in seen} == {"Explore src"} + assert any(isinstance(progress.event, RunDone) for progress in seen) + + +async def test_completed_run_persists_activity_trail(tmp_path, monkeypatch): + """A finished child run lands in the session as one bounded record.""" + provider = FakeProvider( + [ + {"tool_calls": [{"name": "read", "arguments": '{"path": "note.txt"}'}]}, + {"text": "read it"}, + ] + ) + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("agent-run", tmp_path) + runtime = build_runtime(Config(), tmp_path, session=session, store=store) + runtime.ctx.extras["provider"] = provider + (tmp_path / "note.txt").write_text("hello trail", encoding="utf-8") + + outcome = await run_subagent( + runtime.ctx, + runtime.registry, + runtime.agents, + name="explore", + prompt="read note.txt", + description="Read note", + ) + + assert outcome.run_id + runs = store.load_agent_runs(session) + assert len(runs) == 1 + run = runs[0] + assert run["run_id"] == outcome.run_id + assert run["agent"] == "explore" + assert run["description"] == "Read note" + assert run["status"] == "ok" + assert run["answer"] == "read it" + assert run["turns"] == 2 + assert run["tool_calls"][0]["name"] == "read" + assert "hello trail" in run["tool_calls"][0]["result"] + assert run["tool_calls"][0]["is_error"] is False + + +async def test_cancelled_run_persists_cancelled_status(tmp_path, monkeypatch): + provider = BlockingChildProvider({"text": "unused"}) + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("agent-run", tmp_path) + runtime = build_runtime(Config(), tmp_path, session=session, store=store) + runtime.ctx.extras["provider"] = provider + + task = asyncio.ensure_future( + run_subagent( + runtime.ctx, + runtime.registry, + runtime.agents, + name="explore", + prompt="hang", + description="Cancelled run", + ) + ) + await asyncio.wait_for(provider.child_started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + + runs = store.load_agent_runs(session) + assert len(runs) == 1 + assert runs[0]["status"] == "cancelled" + assert runs[0]["description"] == "Cancelled run" + + +async def test_failed_run_persists_error_status(tmp_path, monkeypatch): + monkeypatch.setattr(subagents, "SUBAGENT_TIMEOUT_S", 0.05) + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("agent-run", tmp_path) + runtime = build_runtime(Config(), tmp_path, session=session, store=store) + runtime.ctx.extras["provider"] = NeverProvider() + + with pytest.raises(SubagentError, match="timed out"): + await run_subagent( + runtime.ctx, + runtime.registry, + runtime.agents, + name="explore", + prompt="hang", + description="Never finishes", + ) + + runs = store.load_agent_runs(session) + assert len(runs) == 1 + assert runs[0]["status"] == "error" + assert runs[0]["description"] == "Never finishes" + assert "timed out" in runs[0]["answer"] + + async def test_child_ctx_drops_question_callback(tmp_path, monkeypatch): """The child ctx (dataclasses.replace of the parent's) must not inherit the interactive question callback even if the parent has one installed.""" @@ -417,8 +539,14 @@ async def test_at_agent_runs_directly(tmp_path, monkeypatch): request = provider.requests[0] assert "read-only exploration agent" in request["messages"][0]["content"] assert request["messages"][1] == {"role": "user", "content": "count the files"} - # A side query: nothing persisted to the session. + # A side query: no message history, but the run itself is persisted. assert app.store.load_messages(app.session) == [] + runs = app.store.load_agent_runs(app.session) + assert len(runs) == 1 + assert runs[0]["agent"] == "explore" + assert runs[0]["status"] == "ok" + assert runs[0]["prompt"] == "count the files" + assert runs[0]["answer"] == "42 files" assert app._last_response == "42 files" assert app._status.input_tokens == 7 From 49036f05bc081e818b0190141b462850363acfc5 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Mon, 14 Sep 2026 18:01:29 +0400 Subject: [PATCH 2/8] feat: live agent roster, /runs detail panel, concise tool lines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TUI now keeps child-run activity out of the scrollback and in a compact roster above the composer: per-run number, agent, task description, status and current action, four rows plus +N more. /runs (picker or number/id) opens a live detail panel with the tool trail and answer; Escape closes it, draft and approvals untouched. A finished run leaves one attributed summary line. Tool output is now scannable: successful results render as ✔ name · first line (+N lines), failures keep a labelled head of up to ten lines, and every action line keeps the live ctx/ segment. Completed streamed answers are rendered as Markdown; the task tool reports its run id so the roster pairs a finished run with its answer. --- src/lecode/agent/runner.py | 5 +- src/lecode/agent/tools/task.py | 1 + src/lecode/slash/catalog.py | 1 + src/lecode/slash/handlers.py | 34 ++++- src/lecode/tui/agents.py | 267 +++++++++++++++++++++++++++++++++ src/lecode/tui/app.py | 136 ++++++++++++----- src/lecode/tui/feed.py | 84 +++++++++-- tests/test_subagents.py | 161 +++++++++++++++++++- tests/test_tui_agents.py | 104 +++++++++++++ tests/test_tui_app.py | 16 +- tests/test_tui_feed.py | 52 ++++--- tests/test_tui_question.py | 9 +- 12 files changed, 787 insertions(+), 83 deletions(-) create mode 100644 src/lecode/tui/agents.py create mode 100644 tests/test_tui_agents.py diff --git a/src/lecode/agent/runner.py b/src/lecode/agent/runner.py index 90f8d8b..066d848 100644 --- a/src/lecode/agent/runner.py +++ b/src/lecode/agent/runner.py @@ -26,7 +26,7 @@ import inspect import time from collections.abc import AsyncIterator, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from lecode.agent.review import review, user_request @@ -78,6 +78,8 @@ class ToolResult: name: str content: str is_error: bool + #: Tool-attached metadata (e.g. the task tool's run id for roster lookup). + metadata: dict[str, Any] = field(default_factory=dict) @dataclass(frozen=True) @@ -558,6 +560,7 @@ async def _run_tools( name=call["function"]["name"], content=result.content, is_error=result.is_error, + metadata=result.metadata, ), ) return messages diff --git a/src/lecode/agent/tools/task.py b/src/lecode/agent/tools/task.py index 294818d..23af4f6 100644 --- a/src/lecode/agent/tools/task.py +++ b/src/lecode/agent/tools/task.py @@ -93,6 +93,7 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: outcome.text or "(subagent returned no text)", metadata={ "agent": outcome.agent, + "run_id": outcome.run_id, "turns": outcome.turns, "input_tokens": outcome.input_tokens, "output_tokens": outcome.output_tokens, diff --git a/src/lecode/slash/catalog.py b/src/lecode/slash/catalog.py index f46e8d8..adfe285 100644 --- a/src/lecode/slash/catalog.py +++ b/src/lecode/slash/catalog.py @@ -44,6 +44,7 @@ ("btw", "Side note to the model mid-turn"), ("queue", "Show queued/steered messages"), ("tasks", "List background tasks"), + ("runs", "Inspect agent runs"), ("copy", "Copy the last answer to the clipboard"), ("export", "Export the session as HTML"), ("import", "Import a session file"), diff --git a/src/lecode/slash/handlers.py b/src/lecode/slash/handlers.py index 38b4b73..75579ec 100644 --- a/src/lecode/slash/handlers.py +++ b/src/lecode/slash/handlers.py @@ -739,6 +739,26 @@ async def cmd_queue(app: TuiApp, args: list[str]) -> None: app.feed.info("\n".join(lines)) +async def cmd_runs(app: TuiApp, args: list[str]) -> None: + """``/runs [number|id]``: list this session's agent runs, or open one in + the live detail panel (Escape closes it).""" + if not args: + runs = app.roster.runs() + if not runs: + app.feed.info("(no agent runs this session)") + return + lines = [f"{run.index}. {run.agent} · {run.description} · {run.status}" for run in runs] + lines.append("") + lines.append("open: /runs · close: Esc") + app.feed.info("\n".join(lines)) + return + run = app.roster.resolve(args[0]) + if run is None: + app.feed.error(f"no such agent run: {args[0]}") + return + app.open_agent_run(run.run_id) + + async def cmd_tasks(app: TuiApp, args: list[str]) -> None: """``/tasks``: background tasks (id, kind, status, age, description).""" manager = app.runtime.ctx.extras.get(BACKGROUND_EXTRA) @@ -1389,6 +1409,15 @@ def _complete_resume(app: TuiApp, args: list[str]) -> list[CompletionRow]: return rows +def _complete_runs(app: TuiApp, args: list[str]) -> list[CompletionRow]: + if args: + return [] # one reference only + return [ + (run.run_id, f"{run.index} {run.agent} · {run.description}", run.status) + for run in app.roster.runs() + ] + + def _complete_rewind(app: TuiApp, args: list[str]) -> list[CompletionRow]: if args: return [] @@ -1450,6 +1479,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "tutor": (_complete_tutor, None), "memory": (_complete_memory, None), "wt-exit": (_complete_wt_exit, None), + "runs": (_complete_runs, "no agent runs this session"), } @@ -1490,7 +1520,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: ), ("Permissions", ["permissions", "mode", "toggle"]), ("Worktrees", ["worktree", "wt-merge", "wt-exit"]), - ("Power features", ["loop", "chain", "mcp", "review", "tasks"]), + ("Power features", ["loop", "chain", "mcp", "review", "tasks", "runs"]), ( "Interface", [ @@ -1550,6 +1580,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "hooks": cmd_hooks, "agents": cmd_agents, "queue": cmd_queue, + "runs": cmd_runs, "tasks": cmd_tasks, "btw": cmd_btw, "copy": cmd_copy, @@ -1606,6 +1637,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "tutor": "", "review": "[file…]", "notifications": "[on|off]", + "runs": "[number|id]", } diff --git a/src/lecode/tui/agents.py b/src/lecode/tui/agents.py new file mode 100644 index 0000000..f72eb7b --- /dev/null +++ b/src/lecode/tui/agents.py @@ -0,0 +1,267 @@ +"""Live agent-run roster: identity, status, activity, detail rendering. + +Display state only — nothing here enters the model context. The roster +consumes :class:`SubagentProgress` events and answers three questions: what +is running now, what each run did, and what the open detail panel shows. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass, field + +from rich.text import Text + +from lecode.agent.runner import Done, Error, ToolCall, ToolResult +from lecode.extras.subagents import SubagentProgress +from lecode.permission.patterns import target_of +from lecode.tui.themes import Theme + +#: Rows the compact roster shows before collapsing into ``+N more``. +ROSTER_VISIBLE_ROWS = 4 + +#: Activity entries retained per run (the tail wins). +RUN_ACTIVITY_MAX = 100 + +#: Result lines shown per tool in the detail panel. +RESULT_PREVIEW_LINES = 3 + +#: Height cap of the in-layout detail panel (rows). +DETAIL_MAX_ROWS = 12 + +_STATUS_GLYPHS = { + "running": ("●", "accent"), + "ok": ("✔", "success"), + "error": ("✗", "error"), + "cancelled": ("—", "muted"), +} + + +def _clip(text: str, width: int) -> str: + flat = " ".join(text.split()) + return flat if len(flat) <= width else flat[: width - 1] + "…" + + +def _target(name: str, args: str) -> str: + """Best-effort target (path/command) from a tool call's raw arguments.""" + try: + parsed = json.loads(args) if args.strip() else {} + except json.JSONDecodeError: + parsed = {} + if isinstance(parsed, dict) and parsed: + return target_of(name, parsed) + return " ".join(args.split()) + + +@dataclass +class ActivityEntry: + """One tool call inside a child run, paired with its result when it lands.""" + + call_id: str + name: str + args: str + result: str = "" + is_error: bool = False + running: bool = True + + @property + def target(self) -> str: + return _clip(_target(self.name, self.args), 60) + + +@dataclass +class AgentRun: + """The roster's per-run state.""" + + run_id: str + index: int + agent: str + description: str + status: str = "running" + activity: list[ActivityEntry] = field(default_factory=list) + answer: str = "" + error: str = "" + started_at: float = field(default_factory=time.monotonic) + truncated: bool = False + + @property + def current(self) -> str: + """The one-line current action (what the roster row shows).""" + if self.status != "running": + return {"ok": "done", "error": "failed", "cancelled": "cancelled"}.get( + self.status, self.status + ) + active = next((entry for entry in reversed(self.activity) if entry.running), None) + if active is None: + return "thinking" + return f"{active.name} {active.target}".strip() + + +class AgentRoster: + """Append-only run registry keyed by run id; finished runs stay listed.""" + + def __init__(self) -> None: + self._runs: dict[str, AgentRun] = {} + self._order: list[str] = [] + + def observe(self, progress: SubagentProgress) -> AgentRun: + """Fold one child progress event into its run's state.""" + run = self._runs.get(progress.run_id) + if run is None: + run = AgentRun( + run_id=progress.run_id, + index=len(self._order) + 1, + agent=progress.agent, + description=progress.description, + ) + self._runs[run.run_id] = run + self._order.append(run.run_id) + event = progress.event + if isinstance(event, ToolCall): + run.activity.append( + ActivityEntry(call_id=event.id, name=event.name, args=event.arguments) + ) + if len(run.activity) > RUN_ACTIVITY_MAX: + run.activity = run.activity[-RUN_ACTIVITY_MAX:] + run.truncated = True + elif isinstance(event, ToolResult): + entry = next((e for e in run.activity if e.call_id == event.id), None) + if entry is not None: + entry.result = event.content + entry.is_error = event.is_error + entry.running = False + elif isinstance(event, Error): + run.status = "error" + run.error = event.message + elif isinstance(event, Done) and run.status == "running": + run.status = "ok" + return run + + def finish( + self, run_id: str, *, answer: str = "", error: str = "", is_error: bool = False + ) -> AgentRun | None: + """Close a run with its answer/error once the caller knows the outcome.""" + run = self._runs.get(run_id) + if run is None: + return None + if answer: + run.answer = answer + if error: + run.error = error + if is_error or error: + run.status = "error" + elif run.status == "running": + run.status = "ok" + return run + + def cancel_running(self) -> list[AgentRun]: + """Mark every still-running run cancelled (parent turn was cancelled).""" + cancelled = [run for run in self._runs.values() if run.status == "running"] + for run in cancelled: + run.status = "cancelled" + return cancelled + + def has_running(self) -> bool: + return any(run.status == "running" for run in self._runs.values()) + + def get(self, run_id: str) -> AgentRun | None: + return self._runs.get(run_id) + + def resolve(self, ref: str) -> AgentRun | None: + """Resolve a run by roster number or run-id prefix (exact first).""" + if ref.isdigit(): + wanted = int(ref) + return next((run for run in self.runs() if run.index == wanted), None) + exact = self._runs.get(ref) + if exact is not None: + return exact + matches = [run for run in self.runs() if run.run_id.startswith(ref)] + return matches[0] if len(matches) == 1 else None + + def runs(self) -> list[AgentRun]: + """Runs in start order (stable numbering).""" + return [self._runs[run_id] for run_id in self._order] + + def visible(self, limit: int = ROSTER_VISIBLE_ROWS) -> tuple[list[AgentRun], int]: + """``(shown, hidden_count)``: running first (newest first), then + finished (newest first); stable indices come from the run itself.""" + runs = self.runs() + running = [run for run in reversed(runs) if run.status == "running"] + finished = [run for run in reversed(runs) if run.status != "running"] + ordered = running + finished + return ordered[:limit], max(0, len(ordered) - limit) + + +def _glyph(status: str, theme: Theme) -> tuple[str, str]: + glyph, slot = _STATUS_GLYPHS.get(status, _STATUS_GLYPHS["running"]) + return glyph, getattr(theme, slot) + + +def roster_lines(roster: AgentRoster, theme: Theme, width: int) -> list[Text]: + """The compact roster block (header + up to four rows + overflow).""" + runs = roster.runs() + if not runs: + return [] + running = sum(1 for run in runs if run.status == "running") + done = len(runs) - running + header = f"agents · {running} running" + if done: + header += f" · {done} done" + lines = [Text(header, style=theme.muted)] + shown, hidden = roster.visible() + for run in shown: + glyph, style = _glyph(run.status, theme) + line = Text() + line.append(f" {glyph} ", style=style) + line.append(f"{run.index} {run.agent} ", style=theme.accent) + line.append(_clip(run.description, max(12, width // 3)), style=theme.text) + line.append(f" · {_clip(run.current, max(12, width // 3))}", style=theme.muted) + lines.append(line) + if hidden: + lines.append(Text(f" … +{hidden} more · /runs", style=theme.muted)) + return lines + + +def _preview(text: str, width: int, limit: int = RESULT_PREVIEW_LINES) -> list[str]: + rows = text.splitlines() + shown = [_clip(row, width) for row in rows[:limit]] + if len(rows) > limit: + shown.append(f"… ({len(rows) - limit} more lines)") + return shown + + +def detail_lines(run: AgentRun | None, theme: Theme, width: int) -> list[Text]: + """The per-run detail panel: identity, tool trail, answer/error.""" + if run is None: + return [] + glyph, style = _glyph(run.status, theme) + lines: list[Text] = [] + header = Text() + header.append(f" {glyph} ", style=style) + header.append(f"{run.index} {run.agent} ", style=theme.accent) + header.append(_clip(run.description, max(12, width // 2)), style=theme.text) + header.append(f" · {run.current}", style=theme.muted) + lines.append(header) + if not run.activity: + lines.append(Text(" (no tool calls yet)", style=theme.muted)) + for entry in run.activity: + row = Text() + row.append(f" ⚙ {entry.name} ", style=theme.tool) + row.append(_clip(entry.target, max(12, width - 20)), style=theme.text) + if entry.running: + row.append(" · running", style=theme.muted) + elif entry.is_error: + row.append(" · failed", style=theme.error) + lines.append(row) + if entry.result: + preview_style = theme.error if entry.is_error else theme.muted + for row_text in _preview(entry.result, max(12, width - 6)): + lines.append(Text(f" {row_text}", style=preview_style)) + if run.error and not run.answer: + lines.append(Text(f" error: {_clip(run.error, max(12, width - 14))}", style=theme.error)) + if run.answer: + lines.append(Text(" answer:", style=theme.muted)) + for row_text in _preview(run.answer, max(12, width - 6)): + lines.append(Text(f" {row_text}", style=theme.text)) + return lines diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index 527aee1..1379eb7 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -39,6 +39,7 @@ from prompt_toolkit.utils import get_cwidth from prompt_toolkit.widgets import Frame, TextArea from rich.console import Console +from rich.text import Text from lecode.agent.runner import ( AgentRunner, @@ -109,6 +110,12 @@ SlashCommand, UnknownCommandError, ) +from lecode.tui.agents import ( + DETAIL_MAX_ROWS, + AgentRoster, + detail_lines, + roster_lines, +) from lecode.tui.clipboard import copy_to_clipboard from lecode.tui.feed import Feed from lecode.tui.input import ( @@ -255,6 +262,10 @@ def __init__( # inline through the feed; installed here so tests driving _submit # directly get it too. self._runtime.ctx.extras["subagent_events"] = self._on_child_event + #: Live agent-run roster (compact above the composer; /runs opens detail). + self._roster = AgentRoster() + self._detail_run_id: str | None = None + self._roster_window: Any | None = None # Background-task completions surface as feed info lines as they land. background = self._runtime.ctx.extras.get(BACKGROUND_EXTRA) if background is not None: @@ -639,14 +650,34 @@ def _question_escape(event: Any) -> None: event.current_buffer.reset() self._invalidate() + detail_open = Condition(lambda: self._detail_run_id is not None) + + @kb.add("pageup", filter=detail_open) + def _detail_scroll_up(event: Any) -> None: + if self._roster_window is not None: + self._roster_window.vertical_scroll = max( + 0, self._roster_window.vertical_scroll - 5 + ) + self._invalidate() + + @kb.add("pagedown", filter=detail_open) + def _detail_scroll_down(event: Any) -> None: + if self._roster_window is not None: + self._roster_window.vertical_scroll += 5 + self._invalidate() + @kb.add("escape", filter=~approval_pending & ~question_pending) def _close_completion_menu(event: Any) -> None: - # Approval and question prompts own Escape while pending; otherwise - # dismiss the dropdown, restoring typed text a navigation - # overwrote. The no-match rows (slash and argument pickers) have - # no completion state, so their dismissal is remembered per typed - # text (any edit or Tab brings it back). Longer M-* sequences - # still win over this bare-key handler. + # Approval and question prompts own Escape while pending; an open + # detail panel closes before anything else. Then dismiss the + # dropdown, restoring typed text a navigation overwrote. The + # no-match rows (slash and argument pickers) have no completion + # state, so their dismissal is remembered per typed text (any edit + # or Tab brings it back). Longer M-* sequences still win over this + # bare-key handler. + if self._detail_run_id is not None: + self.close_agent_run() + return buffer = event.current_buffer if buffer.complete_state is not None: buffer.cancel_completion() @@ -908,6 +939,15 @@ def _build_app(self, input: Input | None = None, output: Output | None = None) - height=Dimension(min=1, max=10), ) live_area = ConditionalContainer(live_window, Condition(lambda: bool(self._live_text))) + # Agent roster / detail panel: compact rows above the composer while + # runs are live (or one run is open for inspection). + self._roster_window = Window( + FormattedTextControl(self._roster_text), + wrap_lines=False, + dont_extend_height=True, + height=Dimension(min=1, max=DETAIL_MAX_ROWS), + ) + roster_area = ConditionalContainer(self._roster_window, Condition(self._roster_visible)) # The chatbox: a framed input area directly above the statusline. # Enter submits the text into the transcript above (see _enter). self._chatbox = Frame(self._input_area, title="message") @@ -1028,7 +1068,7 @@ def menu_footer() -> str: picker_menu_visible, ) return Application( - layout=Layout(HSplit([live_area, self._chatbox, picker_panel, toolbar])), + layout=Layout(HSplit([live_area, roster_area, self._chatbox, picker_panel, toolbar])), style=Style.from_dict( { "picker-menu": f"bg:{PICKER_MENU_BG} {self._theme.muted}", @@ -1487,6 +1527,8 @@ async def _run_turn(self, text: MessageContent, *, overlay: str | None = None) - except asyncio.CancelledError: cancelled = True self._feed.info("turn cancelled") + for run in self._roster.cancel_running(): + self._feed.agent_summary(run) except ProviderError: pass # already rendered via the Error event finally: @@ -1555,6 +1597,7 @@ async def _run_subagent_turn(self, name: str, prompt: str) -> None: self._status.state = StatusLineState.IDLE self._status.activity = None if outcome is not None: + self._roster.finish(outcome.run_id, answer=outcome.text) self._feed.assistant_text(outcome.text) self._last_response = outcome.text self._status.input_tokens += outcome.input_tokens @@ -1672,6 +1715,10 @@ def _on_event(self, event: Any) -> None: elif isinstance(event, ToolResult): self._status.context_used += self._estimate(event.content) self._feed.tool_result(event.name, event.content, event.is_error) + run_id = str(event.metadata.get("run_id") or "") if event.metadata else "" + if run_id: + # A finished task tool run: pair its answer with the roster entry. + self._roster.finish(run_id, answer=event.content, is_error=event.is_error) self._activity("thinking") elif isinstance(event, Error): self._feed.error(event.message) @@ -1714,35 +1761,56 @@ def _on_event(self, event: Any) -> None: self._pending_review = event def _on_child_event(self, progress: SubagentProgress) -> None: - """Subagent progress → feed rendering, prefixed with the agent. + """Subagent progress → the live roster; one summary line when it ends. - Token/Reasoning are skipped: they would interleave with the parent's - own stream. + Per-call child detail stays in the roster/detail panel instead of + flooding the scrollback; a finished run leaves one attributed summary + line. Token/Reasoning are ignored: they would interleave with the + parent's own stream. """ - agent = progress.agent - event = progress.event - if isinstance(event, ToolCall): - self._feed.tool_call(f"{agent}/{event.name}", " ".join(event.arguments.split())) - self._activity(f"@{agent} running {event.name}") - elif isinstance(event, ToolResult): - self._feed.tool_result(f"{agent}/{event.name}", event.content, event.is_error) - self._activity(f"@{agent} working") - elif isinstance(event, Error): - self._feed.error(f"{agent}: {event.message}") - elif isinstance(event, Retrying): - self._feed.retrying(event.attempt, event.delay) - elif isinstance(event, LlmCall): - self._feed.llm_call(f"{agent} · {event.model}", event.turn) - elif isinstance(event, LlmResponse): - self._feed.llm_response( - f"{agent} · {event.model}", - event.turn, - event.input_tokens, - event.output_tokens, - event.cost_usd, - ) - elif isinstance(event, Done): - self._feed.info(f"{agent} finished ({event.stop_reason}, {event.turns} turn(s))") + run = self._roster.observe(progress) + if isinstance(progress.event, (Error, Done)): + self._feed.agent_summary(run) + self._invalidate() + + # -- agent roster / detail panel ------------------------------------------- + + @property + def roster(self) -> AgentRoster: + return self._roster + + @property + def detail_run_id(self) -> str | None: + return self._detail_run_id + + def open_agent_run(self, run_id: str) -> bool: + """Show one run's live detail panel; ``False`` for an unknown run.""" + if self._roster.get(run_id) is None: + return False + self._detail_run_id = run_id + if self._roster_window is not None: + self._roster_window.vertical_scroll = 0 + self._invalidate() + return True + + def close_agent_run(self) -> None: + self._detail_run_id = None + self._invalidate() + + def _roster_visible(self) -> bool: + return self._detail_run_id is not None or self._roster.has_running() + + def _term_width(self) -> int: + if self._app is not None: + with contextlib.suppress(Exception): + return self._app.output.get_size().columns + return 80 + + def _roster_text(self) -> list[Text]: + width = self._term_width() + if self._detail_run_id is not None: + return detail_lines(self._roster.get(self._detail_run_id), self._theme, width) + return roster_lines(self._roster, self._theme, width) # -- agents / totals ---------------------------------------------------------- diff --git a/src/lecode/tui/feed.py b/src/lecode/tui/feed.py index 8dcf86b..f5945a2 100644 --- a/src/lecode/tui/feed.py +++ b/src/lecode/tui/feed.py @@ -15,16 +15,20 @@ can record with ``Console(record=True, file=StringIO())``. Logbook style: every discrete line carries a ``[HH:MM:SS]`` timestamp — except -the user-input echo, which prints verbatim — and action lines (tool calls, -tool results) end with the live ``ctx used/window · $cost-so-far`` segment -when a ``metrics`` callable is bound (the TUI binds it to the statusline -state). +the user-input echo, which prints verbatim. Tool lines stay scannable: a call +is one attributed line, a successful result is a one-line summary (``✔ name · +first line (+N lines)``), and only failures keep a head of raw output. Every +action line still ends with the live ``ctx used/window · $cost`` segment when +a ``metrics`` callable is bound — spend is always visible. A completed +streamed answer is rendered as Markdown when a stream sink is bound (the +TUI), raw otherwise. """ from __future__ import annotations from collections.abc import Callable from datetime import datetime +from typing import TYPE_CHECKING from rich.console import Console from rich.markdown import Markdown @@ -34,12 +38,18 @@ from lecode.tui.statusline import context_meter, format_cost, human_tokens from lecode.tui.themes import Theme +if TYPE_CHECKING: + from lecode.tui.agents import AgentRun + #: Max length of a rendered tool-call line. TOOL_CALL_MAX_LEN = 120 -#: Lines of a tool result shown before elision kicks in. +#: Lines of a failed tool result shown before elision kicks in. TOOL_RESULT_HEAD_LINES = 10 +#: Max characters of a successful tool result's one-line summary. +TOOL_RESULT_SUMMARY_MAX_LEN = 140 + #: ``(context_used, context_window, cost_usd)`` at render time. MetricsFn = Callable[[], tuple[int, int, float]] @@ -126,7 +136,8 @@ def _flush_stream(self) -> None: self._stream_parts = [] if self.stream_clear is not None: self.stream_clear() - self._console.print(full, markup=False, highlight=False, soft_wrap=True) + # Completed answers are formatted; live tokens stayed raw in the sink. + self._console.print(Markdown(full)) else: self._console.print() self._stream_printed = False @@ -183,18 +194,40 @@ def tool_call(self, name: str, args_preview: str) -> None: self._console.print(Text(line + self._suffix(), style=self._theme.tool)) def tool_result(self, name: str, content: str, is_error: bool = False) -> None: - """Render a tool result head with ``… (N more lines)`` elision. + """One attributed line: a result summary, or a head of the failure. - The timestamp/metrics marker goes on its own line after the output, - so multi-line output reads top-down and the bookkeeping lands last. + Success is a summary (first non-empty line, clipped, plus a count of + the hidden lines) — routine file contents never flood the transcript. + Failures stay prominent: the label plus up to ten raw lines. """ - lines = content.splitlines() - shown = lines[:TOOL_RESULT_HEAD_LINES] - if len(lines) > TOOL_RESULT_HEAD_LINES: - shown.append(f"… ({len(lines) - TOOL_RESULT_HEAD_LINES} more lines)") - shown.append(f"[{self._stamp()}]{self._suffix()}") - style = self._theme.error if is_error else self._theme.muted - self._console.print(Text("\n".join(shown), style=style)) + if is_error: + rows = content.splitlines() + shown = rows[:TOOL_RESULT_HEAD_LINES] + if len(rows) > TOOL_RESULT_HEAD_LINES: + shown.append(f"… ({len(rows) - TOOL_RESULT_HEAD_LINES} more lines)") + body = "\n".join(shown) or "(no output)" + self._console.print( + Text( + f"[{self._stamp()}] ✗ {name}{self._suffix()}\n{body}", + style=self._theme.error, + ) + ) + return + rows = [row for row in content.splitlines() if row.strip()] + if not rows: + summary = "(no output)" + else: + summary = " ".join(rows[0].split()) + if len(summary) > TOOL_RESULT_SUMMARY_MAX_LEN: + summary = summary[: TOOL_RESULT_SUMMARY_MAX_LEN - 1] + "…" + if len(rows) > 1: + summary += f" (+{len(rows) - 1} lines)" + line = Text() + line.append(f"[{self._stamp()}] ", style=self._theme.muted) + line.append("✔ ", style=self._theme.success) + line.append(f"{name} · {summary}", style=self._theme.text) + line.append(self._suffix(), style=self._theme.muted) + self._console.print(line) def turn_stats( self, @@ -229,6 +262,25 @@ def turn_stats( line = f"[{self._stamp()}] " + " · ".join(parts) self._console.print(Text(line, style=self._theme.muted)) + def agent_summary(self, run: AgentRun) -> None: + """One attributed line for a finished child run (the roster keeps detail).""" + glyph, slot = { + "ok": ("✔", "success"), + "error": ("✗", "error"), + "cancelled": ("—", "muted"), + }.get(run.status, ("✔", "muted")) + parts = [f"{run.agent} · {run.description}"] + if run.activity: + count = len(run.activity) + parts.append(f"{count} tool call{'s' if count != 1 else ''}") + if run.status == "error" and run.error: + parts.append(run.error.splitlines()[0][:80]) + self._console.print( + Text( + f"[{self._stamp()}] {glyph} " + " · ".join(parts), style=getattr(self._theme, slot) + ) + ) + def review(self, model: str, feedback: str) -> None: """Pierre-mode feedback: a labelled block after the stats line.""" self._console.print(Text(f"[{self._stamp()}] ◆ pierre ({model})", style=self._theme.accent)) diff --git a/tests/test_subagents.py b/tests/test_subagents.py index 26fbc3e..94186bd 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -11,7 +11,7 @@ from tests.test_tui_app import make_app, make_blocking_app, wait_for from lecode.agent.builder import build_runtime -from lecode.agent.runner import AgentRunner +from lecode.agent.runner import AgentRunner, ToolResult from lecode.agent.runner import Done as RunDone from lecode.agent.tools import ToolRegistry from lecode.agent.tools.task import make_tool @@ -511,9 +511,18 @@ async def test_cancelling_parent_cancels_child(tmp_path, monkeypatch): # -- TUI integration ------------------------------------------------------------------- -async def test_child_events_render_inline(tmp_path, monkeypatch): +async def test_child_events_feed_roster_and_one_summary_line(tmp_path, monkeypatch): script = [ - {"tool_calls": [{"name": "task", "arguments": '{"prompt": "scan", "agent": "explore"}'}]}, + { + "tool_calls": [ + { + "name": "task", + "arguments": ( + '{"prompt": "scan", "agent": "explore", "description": "Scan repo"}' + ), + } + ] + }, {"tool_calls": [{"name": "list_dir", "arguments": '{"path": "."}'}]}, # child {"text": "scan result"}, # child final {"text": "parent final"}, # parent turn 2 @@ -521,13 +530,153 @@ async def test_child_events_render_inline(tmp_path, monkeypatch): app, _, out = make_app(tmp_path, monkeypatch, script) await app._submit("please scan") await app._turn_task + + runs = app.roster.runs() + assert len(runs) == 1 + run = runs[0] + assert run.description == "Scan repo" + assert run.status == "ok" + assert run.answer == "scan result" + assert [entry.name for entry in run.activity] == ["list_dir"] + rendered = out.getvalue() - assert "⚙ task(" in rendered - assert "explore/list_dir" in rendered - assert "explore finished" in rendered + # Per-call child lines no longer flood the transcript; the run gets one + # attributed summary line and stays inspectable through the roster. + assert "explore/list_dir" not in rendered + assert "Scan repo" in rendered assert "parent final" in rendered +async def test_task_result_event_carries_run_id_metadata(tmp_path, monkeypatch): + provider = FakeProvider( + [ + { + "tool_calls": [ + {"name": "task", "arguments": '{"prompt": "scan", "agent": "explore"}'} + ] + }, + {"text": "scan result"}, + {"text": "parent final"}, + ] + ) + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("agent-run", tmp_path) + runtime = build_runtime(Config(), tmp_path, session=session, store=store) + runtime.ctx.extras["provider"] = provider + events: list[Any] = [] + runner = AgentRunner(provider, runtime.registry, runtime.ctx) + await runner.run([{"role": "user", "content": "go"}], on_event=events.append) + + task_results = [ + event + for event in events + if isinstance(event, ToolResult) and event.name == "task" and not event.is_error + ] + assert task_results + run_id = task_results[0].metadata.get("run_id") + assert run_id + assert run_id == store.load_agent_runs(session)[0]["run_id"] + + +async def test_roster_panel_visible_while_child_runs(tmp_path, monkeypatch): + provider = BlockingChildProvider( + { + "tool_calls": [ + { + "name": "task", + "arguments": ( + '{"prompt": "scan", "agent": "explore", "description": "Scan repo"}' + ), + } + ] + } + ) + app, _, _ = make_app(tmp_path, monkeypatch, []) + app._runner.provider = provider + app._runtime.ctx.extras["provider"] = provider + await app._submit("go") + await wait_for(lambda: provider.child_started.is_set()) + + assert app._roster_visible() + rows = "\n".join(line.plain for line in app._roster_text()) + assert "Scan repo" in rows + + run_id = app.roster.runs()[0].run_id + assert app.open_agent_run(run_id) + detail = "\n".join(line.plain for line in app._roster_text()) + assert "Scan repo" in detail + + app._turn_task.cancel() + await app._turn_task + assert app.roster.runs()[0].status == "cancelled" + app.close_agent_run() + assert not app._roster_visible() + + +async def test_runs_command_lists_and_opens_detail(tmp_path, monkeypatch): + script = [ + { + "tool_calls": [ + { + "name": "task", + "arguments": ( + '{"prompt": "scan", "agent": "explore", "description": "Scan repo"}' + ), + } + ] + }, + {"tool_calls": [{"name": "list_dir", "arguments": '{"path": "."}'}]}, + {"text": "scan result"}, + {"text": "parent final"}, + ] + app, _, out = make_app(tmp_path, monkeypatch, script) + await app._submit("please scan") + await app._turn_task + run_id = app.roster.runs()[0].run_id + + await app.handle_command("/runs") + assert "Scan repo" in out.getvalue() + + await app.handle_command(f"/runs {run_id}") + assert app.detail_run_id == run_id + app.close_agent_run() + assert app.detail_run_id is None + + +async def test_runs_picker_offers_run_ids(tmp_path, monkeypatch): + script = [ + {"tool_calls": [{"name": "task", "arguments": '{"prompt": "scan", "agent": "explore"}'}]}, + {"text": "scan result"}, + {"text": "parent final"}, + ] + app, _, _ = make_app(tmp_path, monkeypatch, script) + await app._submit("please scan") + await app._turn_task + run_id = app.roster.runs()[0].run_id + + from prompt_toolkit.document import Document + + result = app.arg_completion_rows(Document("/runs ", 6)) + assert result is not None + _, _, _, rows = result + assert any(insert == run_id for insert, _, _ in rows) + + +async def test_detail_panel_preserves_draft(tmp_path, monkeypatch): + from prompt_toolkit.input import create_pipe_input + from prompt_toolkit.output import DummyOutput + + app, _, _ = make_app(tmp_path, monkeypatch, []) + with create_pipe_input() as inp: + app._build_app(input=inp, output=DummyOutput()) + app._input_area.buffer.text = "half-written" + app.open_agent_run("r1") + app.close_agent_run() + assert app._input_area.buffer.text == "half-written" + + async def test_at_agent_runs_directly(tmp_path, monkeypatch): script = [{"text": "42 files", "usage": {"input_tokens": 7, "output_tokens": 3}}] app, provider, out = make_app(tmp_path, monkeypatch, script) diff --git a/tests/test_tui_agents.py b/tests/test_tui_agents.py new file mode 100644 index 0000000..0596110 --- /dev/null +++ b/tests/test_tui_agents.py @@ -0,0 +1,104 @@ +"""Tests for the live agent roster and its detail rendering.""" + +from __future__ import annotations + +from lecode.agent.runner import Done, Error, LlmCall, ToolCall, ToolResult +from lecode.extras.subagents import SubagentProgress +from lecode.tui.agents import ( + ROSTER_VISIBLE_ROWS, + AgentRoster, + detail_lines, + roster_lines, +) +from lecode.tui.themes import THEME + + +def progress(run_id: str, event, *, agent: str = "explore", description: str = "Explore src"): + return SubagentProgress(run_id=run_id, agent=agent, description=description, event=event) + + +def _plain(lines) -> str: + return "\n".join(line.plain for line in lines) + + +def test_roster_tracks_identity_and_tool_activity(): + roster = AgentRoster() + roster.observe(progress("r1", LlmCall(model="m", turn=1))) + roster.observe(progress("r1", ToolCall(id="c1", name="read", arguments='{"path": "a.py"}'))) + roster.observe( + progress("r1", ToolResult(id="c1", name="read", content="contents", is_error=False)) + ) + run = roster.get("r1") + assert run is not None + assert run.agent == "explore" + assert run.description == "Explore src" + assert run.status == "running" + assert run.activity[0].name == "read" + assert run.activity[0].result == "contents" + assert run.activity[0].running is False + + +def test_roster_terminal_events_set_status(): + roster = AgentRoster() + roster.observe(progress("ok", Done(stop_reason="done", turns=2))) + roster.observe(progress("bad", Error(message="boom"))) + assert roster.get("ok").status == "ok" + assert roster.get("bad").status == "error" + assert roster.get("bad").error == "boom" + + +def test_roster_visible_prefers_running_runs(): + roster = AgentRoster() + for index in range(5): + roster.observe( + progress(f"r{index}", LlmCall(model="m", turn=1), description=f"run {index}") + ) + roster.observe(progress("r4", Done(stop_reason="done", turns=1))) + visible, hidden = roster.visible(limit=4) + assert [run.run_id for run in visible] == ["r3", "r2", "r1", "r0"] + assert hidden == 1 + assert len(visible) <= ROSTER_VISIBLE_ROWS + + +def test_roster_cancel_running_keeps_finished_runs(): + roster = AgentRoster() + roster.observe(progress("r1", LlmCall(model="m", turn=1))) + roster.observe(progress("r2", LlmCall(model="m", turn=1))) + roster.observe(progress("r2", Done(stop_reason="done", turns=1))) + roster.cancel_running() + assert roster.get("r1").status == "cancelled" + assert roster.get("r2").status == "ok" + + +def test_roster_lines_show_description_state_and_overflow(): + roster = AgentRoster() + for index in range(6): + roster.observe( + progress(f"r{index}", LlmCall(model="m", turn=1), description=f"run {index}") + ) + text = _plain(roster_lines(roster, THEME, width=80)) + assert "run 5" in text + assert "+2 more" in text + assert "running" in text + + +def test_roster_lines_empty_when_no_runs(): + assert roster_lines(AgentRoster(), THEME, width=80) == [] + + +def test_detail_lines_show_tools_and_answer(): + roster = AgentRoster() + roster.observe(progress("r1", ToolCall(id="c1", name="read", arguments='{"path": "a.py"}'))) + roster.observe( + progress( + "r1", ToolResult(id="c1", name="read", content="line one\nline two", is_error=False) + ) + ) + roster.observe(progress("r1", Done(stop_reason="done", turns=1))) + roster.finish("r1", answer="the answer") + text = _plain(detail_lines(roster.get("r1"), THEME, width=80)) + assert "Explore src" in text + assert "read" in text + assert "a.py" in text + assert "line one" in text + assert "the answer" in text diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index 21f0a4a..e319aae 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -88,9 +88,10 @@ async def test_unknown_model_keeps_configured_window(tmp_path, monkeypatch): def test_layout_is_chatbox_above_statusline(tmp_path, monkeypatch): """The input is a framed chatbox directly above the 3-line statusline; - the frame's bottom border is the split between them. A conditional live - region for streamed text sits above the chatbox, and the themed picker - panel anchors below the input for all trigger menus.""" + the frame's bottom border is the split between them. Conditional regions + for streamed text and for the live agent roster/detail panel sit above + the chatbox, and the themed picker panel anchors below the input for all + trigger menus.""" from prompt_toolkit.layout.containers import ConditionalContainer, Window, to_container from prompt_toolkit.widgets import Frame @@ -99,11 +100,12 @@ def test_layout_is_chatbox_above_statusline(tmp_path, monkeypatch): pt_app = app._build_app(input=inp, output=DummyOutput()) assert isinstance(app._chatbox, Frame) and app._chatbox.body is app._input_area children = pt_app.layout.container.children - assert len(children) == 4 + assert len(children) == 5 assert isinstance(children[0], ConditionalContainer) # live stream region - assert children[1] is to_container(app._chatbox) # Frame unwraps to its HSplit - assert isinstance(children[2], ConditionalContainer) # picker panel sizes to its rows - assert isinstance(children[3], Window) and children[3].height == 3 + assert isinstance(children[1], ConditionalContainer) # agent roster/detail panel + assert children[2] is to_container(app._chatbox) # Frame unwraps to its HSplit + assert isinstance(children[3], ConditionalContainer) # picker panel sizes to its rows + assert isinstance(children[4], Window) and children[4].height == 3 assert app._live_buffer is not None diff --git a/tests/test_tui_feed.py b/tests/test_tui_feed.py index e1114c0..a8beac5 100644 --- a/tests/test_tui_feed.py +++ b/tests/test_tui_feed.py @@ -55,14 +55,14 @@ def test_lines_carry_timestamps(theme): ) lines = out.getvalue().splitlines() stamped = [ln for ln in lines if ln.startswith("[")] - # tool call, tool result marker, info, error, turn stats (not the user echo) + # tool call, tool result, info, error, turn stats (not the verbatim echo) assert len(stamped) == 5 assert all(len(ln) >= 10 and ln[1:3].isdigit() and ln[3] == ":" for ln in stamped) - # the tool result marker is its own line after the output - assert lines[lines.index("ok") + 1].startswith("[") + assert not lines[0].startswith("[") def test_metrics_suffix_on_action_lines(theme): + """Live ctx/cost stays on every action line: spend is always visible.""" feed, out = make_feed(theme) feed.metrics = lambda: (12_300, 200_000, 0.0412) feed.user_message("hi") # verbatim: no metrics suffix on the input echo @@ -72,12 +72,6 @@ def test_metrics_suffix_on_action_lines(theme): assert rendered.count("ctx 12.3k/200.0k · $0.0412") == 2 -def test_no_metrics_suffix_when_unbound(theme): - feed, out = make_feed(theme) - feed.user_message("hi") - assert "ctx" not in out.getvalue() - - def test_llm_call_logged(theme): feed, out = make_feed(theme) feed.metrics = lambda: (12_300, 200_000, 0.0412) @@ -182,23 +176,32 @@ def test_tool_call_truncates_long_args(theme): assert len(line) == TOOL_CALL_MAX_LEN -def test_tool_result_elides_long_content(theme): +def test_tool_result_summarizes_long_content(theme): feed, out = make_feed(theme) content = "\n".join(f"line {i}" for i in range(25)) feed.tool_result("bash", content) rendered = out.getvalue() + assert "✔ bash" in rendered assert "line 0" in rendered - assert "line 9" in rendered - assert "line 10" not in rendered - assert "… (15 more lines)" in rendered + assert "line 1" not in rendered # only the first line survives + assert "+24 lines" in rendered -def test_tool_result_short_content_not_elided(theme): +def test_tool_result_single_line_summary(theme): feed, out = make_feed(theme) feed.tool_result("bash", "a\nb\nc") rendered = out.getvalue() - assert "a\nb\nc" in rendered - assert "more lines" not in rendered + assert "✔ bash · a (+2 lines)" in rendered + + +def test_tool_result_error_keeps_head(theme): + feed, out = make_feed(theme) + feed.tool_result("bash", "\n".join(f"e{i}" for i in range(15)), is_error=True) + rendered = out.getvalue() + assert "✗ bash" in rendered + assert "e9" in rendered + assert "e10" not in rendered + assert "… (5 more lines)" in rendered def test_tool_result_error_styled(theme, monkeypatch): @@ -300,10 +303,25 @@ def test_sink_stream_closed_by_llm_response(theme): feed.llm_response("m", 1, 10, 5, 0.001) feed.stream_end() lines = out.getvalue().splitlines() - assert lines[0] == "the answer" + # Markdown pads rendered lines to the console width. + assert lines[0].strip() == "the answer" assert "← m (round 1)" in lines[1] +def test_sink_flush_renders_markdown(theme): + """The completed answer is formatted: headings/lists render as Markdown.""" + feed, out = make_feed(theme) + feed.stream_sink = lambda text: None + feed.stream_clear = lambda: None + feed.stream_start() + feed.stream_token("# Head\n\n- item\n") + feed.stream_end() + rendered = out.getvalue() + assert "Head" in rendered + assert "item" in rendered + assert "# Head" not in rendered + + def test_sink_flush_is_idempotent(theme): feed, out = make_feed(theme) cleared: list[bool] = [] diff --git a/tests/test_tui_question.py b/tests/test_tui_question.py index be41b18..63b31ad 100644 --- a/tests/test_tui_question.py +++ b/tests/test_tui_question.py @@ -281,7 +281,14 @@ async def test_pipe_question_escape_dismisses(tmp_path, monkeypatch): assert await task == 0 rendered = out.getvalue() assert '"dismissed":true' in rendered - assert "best judgment" in rendered + # The transcript shows the tool result summary; the full result (with the + # "best judgment" instruction) is on the persisted tool message. + tool_texts = [ + str(record.message.get("content")) + for record in app.store.load_messages(app.session) + if record.role == "tool" + ] + assert any("best judgment" in text for text in tool_texts) async def test_pipe_question_multi_select_toggle_confirm(tmp_path, monkeypatch): From aba712817088da7df427bfe0c175fee29c746e5e Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Mon, 14 Sep 2026 19:27:02 +0400 Subject: [PATCH 3/8] fix: render the roster window as ANSI, not Rich Text FormattedTextControl hashes its fragments when caching content; Rich Text objects are unhashable, so the first render with a live subagent raised 'cannot use tuple as a dict key'. Convert the roster/detail rows to ANSI through the console, like the statusline toolbar, and add a render-level regression test. --- src/lecode/tui/app.py | 14 ++++++++++---- tests/test_subagents.py | 5 +++-- tests/test_tui_app.py | 27 +++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index 1379eb7..ad151bd 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -39,7 +39,6 @@ from prompt_toolkit.utils import get_cwidth from prompt_toolkit.widgets import Frame, TextArea from rich.console import Console -from rich.text import Text from lecode.agent.runner import ( AgentRunner, @@ -1806,11 +1805,18 @@ def _term_width(self) -> int: return self._app.output.get_size().columns return 80 - def _roster_text(self) -> list[Text]: + def _roster_text(self) -> ANSI: + """Roster/detail rows as prompt_toolkit text (ANSI via Rich, like the + toolbar). Rich ``Text`` objects are unhashable and cannot be the + fragments a ``FormattedTextControl`` caches.""" width = self._term_width() if self._detail_run_id is not None: - return detail_lines(self._roster.get(self._detail_run_id), self._theme, width) - return roster_lines(self._roster, self._theme, width) + lines = detail_lines(self._roster.get(self._detail_run_id), self._theme, width) + else: + lines = roster_lines(self._roster, self._theme, width) + with self._console.capture() as capture: + self._console.print(*lines, sep="\n") + return ANSI(capture.get()) # -- agents / totals ---------------------------------------------------------- diff --git a/tests/test_subagents.py b/tests/test_subagents.py index 94186bd..3b94dea 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -7,6 +7,7 @@ from typing import Any import pytest +from prompt_toolkit.formatted_text import to_formatted_text from tests.fakes import FakeProvider from tests.test_tui_app import make_app, make_blocking_app, wait_for @@ -600,12 +601,12 @@ async def test_roster_panel_visible_while_child_runs(tmp_path, monkeypatch): await wait_for(lambda: provider.child_started.is_set()) assert app._roster_visible() - rows = "\n".join(line.plain for line in app._roster_text()) + rows = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) assert "Scan repo" in rows run_id = app.roster.runs()[0].run_id assert app.open_agent_run(run_id) - detail = "\n".join(line.plain for line in app._roster_text()) + detail = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) assert "Scan repo" in detail app._turn_task.cancel() diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index e319aae..78c543b 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -109,6 +109,33 @@ def test_layout_is_chatbox_above_statusline(tmp_path, monkeypatch): assert app._live_buffer is not None +def test_roster_panel_renders_as_prompt_toolkit_text(tmp_path, monkeypatch): + """A live run must not feed Rich Text into FormattedTextControl: its + fragment cache hashes the fragments, and Rich Text is unhashable.""" + from prompt_toolkit.layout.controls import FormattedTextControl + + from lecode.agent.runner import LlmCall + from lecode.extras.subagents import SubagentProgress + + app, _, _ = make_app(tmp_path, monkeypatch, []) + app._on_child_event( + SubagentProgress( + run_id="r1", + agent="explore", + description="Scan src", + event=LlmCall(model="m", turn=1), + ) + ) + for detail in (False, True): + if detail: + assert app.open_agent_run("r1") is True + content = FormattedTextControl(app._roster_text).create_content(80, None) + rendered = "".join( + fragment[1] for i in range(content.line_count) for fragment in content.get_line(i) + ) + assert "Scan src" in rendered + + def _buffer(app): return app._input_area.buffer From 248bb2a8347bdd96d9756bf7a20f1f16ac232992 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Tue, 15 Sep 2026 14:28:05 +0400 Subject: [PATCH 4/8] feat: add persistent worker subagents --- docs/configuration.md | 16 + src/lecode/agent/builder.py | 10 + src/lecode/agent/runner.py | 27 +- src/lecode/agent/tools/task.py | 45 ++ src/lecode/agent/tools/workers.py | 125 +++++ src/lecode/cli.py | 10 + src/lecode/config/models.py | 9 + src/lecode/extras/workers.py | 775 +++++++++++++++++++++++++++++ src/lecode/extras/worktree.py | 496 ++++++++++++++++++- src/lecode/permission/checker.py | 153 ++++-- src/lecode/session/stats.py | 24 +- src/lecode/session/storage.py | 10 +- src/lecode/slash/catalog.py | 1 + src/lecode/slash/handlers.py | 81 ++- src/lecode/tui/agents.py | 133 ++++- src/lecode/tui/app.py | 271 +++++++++- src/lecode/tui/feed.py | 13 +- src/lecode/tui/permission.py | 66 ++- tests/test_agent_builder.py | 10 + tests/test_agent_runner.py | 34 ++ tests/test_config_loader.py | 13 + tests/test_permission_checker.py | 229 ++++++++- tests/test_session_stats.py | 50 ++ tests/test_session_storage.py | 32 ++ tests/test_slash_features.py | 43 ++ tests/test_subagents.py | 22 +- tests/test_tui_agents.py | 42 ++ tests/test_tui_app.py | 34 ++ tests/test_tui_permission.py | 111 +++++ tests/test_tui_streaming_pty.py | 25 + tests/test_workers.py | 786 ++++++++++++++++++++++++++++++ tests/test_worktree.py | 388 +++++++++++++++ 32 files changed, 3973 insertions(+), 111 deletions(-) create mode 100644 src/lecode/agent/tools/workers.py create mode 100644 src/lecode/extras/workers.py create mode 100644 tests/test_workers.py diff --git a/docs/configuration.md b/docs/configuration.md index 9c8f4e8..0e0a60d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -180,6 +180,22 @@ trouble never blocks the agent. See [memory.md](memory.md). +## `[worktree]` + +| field | default | meaning | +|---|---|---| +| `validation` | `[]` | commands a human-approved caller runs in the worker tree before fast-forward integration | + +For example: + +```toml +[worktree] +validation = ["uv run ruff check", "uv run python -m pytest"] +``` + +The worktree helper does not execute these commands itself. Its caller supplies +the approved command runner; an empty list requires explicit human authorization. + ## `[pierre]` Post-task reviewer: after every completed task, a second model compares the diff --git a/src/lecode/agent/builder.py b/src/lecode/agent/builder.py index 6e9e4eb..648f262 100644 --- a/src/lecode/agent/builder.py +++ b/src/lecode/agent/builder.py @@ -53,6 +53,7 @@ def build_runtime( agent_registry: AgentRegistry | None = None, skill_registry: SkillRegistry | None = None, catalog: Catalog | None = None, + worker_manager: object | None = None, ) -> Runtime: """Build the permission checker, tool context, registry, and system prompt. @@ -93,6 +94,15 @@ def build_runtime( from lecode.agent.tools import task as task_tool tools.append(task_tool.make_tool()) + if session is not None and store is not None: + # Local: workers imports this module to build child runtimes. + from lecode.agent.tools import workers as workers_tool + from lecode.extras.workers import WORKER_EXTRA, WorkerManager + + ctx.extras[WORKER_EXTRA] = worker_manager or WorkerManager( + config, cwd=Path(cwd), root_ctx=ctx, session=session, store=store + ) + tools.append(workers_tool.make_tool()) if config.memory.enabled: memory_store = MemoryStore(memory_root(cwd), max_bytes=config.memory.max_bytes) ctx.extras["memory"] = memory_store diff --git a/src/lecode/agent/runner.py b/src/lecode/agent/runner.py index 066d848..3cbf193 100644 --- a/src/lecode/agent/runner.py +++ b/src/lecode/agent/runner.py @@ -282,6 +282,7 @@ async def run( self.ctx.extras["conversation"] = history # Background tasks finished between runs surface at the start. await self._drain_background(history, on_event) + await self._consume_workers(history, on_event) input_tokens = 0 output_tokens = 0 cost_usd = 0.0 @@ -303,6 +304,7 @@ async def run( break if turns > 0: await self._drain_queues(history, on_event) + await self._consume_workers(history, on_event) if self.config.agent.turn_cooldown_ms > 0: await asyncio.sleep(self.config.agent.turn_cooldown_ms / 1000) @@ -367,6 +369,8 @@ async def run( continuing = True continue final_text = (final_text if continuing else "") + completed.content + if await self._consume_workers(history, on_event): + continue stop_reason = "done" break except asyncio.CancelledError: @@ -539,8 +543,19 @@ async def _run_tools( ) for call in completed.tool_calls ] + manager = self.ctx.extras.get("workers") + worker_id = self.ctx.extras.get("worker_id") + suspend = ( + manager is not None + and worker_id is not None + and all(call["function"]["name"] == "task" for call in completed.tool_calls) + ) try: - pairs = await asyncio.gather(*tasks) + if suspend: + async with manager.suspend(worker_id): + pairs = await asyncio.gather(*tasks) + else: + pairs = await asyncio.gather(*tasks) except asyncio.CancelledError: # Cancel in-flight tools; persist the results that did complete. for task in tasks: @@ -603,6 +618,16 @@ async def _drain_background(self, history: list[ChatMessage], on_event: OnEvent self._persist_message(message) await self._emit(on_event, QueuedMessage(content=note)) + async def _consume_workers(self, history: list[ChatMessage], on_event: OnEvent | None) -> bool: + """Deliver worker inboxes only between model turns, never mid tool batch.""" + manager = self.ctx.extras.get("workers") + if manager is None: + return False + items = manager.consume(self.ctx.extras.get("worker_id"), history) + for item in items: + await self._emit(on_event, QueuedMessage(content=item["text"])) + return bool(items) + # -- automatic compaction ----------------------------------------------------- def _context_window(self) -> int: diff --git a/src/lecode/agent/tools/task.py b/src/lecode/agent/tools/task.py index 23af4f6..6df2c10 100644 --- a/src/lecode/agent/tools/task.py +++ b/src/lecode/agent/tools/task.py @@ -24,6 +24,8 @@ SubagentError, run_subagent, ) +from lecode.extras.workers import WORKER_EXTRA +from lecode.extras.worktree import WorktreeError #: Subagent used when the call does not name one. DEFAULT_AGENT = "explore" @@ -75,6 +77,14 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: agents = ctx.extras.get(AGENTS_EXTRA) if registry is None or agents is None: return ToolResult("error: subagents are unavailable in this context", is_error=True) + # The TUI still owns its transient roster through run_subagent. + manager = ( + None + if ctx.extras.get(SUBAGENT_EVENTS_EXTRA) is not None + else ctx.extras.get(WORKER_EXTRA) + ) + if manager is not None: + return await self._start_worker(args, ctx, manager, prompt) if args.get("run_in_background"): return self._start_background(args, ctx, registry, agents, prompt) try: @@ -101,6 +111,41 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: }, ) + async def _start_worker( + self, args: dict[str, Any], ctx: ToolContext, manager: Any, prompt: str + ) -> ToolResult: + agent = str(args.get("agent") or DEFAULT_AGENT) + description = str(args.get("description") or prompt[:60]) + background = bool(args.get("run_in_background")) + try: + worker = await manager.start( + ctx, + agent=agent, + prompt=prompt, + description=description, + background=background, + ) + if background: + return ToolResult( + f"worker {worker.id} started ({agent}): {description}", + metadata={"worker_id": worker.id, "agent": agent}, + ) + outcome = await manager.wait(worker.id) + except (SubagentError, WorktreeError, RuntimeError) as e: + return ToolResult(f"error: {e}", is_error=True) + return ToolResult( + outcome.final_text or "(subagent returned no text)", + metadata={ + "agent": agent, + "worker_id": worker.id, + "run_id": worker.id, + "turns": outcome.turns, + "input_tokens": outcome.usage_totals.input_tokens, + "output_tokens": outcome.usage_totals.output_tokens, + "cost_usd": outcome.usage_totals.cost_usd, + }, + ) + def _start_background( self, args: dict[str, Any], diff --git a/src/lecode/agent/tools/workers.py b/src/lecode/agent/tools/workers.py new file mode 100644 index 0000000..dd37372 --- /dev/null +++ b/src/lecode/agent/tools/workers.py @@ -0,0 +1,125 @@ +"""Control persisted worker runs from their supervising conversation.""" + +from __future__ import annotations + +from typing import Any + +from lecode.agent.tools.base import Tool, ToolContext, ToolResult +from lecode.extras.subagents import SubagentError +from lecode.extras.workers import WORKER_CURRENT_EXTRA, WORKER_EXTRA + +_ACTIONS = ("list", "send", "stop", "resume", "submit", "question", "integrate", "cleanup") + + +class WorkersTool(Tool): + def __init__(self) -> None: + super().__init__( + name="workers", + description="List and control delegated workers. Workers can manage descendants only.", + parameters={ + "type": "object", + "additionalProperties": False, + "properties": { + "action": {"type": "string", "enum": list(_ACTIONS)}, + "id": {"type": "string"}, + "text": {"type": "string"}, + "interrupt": {"type": "boolean"}, + "tree": {"type": "boolean"}, + }, + "required": ["action"], + }, + ) + + @staticmethod + def _descendants(manager: Any, parent_id: str) -> set[str]: + ids = set() + pending = [parent_id] + while pending: + parent = pending.pop() + for child in manager.children(parent): + ids.add(child.id) + pending.append(child.id) + return ids + + def _validate(self, args: dict[str, Any]) -> str | None: + action = args.get("action") + allowed = { + "list": {"action"}, + "send": {"action", "id", "text", "interrupt"}, + "stop": {"action", "id", "tree"}, + "resume": {"action", "id", "text"}, + "submit": {"action", "id"}, + "question": {"action", "text"}, + "integrate": {"action"}, + "cleanup": {"action"}, + } + if action not in _ACTIONS or set(args) - allowed[action]: + return "invalid workers action or arguments" + if action in {"send", "stop", "resume", "submit"} and not isinstance(args.get("id"), str): + return f"workers {action} needs an id" + if action in {"send", "question"} and not isinstance(args.get("text"), str): + return f"workers {action} needs text" + if not isinstance(args.get("interrupt", False), bool) or not isinstance( + args.get("tree", False), bool + ): + return "interrupt and tree must be booleans" + return None + + async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: + error = self._validate(args) + if error is not None: + return ToolResult(f"error: {error}", is_error=True) + manager = ctx.extras.get(WORKER_EXTRA) + if manager is None: + return ToolResult("error: workers are unavailable in this context", is_error=True) + action = args["action"] + if action in {"integrate", "cleanup"}: + return ToolResult(f"error: workers {action} is unavailable", is_error=True) + current = ctx.extras.get(WORKER_CURRENT_EXTRA) + if action == "list": + workers = ( + manager.list() + if current is None + else [ + worker + for worker in manager.list() + if worker.id in self._descendants(manager, current) + ] + ) + if not workers: + return ToolResult("(no workers)") + return ToolResult( + "\n".join( + f"{worker.id} {worker.state} {worker.agent} {worker.description}" + for worker in workers + ) + ) + id = args.get("id") + if action == "question": + if current is None: + return ToolResult("error: the main agent has no parent", is_error=True) + try: + manager.ask_parent(current, args["text"]) + except (KeyError, SubagentError) as e: + return ToolResult(f"error: {e}", is_error=True) + return ToolResult("question sent to parent") + if current is not None and id not in self._descendants(manager, current): + return ToolResult("error: workers can manage descendants only", is_error=True) + try: + if action == "send": + message_id = await manager.send(id, args["text"], bool(args.get("interrupt"))) + return ToolResult(f"worker {id} message {message_id} queued") + if action == "stop": + await manager.stop(id, bool(args.get("tree"))) + return ToolResult(f"worker {id} stopped") + if action == "resume": + await manager.resume(id, args.get("text")) + return ToolResult(f"worker {id} resumed") + note = await manager.submit(id) + return ToolResult(f"worker {id} submitted", metadata={"notification_id": note["id"]}) + except (KeyError, RuntimeError, SubagentError) as e: + return ToolResult(f"error: {e}", is_error=True) + + +def make_tool() -> Tool: + return WorkersTool() diff --git a/src/lecode/cli.py b/src/lecode/cli.py index 6be7534..987d53d 100644 --- a/src/lecode/cli.py +++ b/src/lecode/cli.py @@ -35,6 +35,7 @@ run_plan_loop, ) from lecode.extras.status_signals import START, STOP, StatusEmitter +from lecode.extras.workers import WORKER_EXTRA from lecode.extras.worktree import WorktreeError, WorktreeInfo, WorktreeManager from lecode.hooks import ( EVENTS, @@ -185,6 +186,9 @@ async def _run_headless( try: return await runner.run(messages) finally: + workers = runner.ctx.extras.get(WORKER_EXTRA) + if workers is not None: + await workers.shutdown() aclose = getattr(provider, "aclose", None) if aclose is not None: await aclose() @@ -482,6 +486,9 @@ async def _loop() -> LoopResult: background = runtime.ctx.extras.get(BACKGROUND_EXTRA) if background is not None: await background.shutdown() + workers = runtime.ctx.extras.get(WORKER_EXTRA) + if workers is not None: + await workers.shutdown() await _aclose(client) signals.emit(START) @@ -585,6 +592,9 @@ async def _chain() -> ChainResult: background = runtime.ctx.extras.get(BACKGROUND_EXTRA) if background is not None: await background.shutdown() + workers = runtime.ctx.extras.get(WORKER_EXTRA) + if workers is not None: + await workers.shutdown() await _aclose(client) signals.emit(START) diff --git a/src/lecode/config/models.py b/src/lecode/config/models.py index 31897eb..67ef166 100644 --- a/src/lecode/config/models.py +++ b/src/lecode/config/models.py @@ -213,6 +213,14 @@ class MemoryConfig(BaseModel): max_bytes: int = 32768 +class WorktreeConfig(BaseModel): + """``[worktree]`` -- worker worktree integration checks.""" + + model_config = ConfigDict(extra="ignore") + + validation: list[str] = Field(default_factory=list) + + class PierreConfig(BaseModel): """``[pierre]`` — post-task reviewer: a second model compares the request with the result and gives feedback after every completed task.""" @@ -276,6 +284,7 @@ class Config(BaseModel): mcp: McpConfig = Field(default_factory=McpConfig) lsp: LspConfig = Field(default_factory=LspConfig) memory: MemoryConfig = Field(default_factory=MemoryConfig) + worktree: WorktreeConfig = Field(default_factory=WorktreeConfig) pierre: PierreConfig = Field(default_factory=PierreConfig) telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) hooks: dict[str, list[str]] = Field(default_factory=dict) diff --git a/src/lecode/extras/workers.py b/src/lecode/extras/workers.py new file mode 100644 index 0000000..cfe7baf --- /dev/null +++ b/src/lecode/extras/workers.py @@ -0,0 +1,775 @@ +"""Persistent child runners. Runner boundary wiring is deliberately external. + +``consume`` may only be called at a safe conversation boundary. ``suspend`` +belongs around the supervisor's whole tool batch, never around individual +concurrent tool calls. Inbox durability does not make external effects atomic. +""" + +from __future__ import annotations + +import asyncio +import inspect +import uuid +from contextlib import asynccontextmanager +from dataclasses import asdict, dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from lecode.agent.builder import build_runtime +from lecode.agent.runner import AgentRunner, LlmResponse, RunResult, UsageTotals +from lecode.agent.tools.base import ToolContext +from lecode.config.models import Config +from lecode.extras.subagents import SubagentError +from lecode.extras.worktree import WorktreeError, WorktreeInfo, WorktreeManager +from lecode.session.model import EventRecord, MessageRecord +from lecode.session.storage import Session, SessionStore + +WORKER_EXTRA = "workers" +WORKER_CURRENT_EXTRA = "worker_id" +MAX_EXECUTING = 10 +MAX_DEPTH = 2 + + +@dataclass +class Worker: + id: str + parent_id: str | None + depth: int + agent: str + origin: str + state: str + session: Session + cwd: Path + description: str = "" + background: bool = False + worktree: WorktreeInfo | None = None + usage_totals: UsageTotals = field(default_factory=UsageTotals) + usage_incomplete: bool = False + dispatch_id: str | None = None + result: RunResult | None = None + error: str | None = None + started_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + + @property + def session_id(self) -> str: + return self.session.id + + +class WorkerManager: + """Own child sessions, tasks, and locks until shutdown. + + ``confirm(question: str) -> bool`` and ``notify(note: dict) -> None`` may + be synchronous or asynchronous. Notify is observational; parent delivery + happens only through ``drain_notifications(parent_id)`` at a safe boundary. + """ + + def __init__( + self, + config: Config, + *, + cwd: Path, + root_ctx: ToolContext, + session: Session | None = None, + store: SessionStore | None = None, + confirm=None, + notify=None, + ) -> None: + self.config = config + self.cwd = Path(cwd) + self.root_ctx = root_ctx + self.store = store or root_ctx.session_store or SessionStore() + self.session = session or root_ctx.session or self.store.create("workers", self.cwd) + self.confirm = confirm + self.notify = notify + #: UI observer for durable state/usage changes; never affects execution. + self.progress = None + self._workers: dict[str, Worker] = {} + self._contexts: dict[str, Any] = {} + self._runtimes: dict[str, Any] = {} + self._tasks: dict[str, asyncio.Task] = {} + self._progress_tasks: set[asyncio.Task] = set() + self._locks: dict[str, Any] = {} + self._closed = False + self._slots = asyncio.Semaphore(MAX_EXECUTING) + self._leases: set[str] = set() + + def get(self, id: str) -> Worker: + return self._workers[id] + + def list(self) -> list[Worker]: + return list(self._workers.values()) + + def children(self, id: str | None) -> list[Worker]: + return [w for w in self.list() if w.parent_id == id] + + def load(self) -> list[Worker]: + """Attach persisted children without scheduling any work; idempotent.""" + if self._closed: + raise RuntimeError("worker manager is shut down") + latest = {item["id"]: item for item in self.store.load_events(self.session, "worker")} + for id, snapshot in latest.items(): + if id in self._workers: + continue + data = dict(snapshot) + child = self.store.open(data.pop("session_id")) + lock = self.store.acquire_lock(child) + try: + data["session"] = child + data["cwd"] = Path(data["cwd"]) + usages = self.store.load_events(child, "worker_usage_checkpoint") + if usages: + data["usage_totals"] = { + key: value for key, value in usages[-1].items() if key != "dispatch_id" + } + data["usage_totals"] = UsageTotals(**data["usage_totals"]) + if data.get("worktree"): + info = data["worktree"] + data["worktree"] = WorktreeInfo( + info["name"], Path(info["path"]), info["branch"] + ) + if data.get("result"): + result = dict(data["result"]) + result["usage_totals"] = UsageTotals(**result["usage_totals"]) + data["result"] = RunResult(**result) + if data["state"] in {"queued", "running", "waiting"}: + data["state"] = "interrupted" + data["usage_incomplete"] = True + worker = Worker(**data) + except BaseException: + if lock is not None: + lock.release() + raise + self._workers[id] = worker + self._locks[id] = lock + self._record_usage(worker) + return self.list() + + def attach(self, session: Session) -> list[Worker]: + """Move this root manager to an idle session and hydrate its workers.""" + if any(worker.state in {"queued", "running", "waiting"} for worker in self.list()): + raise RuntimeError("workers are still active") + for lock in self._locks.values(): + if lock is not None: + lock.release() + self.session = session + self._workers.clear() + self._contexts.clear() + self._runtimes.clear() + self._tasks.clear() + self._locks.clear() + return self.load() + + def _parent_context(self, worker): + if worker.id not in self._contexts: + self._contexts[worker.id] = ( + self._runtime(self.get(worker.parent_id)).ctx if worker.parent_id else self.root_ctx + ) + return self._contexts[worker.id] + + def _record(self, worker: Worker) -> None: + data = { + key: getattr(worker, key) + for key in ( + "id", + "parent_id", + "depth", + "agent", + "origin", + "state", + "description", + "background", + "usage_incomplete", + "dispatch_id", + "error", + "started_at", + ) + } + data.update( + session_id=worker.session_id, + cwd=str(worker.cwd), + usage_totals=asdict(worker.usage_totals), + result=asdict(worker.result) if worker.result else None, + ) + if worker.worktree is not None: + data["worktree"] = {**asdict(worker.worktree), "path": str(worker.worktree.path)} + self.store.append_event(self.session, "worker", data) + if self.progress is not None: + result = self.progress(worker) + if inspect.isawaitable(result): + task = asyncio.create_task(result) + self._progress_tasks.add(task) + task.add_done_callback(self._progress_tasks.discard) + + def _record_usage(self, worker): + # ponytail: scan the root ledger per dispatch; index if session size warrants it. + recorded = [ + r.data + for r in self.store.read_records(self.session) + if isinstance(r, EventRecord) + and r.kind == "worker_usage" + and r.data.get("worker_id") == worker.id + ] + delta = { + key: getattr(worker.usage_totals, key) + - sum(item["usage"].get(key, 0) for item in recorded) + for key in ("input_tokens", "output_tokens", "cost_usd") + } + if ( + recorded + and not any(delta.values()) + and recorded[-1]["dispatch_id"] == worker.dispatch_id + and recorded[-1]["incomplete"] == worker.usage_incomplete + ): + return + self.store.append_event( + self.session, + "worker_usage", + { + "worker_id": worker.id, + "dispatch_id": worker.dispatch_id, + "usage": delta, + "incomplete": worker.usage_incomplete, + }, + ) + + def _agent(self, ctx, name): + agents = ctx.extras.get("agents") + if agents is None or name not in {a.name for a in agents.subagents()}: + raise SubagentError(f"unknown or ineligible subagent: {name}") + return agents.get(name) + + @staticmethod + def _read_only(ctx, definition): + return ctx.permission_checker.for_child(definition.overlay, cwd=ctx.cwd).read_only + + async def _workspace(self, ctx, definition, id): + if self._read_only(ctx, definition): + return Path(ctx.cwd), None + manager = await WorktreeManager.discover(ctx.cwd) + branch = await manager._git("rev-parse", "--abbrev-ref", "HEAD", cwd=ctx.cwd) + if branch == "HEAD": + raise WorktreeError("write workers require an attached branch, not detached HEAD") + destination = Path(await manager._git("rev-parse", "--show-toplevel", cwd=ctx.cwd)) + base = await manager._git("rev-parse", "HEAD", cwd=ctx.cwd) + dirty_root = await manager._git("status", "--porcelain", cwd=self.cwd) + dirty_parent = await manager._git("status", "--porcelain", cwd=ctx.cwd) + if dirty_root or dirty_parent: + question = ( + "Uncommitted changes will not enter the worker's committed-HEAD worktree. Continue?" + ) + approved = self.confirm(question) if self.confirm is not None else False + if inspect.isawaitable(approved): + approved = await approved + if approved is not True: + raise WorktreeError("dirty root/parent requires human confirmation") + info = await manager.create_worker( + id, base_commit=base, dest_path=destination, dest_branch=branch + ) + return info.path, info + + async def start( + self, + ctx: ToolContext, + *, + agent: str, + prompt: str, + description: str = "", + origin: str = "delegated", + background: bool = False, + ) -> Worker: + if self._closed: + raise RuntimeError("worker manager is shut down") + definition = self._agent(ctx, agent) + if ctx.extras.get("provider") is None: + raise SubagentError("no provider available for workers") + parent_id = ctx.extras.get(WORKER_CURRENT_EXTRA) + depth = self.get(parent_id).depth + 1 if parent_id else 1 + if depth > MAX_DEPTH: + raise SubagentError(f"worker depth exceeds {MAX_DEPTH}") + id = uuid.uuid4().hex + cwd, worktree = await self._workspace(ctx, definition, id) + if self._closed: + raise RuntimeError("worker manager is shut down") + session = self.store.create( + f"worker-{id}", + cwd, + agent=agent, + model=definition.model or self.config.agent.subagent_model or self.config.llm.model, + ) + self._locks[id] = self.store.acquire_lock(session) + worker = Worker( + id, + parent_id, + depth, + agent, + origin, + "queued", + session, + cwd, + description=description, + background=background, + worktree=worktree, + ) + self._workers[id] = worker + self._contexts[id] = ctx + self._record(worker) + self._enqueue(worker, prompt) + self._launch(worker) + return worker + + def _runtime(self, worker): + if worker.id not in self._runtimes: + parent = self._parent_context(worker) + definition = self._agent(parent, worker.agent) + config = self.config.model_copy( + update={"pierre": self.config.pierre.model_copy(update={"enabled": False})} + ) + runtime = build_runtime( + config, + worker.cwd, + session=worker.session, + store=self.store, + agent_name=worker.agent, + agent_registry=parent.extras["agents"], + allowed_tools=parent.extras["registry"].names(), + catalog=parent.catalog, + worker_manager=self, + ) + runtime.ctx.permission_checker = parent.permission_checker.for_child( + definition.overlay, + cwd=worker.cwd, + session_perms=runtime.ctx.session_perms, + read_only=self._read_only(parent, definition), + ) + runtime.ctx.auto_approve = parent.auto_approve + callback = parent.approval_callback + if callback is not None: + + async def approval_callback(tool_name, args, reason): + """Keep worker identity with the root approval FIFO.""" + params = inspect.signature(callback).parameters + if "worker" in params or any( + param.kind is inspect.Parameter.VAR_KEYWORD for param in params.values() + ): + result = callback( + tool_name, + args, + reason, + worker=worker.id, + conversation=worker.session.name, + ) + else: + result = callback(tool_name, args, reason) + return await result if inspect.isawaitable(result) else result + + runtime.ctx.approval_callback = approval_callback + runtime.ctx.extras.update( + { + WORKER_EXTRA: self, + WORKER_CURRENT_EXTRA: worker.id, + "provider": parent.extras["provider"], + } + ) + runtime.registry.unregister("ask_user") + self._runtimes[worker.id] = runtime + return self._runtimes[worker.id] + + def _enqueue(self, worker, text): + if self._closed: + raise RuntimeError("worker manager is shut down") + item = {"id": uuid.uuid4().hex, "text": text} + self.store.append_event(worker.session, "worker_inbox", item) + return item["id"] + + def pending(self, id: str) -> list[dict[str, str]]: + worker = self.get(id) + consumed = { + r.usage["worker_inbox_id"] + for r in self.store.read_records(worker.session) + if isinstance(r, MessageRecord) and r.usage and "worker_inbox_id" in r.usage + } + return [ + item + for item in self.store.load_events(worker.session, "worker_inbox") + if item["id"] not in consumed + ] + + @staticmethod + def _outstanding(history): + outstanding = {} + for message in history: + for call in message.get("tool_calls") or []: + outstanding[call["id"]] = call + if message.get("role") == "tool": + outstanding.pop(message.get("tool_call_id"), None) + return outstanding + + @staticmethod + def _notification_text(note: dict[str, Any]) -> str: + if note.get("kind") == "question": + return f"[worker {note['worker_id']} asks] {note['content']}" + if note.get("kind") == "human_message": + return f"[human → worker {note['worker_id']}] {note['content']}" + return f"[worker {note['worker_id']} {note['state']}] {note['content']}" + + def consume(self, id: str | None, history: list[dict]) -> list[dict[str, str]]: + """Append pending inputs durably and to live history at a safe boundary.""" + if self._closed: + raise RuntimeError("worker manager is shut down") + worker = self.get(id) if id is not None else None + session = worker.session if worker is not None else self.session + if self._outstanding(history) or self._outstanding(self.store.load_for_model(session)): + raise RuntimeError("cannot consume inbox with outstanding tool calls") + items = self.pending(id) if id is not None else [] + for item in items: + message = {"role": "user", "content": item["text"]} + self.store.append_message(session, message, usage={"worker_inbox_id": item["id"]}) + history.append(message) + acknowledged = { + item["id"] for item in self.store.load_events(self.session, "worker_notification_ack") + } + notes = [ + note + for note in self.store.load_events(self.session, "worker_notification") + if note["deliver"] and note["parent_id"] == id and note["id"] not in acknowledged + ] + for note in notes: + text = self._notification_text(note) + self.store.append_message( + session, + {"role": "user", "content": text}, + usage={"worker_notification_id": note["id"]}, + ) + self.store.append_event(self.session, "worker_notification_ack", {"id": note["id"]}) + history.append({"role": "user", "content": text}) + items.append({"id": note["id"], "text": text}) + return items + + def _launch(self, worker): + if self._closed: + raise RuntimeError("worker manager is shut down") + existing = self._tasks.get(worker.id) + if existing is not None and not existing.done(): + raise RuntimeError("worker is already active") + worker.state = "queued" + worker.error = None + worker.result = None + worker.dispatch_id = uuid.uuid4().hex + self._record(worker) + self._tasks[worker.id] = asyncio.create_task(self._execute(worker)) + self._tasks[worker.id].add_done_callback(lambda task: self._finished(worker, task)) + + def _finished(self, worker, task): + if ( + not self._closed + and self._tasks.get(worker.id) is task + and worker.state == "completed" + and self.pending(worker.id) + ): + self._launch(worker) + + def _background_descendants(self, id: str) -> list[Worker]: + pending = [id] + descendants: list[Worker] = [] + while pending: + parent = pending.pop() + children = self.children(parent) + descendants.extend(children) + pending.extend(child.id for child in children) + return [worker for worker in descendants if worker.background] + + async def _execute(self, worker): + try: + await self._slots.acquire() + self._leases.add(worker.id) + worker.state = "running" + self._record(worker) + + runtime = self._runtime(worker) + definition = self._agent(self._parent_context(worker), worker.agent) + runner = AgentRunner( + runtime.ctx.extras["provider"], + runtime.registry, + runtime.ctx, + session=worker.session, + store=self.store, + config=runtime.ctx.config, + catalog=runtime.ctx.catalog, + ) + runner.model = ( + definition.model or self.config.agent.subagent_model or self.config.llm.model + ) + while True: + history = self.store.load_for_model(worker.session) + self.consume(worker.id, history) + worker.result = await runner.run( + [{"role": "system", "content": runtime.system_prompt}, *history], + on_event=lambda event: self._event(worker, event), + ) + active = [ + child + for child in self._background_descendants(worker.id) + if child.state in {"queued", "running", "waiting"} + ] + if active: + # Let queued descendants use this supervisor's lease. + async with self.suspend(worker.id): + await asyncio.gather( + *(self.wait(child.id) for child in active), return_exceptions=True + ) + if not self.pending(worker.id): + # Background completions are durable notes, consumed by the + # next child run rather than racing a final response. + acknowledged = { + item["id"] + for item in self.store.load_events(self.session, "worker_notification_ack") + } + if not any( + note["deliver"] + and note["parent_id"] == worker.id + and note["id"] not in acknowledged + for note in self.store.load_events(self.session, "worker_notification") + ): + break + worker.state = "completed" + except asyncio.CancelledError: + worker.state = "stopped" + worker.usage_incomplete = True + except Exception as error: + worker.state = "failed" + worker.error = f"{type(error).__name__}: {error}" + worker.usage_incomplete = True + finally: + if worker.id in self._leases: + self._leases.remove(worker.id) + self._slots.release() + worker.usage_incomplete |= any( + isinstance(record, MessageRecord) + and record.role == "assistant" + and record.usage is None + for record in self.store.read_records(worker.session) + ) + self._record_usage(worker) + self._record(worker) + + if worker.background or worker.origin == "human": + note = self._notification(worker, submitted=False) + if self.notify is not None: + try: + result = self.notify(note) + if inspect.isawaitable(result): + await result + except Exception as error: + self.store.append_event( + self.session, + "worker_notify_error", + { + "worker_id": worker.id, + "error": str(error), + }, + ) + + @asynccontextmanager + async def suspend(self, worker_id: str): + """Yield the supervisor lease while awaiting its entire tool batch. + + Only the worker runner task may suspend itself. Concurrent tool tasks + cannot release a supervisor that is still executing sibling tools. + """ + if worker_id not in self._leases or asyncio.current_task() is not self._tasks.get( + worker_id + ): + raise RuntimeError("suspend must be called once by the worker supervisor") + worker = self.get(worker_id) + self._leases.remove(worker_id) + self._slots.release() + worker.state = "waiting" + self._record(worker) + try: + yield + finally: + acquire = asyncio.create_task(self._slots.acquire()) + try: + await asyncio.shield(acquire) + except asyncio.CancelledError: + await asyncio.shield(acquire) + raise + self._leases.add(worker_id) + worker.state = "running" + self._record(worker) + + def _event(self, worker, event): + if isinstance(event, LlmResponse): + old = worker.usage_totals + worker.usage_totals = UsageTotals( + old.input_tokens + event.input_tokens, + old.output_tokens + event.output_tokens, + old.cost_usd + event.cost_usd, + event.input_tokens or old.context_tokens, + ) + self.store.append_event( + worker.session, + "worker_usage_checkpoint", + { + "dispatch_id": worker.dispatch_id, + **asdict(worker.usage_totals), + }, + ) + self._record(worker) + + async def wait(self, id: str) -> RunResult: + worker = self.get(id) + task = self._tasks.get(id) + while task is not None: + await asyncio.shield(task) + current = self._tasks.get(id) + if current is task: + break + task = current + if worker.result is None or worker.state != "completed": + raise SubagentError(worker.error or f"worker is {worker.state}") + return worker.result + + async def send( + self, id: str, text: str, interrupt: bool = False, *, from_human: bool = False + ) -> str: + worker = self.get(id) + message_id = self._enqueue(worker, text) + if from_human: + self.store.append_event( + self.session, + "worker_notification", + { + "id": uuid.uuid4().hex, + "worker_id": worker.id, + "parent_id": worker.parent_id, + "dispatch_id": worker.dispatch_id, + "agent": worker.agent, + "origin": "human", + "state": worker.state, + "kind": "human_message", + "content": text, + "deliver": True, + }, + ) + if interrupt and worker.state in {"queued", "running", "waiting"}: + await self.stop(id) + await self.resume(id) + elif worker.state == "completed" and (id not in self._tasks or self._tasks[id].done()): + self._launch(worker) + return message_id + + async def stop(self, id: str, tree: bool = False) -> None: + worker = self.get(id) + worker.state = "stopped" + self._record(worker) + task = self._tasks.get(id) + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + if tree: + for child in self.children(id): + await self.stop(child.id, tree=True) + + async def resume(self, id: str, text: str | None = None) -> Worker: + if self._closed: + raise RuntimeError("worker manager is shut down") + worker = self.get(id) + task = self._tasks.get(id) + if task is not None and not task.done(): + raise RuntimeError("worker is already active") + # Finish protocol pairs, not the interrupted external actions. Replaying + # those actions could duplicate an effect whose outcome is unknown. + for call in self._outstanding(self.store.load_for_model(worker.session)).values(): + self.store.append_message( + worker.session, + { + "role": "tool", + "tool_call_id": call["id"], + "name": call["function"]["name"], + "content": ( + "Worker interrupted; tool outcome unknown. Inspect state before retrying." + ), + }, + ) + if text is not None: + self._enqueue(worker, text) + self._launch(worker) + return worker + + def _notification(self, worker, *, submitted): + note = { + "id": f"{worker.dispatch_id}:{'submitted' if submitted else 'completed'}", + "worker_id": worker.id, + "parent_id": worker.parent_id, + "dispatch_id": worker.dispatch_id, + "agent": worker.agent, + "origin": worker.origin, + "state": worker.state, + "content": worker.result.final_text if worker.result else worker.error or worker.state, + "deliver": submitted or (worker.origin == "delegated" and worker.background), + } + existing = self.store.load_events(self.session, "worker_notification") + if not any(item["id"] == note["id"] for item in existing): + self.store.append_event(self.session, "worker_notification", note) + return note + + async def submit(self, id: str) -> dict[str, Any]: + """Explicitly make a human-origin result available to its parent.""" + worker = self.get(id) + if worker.state != "completed" or worker.result is None: + raise SubagentError("only completed workers can be submitted") + return self._notification(worker, submitted=True) + + def ask_parent(self, id: str, text: str) -> dict[str, Any]: + """Durably deliver a worker question without exposing interactive UI.""" + worker = self.get(id) + if worker.parent_id is None: + raise SubagentError("the root worker has no parent") + note = { + "id": uuid.uuid4().hex, + "worker_id": worker.id, + "parent_id": worker.parent_id, + "dispatch_id": worker.dispatch_id, + "agent": worker.agent, + "origin": worker.origin, + "state": worker.state, + "kind": "question", + "content": text, + "deliver": True, + } + self.store.append_event(self.session, "worker_notification", note) + return note + + def drain_notifications(self, parent_id: str | None = None) -> list[dict[str, Any]]: + """Acknowledge parent-bound result notes; human results require submit.""" + acknowledged = { + item["id"] for item in self.store.load_events(self.session, "worker_notification_ack") + } + notes = [ + item + for item in self.store.load_events(self.session, "worker_notification") + if item["deliver"] and item["parent_id"] == parent_id and item["id"] not in acknowledged + ] + for note in notes: + self.store.append_event(self.session, "worker_notification_ack", {"id": note["id"]}) + return notes + + async def shutdown(self) -> None: + self._closed = True + try: + for id, task in self._tasks.items(): + if not task.done(): + await self.stop(id) + for runtime in self._runtimes.values(): + for key in ("background", "lsp"): + resource = runtime.ctx.extras.get(key) + if resource is not None: + result = resource.shutdown() + if inspect.isawaitable(result): + await result + finally: + for lock in self._locks.values(): + if lock is not None: + lock.release() diff --git a/src/lecode/extras/worktree.py b/src/lecode/extras/worktree.py index cba689b..fbe030a 100644 --- a/src/lecode/extras/worktree.py +++ b/src/lecode/extras/worktree.py @@ -11,19 +11,39 @@ worktree branch into the main checkout's current branch. On conflict the merge is left in progress (standard git flow) and the conflicted paths are reported — no auto-resolution, ``git merge --abort`` stays available. + +Worker worktrees: ``create_worker`` pins a base commit and a destination +(path + branch) in a sidecar under ``.lecode/worktrees/.json``; +``integrate`` validates the worker tree and merges it into that pinned +destination. ``discover`` returns the *main* repository root even when +called from inside a linked worktree, so worktrees and sidecars agree +across restarts. """ from __future__ import annotations +import asyncio +import fcntl +import hashlib +import inspect +import json +import os import re +from collections.abc import Callable +from contextlib import asynccontextmanager from dataclasses import dataclass +from datetime import UTC, datetime from pathlib import Path +from typing import Any from lecode.extras.proc import run_proc #: Per-git-command timeout. GIT_TIMEOUT_S = 30.0 +#: Validation output kept in the returned :class:`IntegrationResult`. +_VALIDATION_OUTPUT_LIMIT = 4000 + #: Worktree directory, relative to the repo root. WORKTREE_ROOT = ".lecode/worktrees" @@ -33,6 +53,13 @@ _NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +def _clip_validation(text: str) -> str: + """Keep validation output bounded; the tail usually holds the failure.""" + if len(text) <= _VALIDATION_OUTPUT_LIMIT: + return text + return "[… output clipped …]\n" + text[-_VALIDATION_OUTPUT_LIMIT:] + + class WorktreeError(Exception): """A git or worktree operation failed cleanly.""" @@ -65,15 +92,56 @@ class WorktreeStatus: dirty: bool +@dataclass(frozen=True) +class WorktreeInspection: + """A worker worktree's on-disk state plus its sidecar, if any.""" + + present: bool + info: WorktreeInfo + dirty: bool + merge_in_progress: bool + sidecar: dict[str, Any] | None + + +@dataclass(frozen=True) +class IntegrationResult: + """The outcome of integrating a worker into its pinned destination. + + ``status`` is one of ``integrated``, ``paused``, ``blocked``, + ``conflict``, or ``validation_failed``. + """ + + status: str + detail: str + conflicts: list[str] + validation_output: str + + class WorktreeManager: """Worktree operations over one repository root.""" def __init__(self, repo_root: Path | str) -> None: - self.repo_root = Path(repo_root) + self.repo_root = Path(repo_root).expanduser().resolve() @classmethod async def discover(cls, cwd: Path | str) -> WorktreeManager: - """The manager for the git repository containing ``cwd``.""" + """The manager for the *main* repository containing ``cwd``. + + Inside a linked worktree ``--show-toplevel`` names the worktree, not + the repository that owns it; ``--git-common-dir`` names the main + ``.git``, whose parent is the root we want. + """ + start = Path(cwd).expanduser().resolve() + common = await run_proc( + ["git", "rev-parse", "--git-common-dir"], cwd=cwd, timeout=GIT_TIMEOUT_S + ) + if common.exit_code == 0: + git_dir = Path(common.stdout.strip()) + if not git_dir.is_absolute(): + git_dir = start / git_dir + git_dir = git_dir.resolve() + if git_dir.name == ".git": + return cls(git_dir.parent) result = await run_proc( ["git", "rev-parse", "--show-toplevel"], cwd=cwd, timeout=GIT_TIMEOUT_S ) @@ -89,6 +157,20 @@ async def _git(self, *args: str, cwd: Path | None = None) -> str: raise WorktreeError(f"git {args[0]} failed: {detail or f'exit {result.exit_code}'}") return result.stdout.strip() + async def _common_dir(self, cwd: Path) -> Path: + """Return a resolved common git directory for a working tree.""" + result = await run_proc( + ["git", "rev-parse", "--git-common-dir"], cwd=cwd, timeout=GIT_TIMEOUT_S + ) + if result.exit_code != 0: + raise WorktreeError(f"not a git repository: {cwd}") + common = Path(result.stdout.strip()) + return (common if common.is_absolute() else cwd / common).resolve() + + async def _commit(self, revision: str, *, cwd: Path) -> str: + """Resolve one revision to a commit, without accepting arbitrary refs later.""" + return await self._git("rev-parse", "--verify", f"{revision}^{{commit}}", cwd=cwd) + def _info(self, name: str) -> WorktreeInfo: return WorktreeInfo( name=name, @@ -119,8 +201,8 @@ def _exclude_worktree_root(self) -> None: async def _current_branch(self) -> str: return await self._git("rev-parse", "--abbrev-ref", "HEAD") - async def create(self, name: str) -> WorktreeInfo: - """Create ``.lecode/worktrees/`` on a fresh ``lecode/`` branch.""" + async def _check_new(self, name: str) -> WorktreeInfo: + """Validate ``name`` and that neither its directory nor branch exist.""" if not _NAME_RE.match(name): raise WorktreeError(f"invalid worktree name: {name!r} (letters, digits, . _ -)") info = self._info(name) @@ -133,10 +215,184 @@ async def create(self, name: str) -> WorktreeInfo: ) if branch_ref.exit_code == 0: raise WorktreeError(f"branch already exists: {info.branch}") + return info + + async def create(self, name: str) -> WorktreeInfo: + """Create ``.lecode/worktrees/`` on a fresh ``lecode/`` branch.""" + info = await self._check_new(name) self._exclude_worktree_root() await self._git("worktree", "add", "-b", info.branch, str(info.path)) return info + def _sidecar_path(self, name: str) -> Path: + return self.repo_root / WORKTREE_ROOT / f"{name}.json" + + def write_sidecar(self, name: str, data: dict[str, Any]) -> None: + """Atomically write ``.lecode/worktrees/.json``.""" + if not _NAME_RE.match(name): + raise WorktreeError(f"invalid worktree name: {name!r} (letters, digits, . _ -)") + path = self._sidecar_path(name) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp") + tmp.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, path) + + def read_sidecar(self, name: str) -> dict[str, Any] | None: + """The worker's sidecar, or ``None`` when it was never written.""" + if not _NAME_RE.match(name): + raise WorktreeError(f"invalid worktree name: {name!r} (letters, digits, . _ -)") + try: + text = self._sidecar_path(name).read_text(encoding="utf-8") + except FileNotFoundError: + return None + except OSError as e: + raise WorktreeError(f"cannot read sidecar for '{name}': {e}") from e + try: + data = json.loads(text) + except json.JSONDecodeError as e: + raise WorktreeError(f"invalid sidecar for '{name}': {e}") from e + if not isinstance(data, dict): + raise WorktreeError(f"invalid sidecar for '{name}': expected an object") + return data + + def _validated_sidecar(self, name: str, info: WorktreeInfo) -> dict[str, Any] | None: + """Return a sidecar only when its fixed identity fields are safe to use.""" + data = self.read_sidecar(name) + if data is None: + return None + expected_path = info.path.resolve() + try: + path = Path(data["path"]) + dest_path = Path(data["dest_path"]) + common_dir = Path(data["dest_common_dir"]) + valid = ( + data["name"] == name + and isinstance(data["path"], str) + and path.is_absolute() + and path == path.resolve() == expected_path + and data["branch"] == info.branch + and isinstance(data["base_commit"], str) + and isinstance(data["dest_path"], str) + and dest_path.is_absolute() + and dest_path == dest_path.resolve() + and isinstance(data["dest_branch"], str) + and bool(data["dest_branch"]) + and isinstance(data["dest_common_dir"], str) + and common_dir.is_absolute() + and common_dir == common_dir.resolve() + and ( + data.get("integrated_at") is None or isinstance(data.get("integrated_at"), str) + ) + and ( + data.get("integrated_head") is None + or isinstance(data.get("integrated_head"), str) + ) + ) + except (KeyError, TypeError, ValueError): + valid = False + if not valid: + raise WorktreeError(f"invalid sidecar for '{name}': unsafe worker destination") + return data + + async def _destination_path( + self, dest_path: Path | str, dest_branch: str, *, expected_common: Path + ) -> Path: + """Normalize and verify a destination is this repository's checkout.""" + if not isinstance(dest_branch, str) or not dest_branch: + raise WorktreeError("invalid destination branch") + branch = await run_proc( + ["git", "check-ref-format", "--branch", dest_branch], + cwd=self.repo_root, + timeout=GIT_TIMEOUT_S, + ) + if branch.exit_code != 0: + raise WorktreeError(f"invalid destination branch: {dest_branch!r}") + path = Path(dest_path).expanduser().resolve() + if not path.is_dir(): + raise WorktreeError(f"destination missing: {path}") + root = await run_proc( + ["git", "rev-parse", "--show-toplevel"], cwd=path, timeout=GIT_TIMEOUT_S + ) + if root.exit_code != 0: + raise WorktreeError(f"destination is not a git repository: {path}") + destination = Path(root.stdout.strip()).resolve() + if await self._common_dir(destination) != expected_common: + raise WorktreeError(f"destination is not in this repository: {destination}") + return destination + + async def create_worker( + self, + name: str, + *, + base_commit: str, + dest_path: Path | str, + dest_branch: str, + ) -> WorktreeInfo: + """Create a worker worktree pinned to ``base_commit`` and a destination. + + The branch is ``lecode/``; the pinned destination and base are + recorded in the sidecar so integration resumes deterministically. + """ + info = await self._check_new(name) + common_dir = await self._common_dir(self.repo_root) + destination = await self._destination_path( + dest_path, dest_branch, expected_common=common_dir + ) + base = await self._commit(base_commit, cwd=self.repo_root) + self._exclude_worktree_root() + await self._git("worktree", "add", "-b", info.branch, str(info.path), base) + self.write_sidecar( + name, + { + "name": name, + "path": str(info.path), + "branch": info.branch, + "base_commit": base, + "dest_path": str(destination), + "dest_branch": dest_branch, + "dest_common_dir": str(common_dir), + "integrated_at": None, + "integrated_head": None, + }, + ) + return info + + async def attach(self, name: str) -> WorktreeInfo: + """Resume an existing worker: directory and branch must both exist.""" + info = self._info(name) + if not info.path.is_dir(): + raise WorktreeError(f"no such worktree: {name} (expected at {info.path})") + branch_ref = await run_proc( + ["git", "show-ref", "--verify", f"refs/heads/{info.branch}"], + cwd=self.repo_root, + timeout=GIT_TIMEOUT_S, + ) + if branch_ref.exit_code != 0: + raise WorktreeError(f"worktree '{name}' has no branch {info.branch}") + return info + + async def inspect(self, name: str) -> WorktreeInspection: + """Best-effort worker state; absent worktrees never raise.""" + info = self._info(name) + sidecar = self.read_sidecar(name) + if not info.path.is_dir(): + return WorktreeInspection( + present=False, info=info, dirty=False, merge_in_progress=False, sidecar=sidecar + ) + porcelain = await self._git("status", "--porcelain", cwd=info.path) + merge_head = await run_proc( + ["git", "rev-parse", "-q", "--verify", "MERGE_HEAD"], + cwd=info.path, + timeout=GIT_TIMEOUT_S, + ) + return WorktreeInspection( + present=True, + info=info, + dirty=bool(porcelain), + merge_in_progress=merge_head.exit_code == 0, + sidecar=sidecar, + ) + async def status(self, name: str) -> WorktreeStatus: """Ahead/behind vs the main checkout's branch, plus dirty state.""" info = self._require(name) @@ -176,14 +432,242 @@ async def merge_back(self, name: str) -> MergeResult: merged=False, conflicts=[], message=result.stderr.strip() or "merge failed" ) - async def _conflicted_files(self) -> list[str]: + async def _conflicted_files(self, cwd: Path | None = None) -> list[str]: result = await run_proc( ["git", "diff", "--name-only", "--diff-filter=U"], - cwd=self.repo_root, + cwd=cwd or self.repo_root, timeout=GIT_TIMEOUT_S, ) return [line for line in result.stdout.splitlines() if line.strip()] + async def _destination_problem( + self, dest_path: Path, dest_branch: str, common_dir: Path + ) -> str | None: + """Check the pinned checkout without mutating it.""" + if not dest_path.is_dir(): + raise WorktreeError(f"destination missing: {dest_path}") + root = await run_proc( + ["git", "rev-parse", "--show-toplevel"], cwd=dest_path, timeout=GIT_TIMEOUT_S + ) + if root.exit_code != 0: + raise WorktreeError(f"destination is not a git repository: {dest_path}") + if Path(root.stdout.strip()).resolve() != dest_path: + raise WorktreeError(f"destination path was replaced: {dest_path}") + if await self._common_dir(dest_path) != common_dir: + raise WorktreeError(f"destination repository was replaced: {dest_path}") + ref = await run_proc( + ["git", "show-ref", "--verify", f"refs/heads/{dest_branch}"], + cwd=dest_path, + timeout=GIT_TIMEOUT_S, + ) + if ref.exit_code != 0: + raise WorktreeError(f"destination branch missing: {dest_branch}") + branch = await self._git("rev-parse", "--abbrev-ref", "HEAD", cwd=dest_path) + if branch != dest_branch: + return f"destination is on '{branch}', expected '{dest_branch}'" + if await self._git("status", "--porcelain", cwd=dest_path): + return "destination has uncommitted changes" + return None + + async def _worker_state(self, info: WorktreeInfo, common_dir: Path) -> tuple[str, bool]: + """Return the worker HEAD and dirty state after proving its identity.""" + if not info.path.is_dir(): + raise WorktreeError(f"worker path missing: {info.path}") + if await self._common_dir(info.path) != common_dir: + raise WorktreeError(f"worker repository was replaced: {info.path}") + ref = await run_proc( + ["git", "show-ref", "--verify", f"refs/heads/{info.branch}"], + cwd=info.path, + timeout=GIT_TIMEOUT_S, + ) + if ref.exit_code != 0: + raise WorktreeError(f"worker branch missing: {info.branch}") + branch = await self._git("rev-parse", "--abbrev-ref", "HEAD", cwd=info.path) + if branch != info.branch: + raise WorktreeError( + f"worker branch changed: expected '{info.branch}', found '{branch}'" + ) + return ( + await self._commit("HEAD", cwd=info.path), + bool(await self._git("status", "--porcelain", cwd=info.path)), + ) + + @asynccontextmanager + async def _integration_lock(self, common_dir: Path, dest_path: Path, dest_branch: str): + """Hold a persistent per-destination flock for the whole integration.""" + if not common_dir.is_dir(): + raise WorktreeError(f"destination git directory missing: {common_dir}") + digest = hashlib.sha256(f"{dest_path}\0{dest_branch}".encode()).hexdigest() + lock = (common_dir / f"lecode-integrate-{digest}.lock").open("a+") + try: + await asyncio.to_thread(fcntl.flock, lock.fileno(), fcntl.LOCK_EX) + yield + finally: + await asyncio.to_thread(fcntl.flock, lock.fileno(), fcntl.LOCK_UN) + lock.close() + + async def _is_ancestor(self, ancestor: str, descendant: str, *, cwd: Path) -> bool: + result = await run_proc( + ["git", "merge-base", "--is-ancestor", ancestor, descendant], + cwd=cwd, + timeout=GIT_TIMEOUT_S, + ) + return result.exit_code == 0 + + async def integrate( + self, + name: str, + *, + validation: list[str], + validation_runner: Callable[[str, Path], Any] | None = None, + allow_unvalidated: bool = False, + reviewed_head: str | None = None, + ) -> IntegrationResult: + """Validate a pinned candidate, then fast-forward its destination only.""" + if not validation and not allow_unvalidated: + return IntegrationResult("blocked", "validation required", [], "") + if validation and validation_runner is None: + raise WorktreeError("validation runner required") + if any(not isinstance(cmd, str) or not cmd for cmd in validation): + raise WorktreeError("validation commands must be nonempty strings") + + info = self._require(name) + sidecar = self._validated_sidecar(name, info) + if sidecar is None: + return IntegrationResult("paused", "no destination recorded", [], "") + dest_path = Path(sidecar["dest_path"]) + dest_branch = sidecar["dest_branch"] + common_dir = Path(sidecar["dest_common_dir"]) + if common_dir != await self._common_dir(self.repo_root): + raise WorktreeError("worker sidecar belongs to a different repository") + + async with self._integration_lock(common_dir, dest_path, dest_branch): + # Re-read after acquiring the stable lock so a replaced sidecar cannot retarget us. + if self._validated_sidecar(name, info) != sidecar: + raise WorktreeError(f"sidecar changed while integrating '{name}'") + problem = await self._destination_problem(dest_path, dest_branch, common_dir) + if problem is not None: + return IntegrationResult("paused", problem, [], "") + worker_head, worker_dirty = await self._worker_state(info, common_dir) + if worker_dirty: + return IntegrationResult("blocked", "uncommitted changes", [], "") + if reviewed_head is not None and worker_head != reviewed_head: + raise WorktreeError("worker head differs from reviewed head") + base_commit = await self._commit(sidecar["base_commit"], cwd=info.path) + if not await self._is_ancestor(base_commit, worker_head, cwd=info.path): + raise WorktreeError("worker branch no longer descends from its pinned base") + dest_head = await self._commit("HEAD", cwd=dest_path) + + if not await self._is_ancestor(dest_head, worker_head, cwd=info.path): + merge = await run_proc( + ["git", "merge", "--no-edit", dest_head], cwd=info.path, timeout=GIT_TIMEOUT_S + ) + if merge.exit_code != 0: + conflicts = await self._conflicted_files(cwd=info.path) + if conflicts: + return IntegrationResult( + "conflict", + f"conflicts merging {dest_branch} into {info.branch}", + conflicts, + "", + ) + return IntegrationResult( + "blocked", merge.stderr.strip() or "merge failed", [], "" + ) + candidate, worker_dirty = await self._worker_state(info, common_dir) + if worker_dirty: + return IntegrationResult("blocked", "uncommitted changes", [], "") + + validation_output = "" + for cmd in validation: + result = validation_runner(cmd, info.path) # type: ignore[misc] + if inspect.isawaitable(result): + result = await result + try: + exit_code, output = result + except (TypeError, ValueError) as error: + raise WorktreeError( + "validation runner must return (exit_code, output)" + ) from error + if not isinstance(exit_code, int) or not isinstance(output, str): + raise WorktreeError("validation runner must return (int, str)") + validation_output += output + if exit_code != 0: + return IntegrationResult( + "validation_failed", + f"validation failed: {cmd}", + [], + _clip_validation(validation_output), + ) + validation_output = _clip_validation(validation_output) + + problem = await self._destination_problem(dest_path, dest_branch, common_dir) + if problem is not None: + return IntegrationResult("paused", problem, [], validation_output) + if await self._commit("HEAD", cwd=dest_path) != dest_head: + return IntegrationResult( + "paused", "destination changed during validation", [], validation_output + ) + current_worker, worker_dirty = await self._worker_state(info, common_dir) + if worker_dirty or current_worker != candidate: + return IntegrationResult( + "blocked", "worker changed during validation", [], validation_output + ) + if not await self._is_ancestor(dest_head, candidate, cwd=info.path): + raise WorktreeError("validated candidate does not contain the pinned destination") + + merge = await run_proc( + ["git", "merge", "--ff-only", candidate], cwd=dest_path, timeout=GIT_TIMEOUT_S + ) + if merge.exit_code != 0: + return IntegrationResult( + "blocked", merge.stderr.strip() or "fast-forward failed", [], validation_output + ) + integrated_head = await self._commit("HEAD", cwd=dest_path) + if integrated_head != candidate: + return IntegrationResult( + "blocked", "destination changed while fast-forwarding", [], validation_output + ) + sidecar["integrated_at"] = datetime.now(UTC).isoformat() + sidecar["integrated_head"] = integrated_head + self.write_sidecar(name, sidecar) + return IntegrationResult( + "integrated", f"fast-forwarded {dest_branch} to {candidate}", [], validation_output + ) + + async def cleanup_worker(self, name: str, *, discard: bool = False) -> WorktreeInfo: + """Remove an integrated worker, or explicitly discard one.""" + info = self._require(name) + if discard: + await self._git("worktree", "remove", "--force", str(info.path)) + await self._git("branch", "-D", info.branch) + return info + + sidecar = self._validated_sidecar(name, info) + if sidecar is None or sidecar.get("integrated_at") is None: + raise WorktreeError( + f"worktree '{name}' is not integrated (use discard=True to discard)" + ) + dest_path = Path(sidecar["dest_path"]) + dest_branch = sidecar["dest_branch"] + common_dir = Path(sidecar["dest_common_dir"]) + problem = await self._destination_problem(dest_path, dest_branch, common_dir) + if problem is not None: + raise WorktreeError(problem) + _worker_head, worker_dirty = await self._worker_state(info, common_dir) + if worker_dirty: + raise WorktreeError( + f"worktree '{name}' has uncommitted changes (use discard=True to discard)" + ) + dest_head = await self._commit("HEAD", cwd=dest_path) + if not await self._is_ancestor(info.branch, dest_head, cwd=dest_path): + raise WorktreeError( + f"worktree '{name}' is not fully integrated (use discard=True to discard)" + ) + await self._git("worktree", "remove", str(info.path)) + await self._git("branch", "-d", info.branch, cwd=dest_path) + return info + async def exit_worktree( self, name: str, *, delete_branch: bool = False, force: bool = False ) -> WorktreeInfo: diff --git a/src/lecode/permission/checker.py b/src/lecode/permission/checker.py index f3850cd..eba8892 100644 --- a/src/lecode/permission/checker.py +++ b/src/lecode/permission/checker.py @@ -1,18 +1,14 @@ """The permission checker: pure, synchronous, no I/O. -Decision pipeline in :meth:`PermissionChecker.check`: - -1. extract the match target (:func:`~lecode.permission.patterns.target_of`) -2. deny rules (global + overlay extras) are unbypassable → Deny -3. overlay ``denied_tools`` → Deny -4. overlay extra allow/ask rules, last match wins -5. global allow rules then ask rules, last match wins -6. session "allow always" grants → Allow -7. mode fallback (two modes: ``yolo`` allows everything, ``readonly`` allows - read-class tools only; MCP read-equivalence makes Exa/context7/grep.app - tools read-class in both modes) -8. doom-loop escalation wraps the result: 3rd identical consecutive call - turns Allow into Ask with a coach reason; 4th+ is Deny +Root policy evaluates deny rules, explicit read-only enforcement, allow/ask +rules, session grants, then mode fallback. Overlay rules and fallback are +intersected with that result and every ancestor policy: Deny > Ask > Allow. +Path rules match the original target and normalized absolute/relative paths +in the executing checker's cwd, including when evaluating ancestor policy. + +Only the executing checker's doom tracker records the call: the 3rd identical +consecutive call turns Allow into Ask; the 4th+ is Deny. Agent views share a +tracker and grants; children have independent trackers and grant stores. The seam for lifecycle hooks (Phase 7): hooks receive the returned :class:`CheckResult` and may only narrow it (Allow → Ask/Deny, Ask → Deny), @@ -21,6 +17,7 @@ from __future__ import annotations +import os.path from dataclasses import dataclass, field from enum import StrEnum from pathlib import Path @@ -99,8 +96,9 @@ class Deny: class AgentOverlay: """Per-agent permission narrowing layered over the global config. - ``mode`` replaces the fallback mode; ``extra_rules`` are evaluated before - global rules; ``denied_tools`` always Deny. Overlays can only narrow. + ``mode`` replaces the overlay layer's fallback; ``extra_rules`` precede + global rules within that layer; ``denied_tools`` always Deny. The result + is capped by the base and ancestor policies, so overlays can only narrow. """ mode: PermissionMode | None = None @@ -155,6 +153,8 @@ def __init__( cwd: Path | None = None, _overlay: AgentOverlay | None = None, _doom: _DoomTracker | None = None, + *, + read_only: bool = False, ) -> None: self._config = config self._rules = config.permissions.rules @@ -163,21 +163,76 @@ def __init__( self._cwd = Path(cwd) if cwd is not None else Path.cwd() self._overlay = _overlay self._doom = _doom or _DoomTracker() + self._read_only = read_only + self._parent: PermissionChecker | None = None @property def mode(self) -> PermissionMode: - """The mode used by the fallback step.""" + """The narrowest fallback mode in this policy's ancestry.""" + if (self._overlay is not None and self._overlay.mode == "readonly") or ( + self._parent is not None and self._parent.mode == "readonly" + ): + return "readonly" return self._mode + @property + def read_only(self) -> bool: + """Whether policy guarantees denial of non read-class tools. + + A readonly fallback with writable rule/grant exceptions is not a safe + shared-checkout policy. Conservatively treat Ask as writable too. + """ + if self._read_only or (self._parent is not None and self._parent.read_only): + return True + mode = self._overlay.mode if self._overlay and self._overlay.mode else self._mode + if mode != "readonly": + return False + rules = [self._rules] + if self._overlay is not None: + rules.append(self._overlay.extra_rules) + return not ( + any( + entries and not self._is_read_class(tool) + for ruleset in rules + for table in (ruleset.allow, ruleset.ask) + for tool, entries in table.items() + ) + or any(not self._is_read_class(tool) for tool, _ in self._session.grants) + ) + def set_mode(self, mode: PermissionMode) -> None: """Switch the fallback mode (``/permissions``); overlays still win.""" self._mode = mode - def for_agent(self, overlay: AgentOverlay) -> PermissionChecker: - """A derived checker with the overlay applied (shares doom tracking).""" - return PermissionChecker( - self._config, self._session, self._mode, self._cwd, overlay, self._doom + def for_agent(self, overlay: AgentOverlay, *, read_only: bool = False) -> PermissionChecker: + """Narrow this policy for an agent, sharing session grants and doom tracking.""" + derived = self.for_child(overlay, session_perms=self._session, read_only=read_only) + derived._doom = self._doom + return derived + + def for_child( + self, + overlay: AgentOverlay | None = None, + *, + cwd: Path | None = None, + session_perms: SessionPermissions | None = None, + read_only: bool = False, + ) -> PermissionChecker: + """Derive an isolated child capped by every ancestor's full policy. + + Uses the supplied child grants (fresh empty grants by default) and a + fresh doom tracker. Ancestor evaluation never records a parent call. + """ + derived = PermissionChecker( + self._config, + session_perms, + self._mode, + self._cwd if cwd is None else cwd, + overlay, + read_only=self._read_only or read_only, ) + derived._parent = self + return derived def check( self, @@ -186,20 +241,38 @@ def check( agent_overlay: AgentOverlay | None = None, ) -> CheckResult: """Decide Allow/Ask/Deny for one tool call. Pure; prompting is the TUI's job.""" - overlay = agent_overlay or self._overlay + if agent_overlay is not None: + return self.for_agent(agent_overlay).check(tool_name, args) target = target_of(tool_name, args) - result = self._base_decision(tool_name, target, overlay) + targets = (target,) + if target and ( + tool_name in ("read", "write", "edit", "list_dir") + or (tool_name in ("grep", "find_files") and args.get("path")) + ): + absolute = os.path.abspath(self._cwd / target) + targets = (target, absolute, os.path.relpath(absolute, self._cwd)) + result = self._policy_decision(tool_name, targets) return self._apply_doom_loop(tool_name, args, result) # -- pipeline steps ------------------------------------------------------ + def _policy_decision(self, tool_name: str, targets: tuple[str, ...]) -> CheckResult: + """Intersect every policy layer without recording an ancestor call.""" + results = [self._base_decision(tool_name, targets, None)] + if self._overlay is not None: + results.append(self._base_decision(tool_name, targets, self._overlay)) + if self._parent is not None: + results.append(self._parent._policy_decision(tool_name, targets)) + priority = {Decision.ALLOW: 0, Decision.ASK: 1, Decision.DENY: 2} + return max(results, key=lambda result: priority[result.decision]) + def _base_decision( - self, tool_name: str, target: str, overlay: AgentOverlay | None + self, tool_name: str, targets: tuple[str, ...], overlay: AgentOverlay | None ) -> CheckResult: # Deny rules are unbypassable: global table + overlay extras. - deny = self._last_match(self._rules.deny, tool_name, target) + deny = self._last_match(self._rules.deny, tool_name, targets) if overlay is not None: - deny = self._last_match(overlay.extra_rules.deny, tool_name, target) or deny + deny = self._last_match(overlay.extra_rules.deny, tool_name, targets) or deny if deny is not None: return CheckResult(Decision.DENY, f"deny rule matched: {deny.pattern}", deny) @@ -207,44 +280,50 @@ def _base_decision( if overlay is not None and tool_name in overlay.denied_tools: return CheckResult(Decision.DENY, f"tool denied by agent overlay: {tool_name}") + # Read-only checkers can never be widened by overlay rules, global + # rules, session grants, or the mode fallback. + if self._read_only and not self._is_read_class(tool_name): + return CheckResult(Decision.DENY, f"read-only: {tool_name} is not a read-class tool") + # Overlay extra allow/ask rules first (last match wins within them). if overlay is not None: - extra = self._last_match_allow_ask(overlay.extra_rules, tool_name, target) + extra = self._last_match_allow_ask(overlay.extra_rules, tool_name, targets) if extra is not None: return extra # Global rules: allow table walked first, then ask table; last match wins. - matched = self._last_match_allow_ask(self._rules, tool_name, target) + matched = self._last_match_allow_ask(self._rules, tool_name, targets) if matched is not None: return matched # Session "allow always" grants (can never reach a deny: handled above). - grant = self._session.matching_grant(tool_name, target) - if grant is not None: - return CheckResult(Decision.ALLOW, f"session grant: {grant}") + for target in targets: + grant = self._session.matching_grant(tool_name, target) + if grant is not None: + return CheckResult(Decision.ALLOW, f"session grant: {grant}") # Mode fallback. mode = overlay.mode if overlay and overlay.mode else self._mode - return self._mode_fallback(mode, tool_name, target) + return self._mode_fallback(mode, tool_name) def _last_match( - self, table: dict[str, list[PermissionRule]], tool_name: str, target: str + self, table: dict[str, list[PermissionRule]], tool_name: str, targets: tuple[str, ...] ) -> PermissionRule | None: match = None for rule in table.get(tool_name, []): - if rule_matches(rule, target): + if any(rule_matches(rule, target) for target in targets): match = rule return match def _last_match_allow_ask( - self, rules: PermissionRuleSet, tool_name: str, target: str + self, rules: PermissionRuleSet, tool_name: str, targets: tuple[str, ...] ) -> CheckResult | None: result: CheckResult | None = None for rule in rules.allow.get(tool_name, []): - if rule_matches(rule, target): + if any(rule_matches(rule, target) for target in targets): result = CheckResult(Decision.ALLOW, f"allow rule matched: {rule.pattern}", rule) for rule in rules.ask.get(tool_name, []): - if rule_matches(rule, target): + if any(rule_matches(rule, target) for target in targets): result = CheckResult(Decision.ASK, f"ask rule matched: {rule.pattern}", rule) return result @@ -269,7 +348,7 @@ def _apply_doom_loop( def _is_read_class(self, tool_name: str) -> bool: return tool_name in READ_TOOLS or is_read_equiv_mcp(tool_name) - def _mode_fallback(self, mode: PermissionMode, tool_name: str, target: str) -> CheckResult: + def _mode_fallback(self, mode: PermissionMode, tool_name: str) -> CheckResult: reason = f"mode: {mode}" if mode == "yolo": return CheckResult(Decision.ALLOW, reason) diff --git a/src/lecode/session/stats.py b/src/lecode/session/stats.py index 0b05e9b..4494da2 100644 --- a/src/lecode/session/stats.py +++ b/src/lecode/session/stats.py @@ -26,6 +26,8 @@ class Stats: created_at: str last_active: str | None tombstone_count: int + #: True when a counted worker usage event is marked incomplete. + usage_incomplete: bool = False def _usage_tokens(usage: dict) -> tuple[int, int]: @@ -44,6 +46,7 @@ def session_stats(store: SessionStore, session: Session, catalog: Catalog | None input_tokens = 0 output_tokens = 0 cost_usd = 0.0 + usage_incomplete = False for record in messages: role_counts[record.role] = role_counts.get(record.role, 0) + 1 @@ -62,14 +65,22 @@ def session_stats(store: SessionStore, session: Session, catalog: Catalog | None continue cost_usd += (in_tok * pricing.prompt + out_tok * pricing.completion) / 1_000_000 - # Pierre reviews carry their own usage on the event record. + # Pierre reviews and worker dispatches carry their own usage on the event. for record in records: - if isinstance(record, EventRecord) and record.kind == "pierre": + if not isinstance(record, EventRecord): + continue + if record.kind == "pierre": usage = record.data.get("usage") or {} - in_tok, out_tok = _usage_tokens(usage) - input_tokens += in_tok - output_tokens += out_tok - cost_usd += float(usage.get("cost_usd") or 0.0) + elif record.kind == "worker_usage": + usage = record.data.get("usage") or record.data + if record.data.get("incomplete"): + usage_incomplete = True + else: + continue + in_tok, out_tok = _usage_tokens(usage) + input_tokens += in_tok + output_tokens += out_tok + cost_usd += float(usage.get("cost_usd") or 0.0) timestamps = [ r.ts for r in records if isinstance(r, MessageRecord | EventRecord | TombstoneRecord) @@ -93,4 +104,5 @@ def session_stats(store: SessionStore, session: Session, catalog: Catalog | None created_at=session.meta.created_at, last_active=max(timestamps) if timestamps else None, tombstone_count=sum(isinstance(r, TombstoneRecord) for r in records), + usage_incomplete=usage_incomplete, ) diff --git a/src/lecode/session/storage.py b/src/lecode/session/storage.py index 4ae942b..8d360f9 100644 --- a/src/lecode/session/storage.py +++ b/src/lecode/session/storage.py @@ -440,18 +440,22 @@ def load_messages(self, session: Session) -> list[MessageRecord]: if isinstance(r, MessageRecord) and not self._is_hidden(r.seq, tombstones) ] - def load_agent_runs(self, session: Session) -> list[dict[str, Any]]: - """Agent-run activity records with tombstones applied, in append order.""" + def load_events(self, session: Session, kind: str) -> list[dict[str, Any]]: + """Event records of ``kind`` with tombstones applied, in append order.""" records = self.read_records(session) tombstones = self._active_tombstones(records) return [ dict(r.data) for r in records if isinstance(r, EventRecord) - and r.kind == "agent_run" + and r.kind == kind and not self._is_hidden(r.seq, tombstones) ] + def load_agent_runs(self, session: Session) -> list[dict[str, Any]]: + """Agent-run activity records with tombstones applied, in append order.""" + return self.load_events(session, "agent_run") + def undo(self, session: Session) -> TombstoneRecord | None: """Hide the last user turn (the user message and everything after it).""" visible = self.load_messages(session) diff --git a/src/lecode/slash/catalog.py b/src/lecode/slash/catalog.py index adfe285..05251c4 100644 --- a/src/lecode/slash/catalog.py +++ b/src/lecode/slash/catalog.py @@ -45,6 +45,7 @@ ("queue", "Show queued/steered messages"), ("tasks", "List background tasks"), ("runs", "Inspect agent runs"), + ("agent", "Control a worker"), ("copy", "Copy the last answer to the clipboard"), ("export", "Export the session as HTML"), ("import", "Import a session file"), diff --git a/src/lecode/slash/handlers.py b/src/lecode/slash/handlers.py index 75579ec..12be7fb 100644 --- a/src/lecode/slash/handlers.py +++ b/src/lecode/slash/handlers.py @@ -20,6 +20,7 @@ from lecode.extras.loop_mode import DEFAULT_MAX_ITERATIONS from lecode.extras.proc import run_proc from lecode.extras.status_signals import GIT_CONFLICT +from lecode.extras.subagents import SubagentError from lecode.extras.worktree import WorktreeError, WorktreeManager from lecode.hooks import hooks_status from lecode.memory import MemoryStore, memory_command, memory_root @@ -204,6 +205,7 @@ async def cmd_session(app: TuiApp, args: list[str]) -> None: """``/session``: metadata plus token/cost stats.""" session = app.session stats = session_stats(app.store, session) + input_tokens, output_tokens, cost_usd, incomplete = app.current_session_usage() roles = " · ".join(f"{role} x{count}" for role, count in sorted(stats.role_counts.items())) lines = [ f"session: {session.name} ({session.id})", @@ -213,7 +215,8 @@ async def cmd_session(app: TuiApp, args: list[str]) -> None: f"agent: {session.meta.agent} · model: {session.meta.model or app.config.llm.model}", f"messages: {stats.message_count} ({roles or 'none'})" f" · tombstones: {stats.tombstone_count}", - f"tokens: {stats.input_tokens} in / {stats.output_tokens} out · cost ${stats.cost_usd:.4f}", + f"tokens: {input_tokens} in / {output_tokens} out · cost ${cost_usd:.4f}" + + (" (incomplete)" if incomplete else ""), ] app.feed.info("\n".join(lines)) @@ -759,6 +762,60 @@ async def cmd_runs(app: TuiApp, args: list[str]) -> None: app.open_agent_run(run.run_id) +async def cmd_agent(app: TuiApp, args: list[str]) -> None: + """``/agent [send|stop|resume|submit|focus]``: human worker control.""" + manager = app.worker_manager + if manager is None: + app.feed.info("(no workers this session)") + return + if not args: + workers = [run for run in app.roster.runs() if run.worker] + if not workers: + app.feed.info("(no workers this session)") + return + app.feed.info( + "\n".join( + f"{run.index}. {run.agent} · {run.status} · {run.description} · {run.run_id}" + for run in workers + ) + + "\n\ncontrol: /agent [send TEXT|stop [tree]|resume [TEXT]|submit|focus]" + ) + return + worker = app.resolve_worker(args[0]) + if worker is None: + app.feed.error(f"no such worker: {args[0]}") + return + if len(args) == 1: + app.open_agent_run(worker.id) + return + action = args[1] + rest = args[2:] + try: + if action == "send": + if not rest: + raise ValueError("usage: /agent send ") + await manager.send(worker.id, " ".join(rest), from_human=True) + app.feed.info(f"sent to @{worker.agent} ({worker.id[:8]})") + elif action == "stop": + await manager.stop(worker.id, tree=bool(rest and rest[0] == "tree")) + app.feed.info(f"worker {worker.id[:8]} stopped") + elif action == "resume": + await manager.resume(worker.id, " ".join(rest) if rest else None) + app.feed.info(f"worker {worker.id[:8]} resumed") + elif action == "submit": + await app.submit_worker(worker.id) + app.feed.info(f"worker {worker.id[:8]} submitted") + elif action == "focus": + app.focus_worker(worker.id) + app.feed.info(f"composer focused on @{worker.agent} ({worker.id[:8]}); Esc returns") + else: + raise ValueError( + "usage: /agent [send TEXT|stop [tree]|resume [TEXT]|submit|focus]" + ) + except (KeyError, RuntimeError, SubagentError, ValueError) as e: + app.feed.error(str(e)) + + async def cmd_tasks(app: TuiApp, args: list[str]) -> None: """``/tasks``: background tasks (id, kind, status, age, description).""" manager = app.runtime.ctx.extras.get(BACKGROUND_EXTRA) @@ -1418,6 +1475,23 @@ def _complete_runs(app: TuiApp, args: list[str]) -> list[CompletionRow]: ] +def _complete_agent(app: TuiApp, args: list[str]) -> list[CompletionRow]: + if not args: + return [ + (run.run_id, f"{run.index} {run.agent} · {run.description}", run.status) + for run in app.roster.runs() + if run.worker + ] + if len(args) == 1: + return [ + (action, action, "worker control") + for action in ("send", "stop", "resume", "submit", "focus") + ] + if len(args) == 2 and args[1] == "stop": + return [("tree", "tree", "stop descendants too")] + return [] + + def _complete_rewind(app: TuiApp, args: list[str]) -> list[CompletionRow]: if args: return [] @@ -1480,6 +1554,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "memory": (_complete_memory, None), "wt-exit": (_complete_wt_exit, None), "runs": (_complete_runs, "no agent runs this session"), + "agent": (_complete_agent, "no workers this session"), } @@ -1520,7 +1595,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: ), ("Permissions", ["permissions", "mode", "toggle"]), ("Worktrees", ["worktree", "wt-merge", "wt-exit"]), - ("Power features", ["loop", "chain", "mcp", "review", "tasks", "runs"]), + ("Power features", ["loop", "chain", "mcp", "review", "tasks", "runs", "agent"]), ( "Interface", [ @@ -1581,6 +1656,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "agents": cmd_agents, "queue": cmd_queue, "runs": cmd_runs, + "agent": cmd_agent, "tasks": cmd_tasks, "btw": cmd_btw, "copy": cmd_copy, @@ -1638,6 +1714,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "review": "[file…]", "notifications": "[on|off]", "runs": "[number|id]", + "agent": "[id|number] [send TEXT|stop [tree]|resume [TEXT]|submit|focus]", } diff --git a/src/lecode/tui/agents.py b/src/lecode/tui/agents.py index f72eb7b..60217cc 100644 --- a/src/lecode/tui/agents.py +++ b/src/lecode/tui/agents.py @@ -9,13 +9,16 @@ import json import time +from contextlib import suppress from dataclasses import dataclass, field +from datetime import UTC, datetime from rich.text import Text from lecode.agent.runner import Done, Error, ToolCall, ToolResult from lecode.extras.subagents import SubagentProgress from lecode.permission.patterns import target_of +from lecode.tui.statusline import format_cost, human_tokens from lecode.tui.themes import Theme #: Rows the compact roster shows before collapsing into ``+N more``. @@ -32,9 +35,15 @@ _STATUS_GLYPHS = { "running": ("●", "accent"), + "queued": ("○", "muted"), + "waiting": ("◌", "warning"), + "done": ("✔", "success"), "ok": ("✔", "success"), "error": ("✗", "error"), + "failed": ("✗", "error"), "cancelled": ("—", "muted"), + "stopped": ("■", "muted"), + "interrupted": ("!", "warning"), } @@ -84,6 +93,17 @@ class AgentRun: error: str = "" started_at: float = field(default_factory=time.monotonic) truncated: bool = False + parent_id: str | None = None + depth: int = 0 + worker: bool = False + origin: str = "delegated" + cost_usd: float = 0.0 + usage_incomplete: bool = False + context_used: int = 0 + context_window: int = 0 + elapsed_s: float = 0.0 + subtree_cost_usd: float = 0.0 + subtree_usage_incomplete: bool = False @property def current(self) -> str: @@ -138,6 +158,56 @@ def observe(self, progress: SubagentProgress) -> AgentRun: run.status = "ok" return run + def sync_worker(self, worker, *, context_window: int = 0) -> AgentRun: + """Mirror the persistent worker state into the UI-only roster.""" + run = self._runs.get(worker.id) + if run is None: + run = AgentRun( + run_id=worker.id, + index=len(self._order) + 1, + agent=worker.agent, + description=worker.description, + parent_id=worker.parent_id, + depth=worker.depth, + worker=True, + origin=worker.origin, + ) + self._runs[run.run_id] = run + self._order.append(run.run_id) + run.status = {"completed": "done", "failed": "error"}.get(worker.state, worker.state) + run.error = worker.error or "" + run.answer = worker.result.final_text if worker.result is not None else "" + run.cost_usd = worker.usage_totals.cost_usd + run.usage_incomplete = worker.usage_incomplete + run.context_used = worker.usage_totals.context_tokens + run.context_window = context_window + started_at = getattr(worker, "started_at", "") + now = datetime.now(UTC) + with suppress(TypeError, ValueError): + run.elapsed_s = max(0.0, (now - datetime.fromisoformat(started_at)).total_seconds()) + self._refresh_subtree_costs() + return run + + def _refresh_subtree_costs(self) -> None: + workers = [run for run in self._runs.values() if run.worker] + children: dict[str | None, list[AgentRun]] = {} + for run in workers: + children.setdefault(run.parent_id, []).append(run) + + def total(run: AgentRun) -> tuple[float, bool]: + cost = run.cost_usd + incomplete = run.usage_incomplete + for child in children.get(run.run_id, []): + child_cost, child_incomplete = total(child) + cost += child_cost + incomplete |= child_incomplete + run.subtree_cost_usd = cost + run.subtree_usage_incomplete = incomplete + return cost, incomplete + + for run in children.get(None, []): + total(run) + def finish( self, run_id: str, *, answer: str = "", error: str = "", is_error: bool = False ) -> AgentRun | None: @@ -163,7 +233,14 @@ def cancel_running(self) -> list[AgentRun]: return cancelled def has_running(self) -> bool: - return any(run.status == "running" for run in self._runs.values()) + return any(run.status in {"queued", "running", "waiting"} for run in self._runs.values()) + + def has_workers(self) -> bool: + return any(run.worker for run in self._runs.values()) + + def worker_total(self) -> tuple[float, bool]: + workers = [run for run in self._runs.values() if run.worker] + return sum(run.cost_usd for run in workers), any(run.usage_incomplete for run in workers) def get(self, run_id: str) -> AgentRun | None: return self._runs.get(run_id) @@ -187,8 +264,25 @@ def visible(self, limit: int = ROSTER_VISIBLE_ROWS) -> tuple[list[AgentRun], int """``(shown, hidden_count)``: running first (newest first), then finished (newest first); stable indices come from the run itself.""" runs = self.runs() - running = [run for run in reversed(runs) if run.status == "running"] - finished = [run for run in reversed(runs) if run.status != "running"] + if self.has_workers(): + by_parent: dict[str | None, list[AgentRun]] = {} + workers = [run for run in runs if run.worker] + for run in workers: + by_parent.setdefault(run.parent_id, []).append(run) + ordered: list[AgentRun] = [] + + def visit(parent_id: str | None) -> None: + for child in by_parent.get(parent_id, []): + ordered.append(child) + visit(child.run_id) + + visit(None) + ordered.extend(run for run in runs if not run.worker) + return ordered[:limit], max(0, len(ordered) - limit) + running = [run for run in reversed(runs) if run.status in {"queued", "running", "waiting"}] + finished = [ + run for run in reversed(runs) if run.status not in {"queued", "running", "waiting"} + ] ordered = running + finished return ordered[:limit], max(0, len(ordered) - limit) @@ -203,20 +297,30 @@ def roster_lines(roster: AgentRoster, theme: Theme, width: int) -> list[Text]: runs = roster.runs() if not runs: return [] - running = sum(1 for run in runs if run.status == "running") + running = sum(1 for run in runs if run.status in {"queued", "running", "waiting"}) done = len(runs) - running - header = f"agents · {running} running" + header = ( + f"workers · {running} running" if roster.has_workers() else f"agents · {running} running" + ) if done: header += f" · {done} done" + if roster.has_workers(): + cost, incomplete = roster.worker_total() + header += f" · {format_cost(cost)}" + (" incomplete" if incomplete else "") lines = [Text(header, style=theme.muted)] shown, hidden = roster.visible() for run in shown: glyph, style = _glyph(run.status, theme) line = Text() - line.append(f" {glyph} ", style=style) + indent = " " * run.depth if run.worker else "" + line.append(f" {indent}{glyph} ", style=style) line.append(f"{run.index} {run.agent} ", style=theme.accent) line.append(_clip(run.description, max(12, width // 3)), style=theme.text) - line.append(f" · {_clip(run.current, max(12, width // 3))}", style=theme.muted) + detail = _clip(run.current, max(12, width // 3)) + if run.worker: + detail += f" · {format_cost(run.cost_usd)}" + ("?" if run.usage_incomplete else "") + detail += f" · {run.elapsed_s:.0f}s" + line.append(f" · {detail}", style=theme.muted) lines.append(line) if hidden: lines.append(Text(f" … +{hidden} more · /runs", style=theme.muted)) @@ -243,6 +347,21 @@ def detail_lines(run: AgentRun | None, theme: Theme, width: int) -> list[Text]: header.append(_clip(run.description, max(12, width // 2)), style=theme.text) header.append(f" · {run.current}", style=theme.muted) lines.append(header) + if run.worker: + cost = format_cost(run.cost_usd) + (" (incomplete)" if run.usage_incomplete else "") + subtree = format_cost(run.subtree_cost_usd) + ( + " (incomplete)" if run.subtree_usage_incomplete else "" + ) + context = "unknown" + if run.context_window: + context = f"{human_tokens(run.context_used)}/{human_tokens(run.context_window)}" + lines.append( + Text( + f" cost: {cost} · subtree: {subtree} · ctx: {context}" + f" · elapsed: {run.elapsed_s:.0f}s", + style=theme.muted, + ) + ) if not run.activity: lines.append(Text(" (no tool calls yet)", style=theme.muted)) for entry in run.activity: diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index ad151bd..080ffd4 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -71,6 +71,8 @@ SubagentProgress, run_subagent, ) +from lecode.extras.workers import WORKER_EXTRA +from lecode.extras.worktree import WorktreeError from lecode.hooks import ( INTERRUPT, NOTIFICATION, @@ -263,6 +265,12 @@ def __init__( self._runtime.ctx.extras["subagent_events"] = self._on_child_event #: Live agent-run roster (compact above the composer; /runs opens detail). self._roster = AgentRoster() + #: A focused worker shares this composer; it never opens another TUI. + self._focused_worker_id: str | None = None + self._composer_drafts: dict[str, str] = {} + self._root_stopped = False + self._worker_wake_task: asyncio.Task[None] | None = None + self._worker_usage = (0, 0, 0.0) self._detail_run_id: str | None = None self._roster_window: Any | None = None # Background-task completions surface as feed info lines as they land. @@ -290,6 +298,11 @@ def __init__( self._status.output_tokens = stats.output_tokens self._status.cost_usd = stats.cost_usd self._status.context_used = stats.context_tokens + self._worker_manager = runtime.ctx.extras.get(WORKER_EXTRA) + if self._worker_manager is not None: + self._worker_manager.notify = self._on_worker_notification + self._worker_manager.progress = self._on_worker_progress + self._hydrate_workers() # Logbook suffix: the feed reads live context/cost from the statusline. self._feed.metrics = lambda: ( self._status.context_used, @@ -392,6 +405,7 @@ def catalog(self) -> Any: def submit_prompt(self, text: str, *, echo: str | None = None) -> None: """Render and queue/start a user prompt (skill commands, ``/retry``).""" + self._root_stopped = False self._feed.user_message(text if echo is None else echo) self._enqueue_or_start(text) @@ -421,6 +435,9 @@ def switch_session(self, session: Session) -> bool: """ from lecode.session.storage import SessionInUseError + if self.workers_active(): + self._feed.error("stop or finish all workers before switching sessions") + return False try: new_lock = self._store.acquire_lock(session) except SessionInUseError as e: @@ -450,6 +467,14 @@ def switch_session(self, session: Session) -> bool: self._input_area.history = self._input_history self._status.session_name = session.name self._status.agent = self._agent_name + if self._worker_manager is not None: + self._worker_manager.attach(session) + self._roster = AgentRoster() + self._detail_run_id = None + self._focused_worker_id = None + self._composer_drafts.clear() + self._worker_usage = (0, 0, 0.0) + self._hydrate_workers() self._invalidate() hooks = self._runtime.hooks if hooks is not None and ( @@ -475,7 +500,115 @@ def reload_history(self) -> None: def turn_busy(self) -> bool: """Whether a turn or plan loop is in flight (switching commands refuse).""" - return self._turn_running() or self.loop_running() + return self._turn_running() or self.loop_running() or self.workers_active() + + def workers_active(self) -> bool: + return self._worker_manager is not None and any( + worker.state in {"queued", "running", "waiting"} + for worker in self._worker_manager.list() + ) + + def resolve_worker(self, ref: str) -> Any | None: + """Resolve a worker by stable roster number, exact id, or unique prefix.""" + if self._worker_manager is None: + return None + run = self._roster.resolve(ref) + if run is not None and run.worker: + return self._worker_manager.get(run.run_id) + matches = [worker for worker in self._worker_manager.list() if worker.id.startswith(ref)] + return matches[0] if len(matches) == 1 else None + + def focus_worker(self, worker_id: str | None) -> bool: + """Switch the shared composer to a worker, retaining both drafts.""" + if worker_id is not None and self.resolve_worker(worker_id) is None: + return False + current = self._focused_worker_id or "main" + if self._input_area is not None: + self._composer_drafts[current] = self._input_area.text + self._focused_worker_id = worker_id + if self._input_area is not None: + draft = self._composer_drafts.get(worker_id or "main", "") + self._input_area.buffer.set_document(Document(draft, len(draft)), bypass_readonly=True) + self._sync_queue_status() + return True + + def _hydrate_workers(self) -> None: + """Load persisted workers once per attached root session.""" + if self._worker_manager is None: + return + for worker in self._worker_manager.load(): + self._roster.sync_worker(worker, context_window=self._status.context_window) + self._worker_usage = self._worker_totals() + + def _worker_totals(self) -> tuple[int, int, float]: + if self._worker_manager is None: + return (0, 0, 0.0) + totals = [worker.usage_totals for worker in self._worker_manager.list()] + return ( + sum(total.input_tokens for total in totals), + sum(total.output_tokens for total in totals), + sum(total.cost_usd for total in totals), + ) + + def current_session_usage(self) -> tuple[int, int, float, bool]: + """Session totals with live worker usage substituted for persisted deltas.""" + stats = session_stats(self._store, self._session) + persisted = [ + event.get("usage", event) + for event in self._store.load_events(self._session, "worker_usage") + ] + live = self._worker_totals() + recorded = ( + sum(int(item.get("input_tokens") or 0) for item in persisted), + sum(int(item.get("output_tokens") or 0) for item in persisted), + sum(float(item.get("cost_usd") or 0.0) for item in persisted), + ) + incomplete = ( + stats.usage_incomplete + or any(worker.usage_incomplete for worker in self._worker_manager.list()) + if self._worker_manager is not None + else stats.usage_incomplete + ) + return ( + stats.input_tokens + live[0] - recorded[0], + stats.output_tokens + live[1] - recorded[1], + stats.cost_usd + live[2] - recorded[2], + incomplete, + ) + + def _on_worker_progress(self, worker: Any) -> None: + """WorkerManager callback: update exactly that roster entry and live totals.""" + before = self._worker_usage + self._roster.sync_worker(worker, context_window=self._status.context_window) + after = self._worker_totals() + self._status.input_tokens += after[0] - before[0] + self._status.output_tokens += after[1] - before[1] + self._status.cost_usd += after[2] - before[2] + self._worker_usage = after + self._invalidate() + + async def _on_worker_notification(self, note: dict[str, Any]) -> None: + """Render completed-worker attribution and wake only delegated root work.""" + worker = self.resolve_worker(note["worker_id"]) + if worker is None: + return + run = self._roster.sync_worker(worker, context_window=self._status.context_window) + self._feed.agent_summary(run) + if ( + (note["origin"] == "delegated" or note["deliver"]) + and note["parent_id"] is None + and not self._root_stopped + and not self._turn_running() + and (self._worker_wake_task is None or self._worker_wake_task.done()) + ): + self._worker_wake_task = asyncio.create_task(self._wake_for_workers()) + self._invalidate() + + async def _wake_for_workers(self) -> None: + """Coalesce same-loop worker completions into one parent wake-up.""" + await asyncio.sleep(0) + if not self._root_stopped and not self._turn_running() and not self._quit: + self._enqueue_or_start("Worker updates are available.") def loop_running(self) -> bool: return self._loop_task is not None and not self._loop_task.done() @@ -607,21 +740,21 @@ def question_nav() -> bool: @kb.add("y", filter=approval_pending) def _approve_once(event: Any) -> None: - self._approval.resolve(AllowOnce()) + self._resolve_approval(AllowOnce()) @kb.add("a", filter=approval_pending) def _approve_always(event: Any) -> None: pending = self._approval.pending pattern = pending.target if pending is not None else "*" - self._approval.resolve(AllowAlways(pattern=pattern)) + self._resolve_approval(AllowAlways(pattern=pattern)) @kb.add("n", filter=approval_pending) def _deny(event: Any) -> None: - self._approval.resolve(Deny()) + self._resolve_approval(Deny()) @kb.add("escape", filter=approval_pending) def _deny_escape(event: Any) -> None: - self._approval.resolve(Deny()) + self._resolve_approval(Deny()) @kb.add("up", filter=question_picker) def _question_up(event: Any) -> None: @@ -674,6 +807,10 @@ def _close_completion_menu(event: Any) -> None: # state, so their dismissal is remembered per typed text (any edit # or Tab brings it back). Longer M-* sequences still win over this # bare-key handler. + if self._focused_worker_id is not None: + worker = self.resolve_worker(self._focused_worker_id) + self.focus_worker(worker.parent_id if worker is not None else None) + return if self._detail_run_id is not None: self.close_agent_run() return @@ -735,7 +872,7 @@ def _newline(event: Any) -> None: @kb.add("c-c") def _ctrl_c(event: Any) -> None: if self._approval.is_pending: - self._approval.resolve(Deny()) + self._resolve_approval(Deny()) return if self._question.is_pending: self._question.dismiss() @@ -1089,6 +1226,7 @@ def menu_footer() -> str: async def run(self, *, input: Input | None = None, output: Output | None = None) -> int: """Run the interactive loop until quit; returns the exit code.""" + self._hydrate_workers() self._status.git = await self._git.get(self._cwd) self._app = self._build_app(input=input, output=output) self._runtime.ctx.approval_callback = self._request_approval @@ -1130,6 +1268,8 @@ async def run(self, *, input: Input | None = None, output: Output | None = None) background = self._runtime.ctx.extras.get(BACKGROUND_EXTRA) if background is not None: await background.shutdown() + if self._worker_manager is not None: + await self._worker_manager.shutdown() mcp = self._runtime.ctx.extras.get(MCP_EXTRA) if mcp is not None: await mcp.shutdown() @@ -1222,6 +1362,16 @@ async def _submit(self, text: str, *, steer: bool = False) -> None: elif text.startswith(".") and self._submit_persona(text, steer=steer): pass # .persona : handled (persona system-prompt overlay) else: + if self._focused_worker_id is not None: + worker = self.resolve_worker(self._focused_worker_id) + if worker is None: + self._feed.error("focused worker no longer exists") + self.focus_worker(None) + return + self._feed.user_message(f"[@{worker.agent}] {text}") + await self._worker_manager.send(worker.id, text, steer, from_human=True) + self._feed.info(f"sent to @{worker.agent} ({worker.id[:8]})") + return if self.loop_running(): self._feed.info("a plan loop is running — /loop stop first") return @@ -1235,14 +1385,15 @@ async def _submit(self, text: str, *, steer: bool = False) -> None: invocable = {a.name for a in self._runtime.agents.subagents()} targets = [name for name in mentions if name in invocable] if targets and not self._turn_running(): - # Direct @agent dispatch: a side query run by the subagent. - # While a turn runs, mentions keep the note behavior in - # _prepare_message (the message queues as normal input). self._feed.user_message(text) - self._turn_task = asyncio.ensure_future( - self._run_subagent_turn(targets[0], cleaned or text) - ) + worker = await self._start_human_worker(targets[0], cleaned or text) + if worker is not None: + # Compatibility: direct @agent work remains awaitable through + # the established turn-task seam, but execution stays in the + # persistent human-origin worker manager. + self._turn_task = asyncio.create_task(self._wait_for_worker(worker.id)) return + self._root_stopped = False prepared = self._prepare_message(text) echo = self._attachment_echo() content = self._with_attachments(prepared) @@ -1254,6 +1405,38 @@ async def _submit(self, text: str, *, steer: bool = False) -> None: self._feed.user_message(text + echo) self._enqueue_or_start(content, steer=steer) + async def _start_human_worker(self, agent: str, prompt: str) -> Any | None: + """Start direct ``@agent`` work as a persistent human-origin worker.""" + if self._worker_manager is None: + self._feed.error("workers are unavailable") + return None + try: + worker = await self._worker_manager.start( + self._runtime.ctx, + agent=agent, + prompt=prompt, + description=prompt.splitlines()[0][:80], + origin="human", + background=True, + ) + except (RuntimeError, SubagentError, WorktreeError) as e: + self._feed.error(str(e)) + return None + self._roster.sync_worker(worker, context_window=self._status.context_window) + self._feed.info( + f"worker {worker.id[:8]} started for @{agent} · /agent {worker.id[:8]} focus" + ) + self._invalidate() + return worker + + async def _wait_for_worker(self, worker_id: str) -> None: + """Keep direct worker execution compatible with the legacy turn seam.""" + assert self._worker_manager is not None + try: + await self._worker_manager.wait(worker_id) + except SubagentError as e: + self._feed.error(str(e)) + def _with_attachments(self, text: str) -> MessageContent | None: """Attach pending attachments to ``text``; ``None`` = blocked (modality).""" attachments = self._attachments.list() @@ -1398,15 +1581,43 @@ def _activity(self, label: str | None) -> None: self._status.activity = label self._invalidate() + def _show_approval_head(self) -> None: + """Print the FIFO head once, keeping worker attribution with the ask.""" + pending = self._approval.pending + if pending is None: + return + shown = getattr(self, "_shown_approval", None) + if shown is pending.future: + return + self._shown_approval = pending.future + attribution = ( + f"[worker {pending.worker[:8]} · {pending.conversation}] " if pending.worker else "" + ) + self._feed.permission(attribution + approval_prompt_text(pending.tool_name, pending.target)) + + def _resolve_approval(self, decision: ApprovalDecision) -> None: + self._approval.resolve(decision) + self._shown_approval = None + self._show_approval_head() + self._invalidate() + async def _request_approval( - self, tool_name: str, args: dict[str, Any], reason: str + self, + tool_name: str, + args: dict[str, Any], + reason: str, + *, + worker: str | None = None, + conversation: str = "main", ) -> ApprovalDecision: """``ctx.approval_callback``: inline y/a/n/ESC ask during a turn.""" target = target_of(tool_name, args) if reason: self._feed.info(reason) # doom-loop coach reasons land here too - self._feed.permission(approval_prompt_text(tool_name, target)) - future = self._approval.request(tool_name, target, reason) + future = self._approval.request( + tool_name, target, reason, worker=worker, conversation=conversation + ) + self._show_approval_head() self._status.state = StatusLineState.AWAITING_APPROVAL self._invalidate() self._spawn(self._notifier.approval_needed(tool_name)) @@ -1414,8 +1625,14 @@ async def _request_approval( try: return await future finally: - self._approval.cancel() - self._status.state = StatusLineState.RUNNING + self._approval.cancel(future) + self._shown_approval = None + self._show_approval_head() + self._status.state = ( + StatusLineState.AWAITING_APPROVAL + if self._approval.is_pending + else StatusLineState.RUNNING + ) self._invalidate() def _render_question(self) -> None: @@ -1456,6 +1673,7 @@ def _turn_running(self) -> bool: def cancel_turn(self) -> bool: """Cancel the in-flight turn (Ctrl-C); ``True`` if one was cancelled.""" if self._turn_running(): + self._root_stopped = True self._turn_task.cancel() return True return False @@ -1501,7 +1719,12 @@ def _chatbox_title(self) -> str: parts.append("queue: " + ", ".join(self._queue_label(m) for m in queued)) if steered: parts.append("steer: " + ", ".join(self._queue_label(m) for m in steered)) - return "message" if not parts else "message · " + " · ".join(parts) + recipient = "message" + if self._focused_worker_id is not None: + worker = self.resolve_worker(self._focused_worker_id) + if worker is not None: + recipient = f"to @{worker.agent} ({worker.id[:8]})" + return recipient if not parts else recipient + " · " + " · ".join(parts) async def _run_turn(self, text: MessageContent, *, overlay: str | None = None) -> None: message: dict[str, Any] = {"role": "user", "content": text} @@ -1778,6 +2001,18 @@ def _on_child_event(self, progress: SubagentProgress) -> None: def roster(self) -> AgentRoster: return self._roster + @property + def worker_manager(self) -> Any | None: + return self._worker_manager + + async def submit_worker(self, worker_id: str) -> dict[str, Any]: + """Submit a human worker and apply its root wake-up policy.""" + if self._worker_manager is None: + raise RuntimeError("workers are unavailable") + note = await self._worker_manager.submit(worker_id) + await self._on_worker_notification(note) + return note + @property def detail_run_id(self) -> str | None: return self._detail_run_id diff --git a/src/lecode/tui/feed.py b/src/lecode/tui/feed.py index f5945a2..8364d16 100644 --- a/src/lecode/tui/feed.py +++ b/src/lecode/tui/feed.py @@ -265,16 +265,27 @@ def turn_stats( def agent_summary(self, run: AgentRun) -> None: """One attributed line for a finished child run (the roster keeps detail).""" glyph, slot = { + "done": ("✔", "success"), "ok": ("✔", "success"), "error": ("✗", "error"), + "failed": ("✗", "error"), "cancelled": ("—", "muted"), + "stopped": ("■", "muted"), + "interrupted": ("!", "warning"), }.get(run.status, ("✔", "muted")) - parts = [f"{run.agent} · {run.description}"] + identity = f"{run.agent} · {run.description}" + if run.worker: + identity += f" · worker {run.run_id[:8]}" + parts = [identity] if run.activity: count = len(run.activity) parts.append(f"{count} tool call{'s' if count != 1 else ''}") if run.status == "error" and run.error: parts.append(run.error.splitlines()[0][:80]) + if run.worker: + parts.append( + format_cost(run.cost_usd) + (" incomplete" if run.usage_incomplete else "") + ) self._console.print( Text( f"[{self._stamp()}] {glyph} " + " · ".join(parts), style=getattr(self._theme, slot) diff --git a/src/lecode/tui/permission.py b/src/lecode/tui/permission.py index 36836c7..30fc8d7 100644 --- a/src/lecode/tui/permission.py +++ b/src/lecode/tui/permission.py @@ -4,12 +4,17 @@ ``awaiting approval`` state; keypresses are intercepted by the main app's keybindings, filtered on :attr:`ApprovalPrompt.is_pending`, so no nested prompt_toolkit application ever fights over stdin. + +Concurrent askers (parent turn + workers) queue FIFO. Only the head is +resolved by keypresses; queued entries keep waiting until they reach the +front. Cancelling an awaiter removes exactly that entry. """ from __future__ import annotations import asyncio from dataclasses import dataclass, field +from functools import partial from lecode.permission import ApprovalDecision @@ -23,6 +28,9 @@ class PendingApproval: target: str reason: str future: asyncio.Future[ApprovalDecision] = field(repr=False) + #: Attribution for "[w2 bash]" style prompts; None = the main turn. + worker: str | None = None + conversation: str = "main" def approval_prompt_text(tool_name: str, target: str) -> str: @@ -32,32 +40,56 @@ def approval_prompt_text(tool_name: str, target: str) -> str: class ApprovalPrompt: - """At most one pending approval; resolved by keypress or cancelled.""" + """FIFO queue of pending approvals; keypresses resolve only the head.""" def __init__(self) -> None: - self._pending: PendingApproval | None = None + self._queue: list[PendingApproval] = [] @property def pending(self) -> PendingApproval | None: - return self._pending + return self._queue[0] if self._queue else None @property def is_pending(self) -> bool: - return self._pending is not None - - def request(self, tool_name: str, target: str, reason: str) -> asyncio.Future[ApprovalDecision]: + return bool(self._queue) + + def request( + self, + tool_name: str, + target: str, + reason: str, + *, + worker: str | None = None, + conversation: str = "main", + ) -> asyncio.Future[ApprovalDecision]: future: asyncio.Future[ApprovalDecision] = asyncio.get_running_loop().create_future() - self._pending = PendingApproval(tool_name, target, reason, future) + entry = PendingApproval(tool_name, target, reason, future, worker, conversation) + self._queue.append(entry) + future.add_done_callback(partial(self._on_future_done, entry)) return future def resolve(self, decision: ApprovalDecision) -> None: - pending = self._pending - if pending is not None and not pending.future.done(): - pending.future.set_result(decision) - self._pending = None - - def cancel(self) -> None: - pending = self._pending - if pending is not None and not pending.future.done(): - pending.future.cancel() - self._pending = None + entry = self.pending + if entry is not None and not entry.future.done(): + entry.future.set_result(decision) + self._queue.pop(0) + + def cancel(self, future: asyncio.Future[ApprovalDecision] | None = None) -> None: + """Cancel the entry behind ``future`` only, or all outstanding when omitted.""" + if future is None: + queue, self._queue = self._queue, [] + for item in queue: + if not item.future.done(): + item.future.cancel() + return + for i, item in enumerate(self._queue): + if item.future is future: + del self._queue[i] + future.cancel() + return + + def _on_future_done( + self, entry: PendingApproval, future: asyncio.Future[ApprovalDecision] + ) -> None: + if future.cancelled() and entry in self._queue: + self._queue.remove(entry) diff --git a/tests/test_agent_builder.py b/tests/test_agent_builder.py index 5695773..14c3ba5 100644 --- a/tests/test_agent_builder.py +++ b/tests/test_agent_builder.py @@ -68,6 +68,16 @@ def test_session_grants_loaded(cwd, tmp_path): assert "session grant" in check.reason +def test_session_runtime_installs_workers(cwd, tmp_path): + from lecode.session.storage import SessionStore + + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("workers", cwd) + runtime = build_runtime(Config(), cwd, session=session, store=store) + assert "workers" in runtime.registry.names() + assert runtime.ctx.extras["workers"].session is session + + def test_agent_name_applies_overlay(cwd): runtime = build_runtime(Config(), cwd, agent_name="plan") checker = runtime.ctx.permission_checker diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 4f23131..1351f91 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from contextlib import asynccontextmanager import pytest from tests.fakes import FakeProvider, sample_catalog @@ -56,6 +57,25 @@ async def run(self, args, ctx) -> ToolExecResult: raise AssertionError("unreachable") +class TaskLikeTool(EchoTool): + def __init__(self) -> None: + super().__init__() + self.name = "task" + + +class LeaseManager: + def __init__(self) -> None: + self.suspensions = 0 + + def consume(self, id, history): + return [] + + @asynccontextmanager + async def suspend(self, id): + self.suspensions += 1 + yield + + def make_runner(tool_ctx, script, **kwargs) -> tuple[AgentRunner, FakeProvider]: provider = FakeProvider(script) kwargs.setdefault("catalog", sample_catalog()) @@ -136,6 +156,20 @@ async def test_tool_round_trip(tool_ctx): assert result_events[0].is_error is False +@pytest.mark.parametrize("names,expected", [(["task"], 1), (["task", "echo"], 0)]) +async def test_worker_suspends_only_all_task_batches(tool_ctx, names, expected): + manager = LeaseManager() + tool_ctx.extras.update({"workers": manager, "worker_id": "worker"}) + registry = ToolRegistry([TaskLikeTool(), EchoTool()]) + calls = [ + {"id": f"c{i}", "name": name, "arguments": '{"text":"ok"}'} for i, name in enumerate(names) + ] + provider = FakeProvider([{"tool_calls": calls}, {"text": "done"}]) + runner = AgentRunner(provider, registry, tool_ctx) + assert (await runner.run([{"role": "user", "content": "go"}])).final_text == "done" + assert manager.suspensions == expected + + async def test_parallel_tool_calls_paired_by_id(tool_ctx): script = [ { diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 5bd70d0..483a9f2 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -5,6 +5,7 @@ import json import pytest +from pydantic import ValidationError from lecode.config.loader import deep_merge, load_config from lecode.config.migrations import MIGRATIONS @@ -47,6 +48,7 @@ def test_defaults_validate_from_empty(): assert config.mcp.enable_exa is True assert config.mcp.enable_context7 is False assert config.memory.max_bytes == 32768 + assert config.worktree.validation == [] assert config.telemetry.enabled is False assert config.telemetry.sentry_dsn is None assert config.telemetry.otlp_endpoint is None @@ -72,6 +74,17 @@ def test_telemetry_section_parses(global_dir, tmp_path): assert result.warnings == [] +def test_worktree_validation_parses_and_requires_a_list(global_dir, tmp_path): + (global_dir / "config.toml").write_text( + '[worktree]\nvalidation = ["uv run pytest tests/test_worktree.py"]\n' + ) + assert load_config(cwd=tmp_path).config.worktree.validation == [ + "uv run pytest tests/test_worktree.py" + ] + with pytest.raises(ValidationError): + Config.model_validate({"worktree": {"validation": "pytest"}}) + + def test_first_run_creates_default_config(global_dir, tmp_path): result = load_config(cwd=tmp_path) created = global_dir / "config.toml" diff --git a/tests/test_permission_checker.py b/tests/test_permission_checker.py index 6068bb5..3a26fcd 100644 --- a/tests/test_permission_checker.py +++ b/tests/test_permission_checker.py @@ -14,9 +14,12 @@ ) -def _checker(permissions: dict, session_perms=None, cwd=None) -> PermissionChecker: +def _checker(permissions: dict, session_perms=None, cwd=None, read_only=False) -> PermissionChecker: return PermissionChecker( - Config.model_validate({"permissions": permissions}), session_perms, cwd=cwd + Config.model_validate({"permissions": permissions}), + session_perms, + cwd=cwd, + read_only=read_only, ) @@ -272,3 +275,225 @@ def test_overlay_shares_doom_tracking(): plan.check("read", args) third = checker.check("read", args) assert third.decision == Decision.ASK # count shared across derived checkers + + +# -- read-only enforcement ------------------------------------------------------ + + +def test_read_only_denies_writes_even_in_yolo(): + checker = _checker({"mode": "yolo"}, read_only=True) + assert checker.read_only is True + assert checker.check("bash", {"command": "ls"}).decision == Decision.DENY + assert checker.check("write", {"path": "a.py"}).decision == Decision.DENY + assert checker.check("edit", {"path": "a.py"}).decision == Decision.DENY + assert "read-only" in checker.check("bash", {"command": "ls"}).reason + + +def test_read_only_not_widened_by_overlay_allow_rule(): + overlay = AgentOverlay(extra_rules=_ruleset(allow={"bash": [{"pattern": "*"}]})) + checker = _checker({"mode": "yolo"}, read_only=True).for_agent(overlay) + assert checker.check("bash", {"command": "ls"}).decision == Decision.DENY + + +def test_read_only_not_widened_by_session_grant(): + perms = SessionPermissions([("bash", "*")]) + checker = _checker({"mode": "yolo"}, session_perms=perms, read_only=True) + assert checker.check("bash", {"command": "ls"}).decision == Decision.DENY + + +@pytest.mark.parametrize( + ("tool", "args"), + [ + ("read", {"path": "a.py"}), + ("grep", {"pattern": "x"}), + ("list_dir", {"path": "."}), + ("task", {"prompt": "inspect"}), + ], +) +def test_read_only_allows_read_class_tools(tool, args): + checker = _checker({"mode": "yolo"}, read_only=True) + assert checker.check(tool, args).decision == Decision.ALLOW + + +def test_for_agent_propagates_read_only(): + checker = _checker({"mode": "yolo"}, read_only=True) + derived = checker.for_agent(BUILTIN_AGENT_OVERLAYS["build"]) + assert derived.read_only is True + assert derived.check("write", {"path": "a"}).decision == Decision.DENY + + +def test_for_agent_read_only_narrows_writable_checker(): + checker = _checker({"mode": "yolo"}) + derived = checker.for_agent(BUILTIN_AGENT_OVERLAYS["build"], read_only=True) + assert derived.read_only is True + assert checker.read_only is False + assert checker.check("write", {"path": "a"}).decision == Decision.ALLOW + assert derived.check("write", {"path": "a"}).decision == Decision.DENY + + +def test_for_agent_cannot_widen_read_only_checker(): + checker = _checker({"mode": "yolo"}, read_only=True) + derived = checker.for_agent(BUILTIN_AGENT_OVERLAYS["build"], read_only=False) + assert derived.read_only is True + assert derived.check("bash", {"command": "ls"}).decision == Decision.DENY + + +def test_nested_agent_overlays_retain_each_ancestor_restriction(): + parent = _checker( + {"mode": "yolo", "rules": {"ask": {"read": [{"pattern": "global/*"}]}}} + ).for_agent( + AgentOverlay( + denied_tools=("bash",), + extra_rules=_ruleset( + deny={"read": [{"pattern": "denied/*"}]}, + ask={"read": [{"pattern": "private/*"}]}, + ), + ) + ) + child = parent.for_agent( + AgentOverlay(extra_rules=_ruleset(allow={"read": [{"pattern": "*"}]})) + ).for_agent(AgentOverlay(mode="yolo")) + for path, expected in ( + ("global/a", Decision.ASK), + ("denied/a", Decision.DENY), + ("private/a", Decision.ASK), + ("public/a", Decision.ALLOW), + ): + assert child.check("read", {"path": path}).decision == expected + assert child.check("bash", {"command": "ls"}).decision == Decision.DENY + + +def test_child_rules_and_grants_cannot_override_ancestor_ask_or_deny(): + parent = _checker({"mode": "yolo"}).for_agent( + AgentOverlay( + denied_tools=("write",), + extra_rules=_ruleset( + ask={"read": [{"pattern": "private/*"}]}, + deny={"read": [{"pattern": "denied/*"}]}, + ), + ) + ) + grants = SessionPermissions([("read", "*"), ("write", "*")]) + child = parent.for_child( + AgentOverlay(extra_rules=_ruleset(allow={"read": [{"pattern": "*"}]})), + session_perms=grants, + ).for_child(AgentOverlay(mode="yolo")) + assert child.check("read", {"path": "private/a"}).decision == Decision.ASK + assert child.check("read", {"path": "denied/a"}).decision == Decision.DENY + assert child.check("write", {"path": "a"}).decision == Decision.DENY + assert child.check("read", {"path": "public/a"}).decision == Decision.ALLOW + + +@pytest.mark.parametrize("derive", ["for_agent", "for_child"]) +def test_readonly_overlay_remains_effective_through_writable_descendants(derive): + parent = _checker({"mode": "yolo"}).for_agent(AgentOverlay(mode="readonly")) + child = getattr(parent, derive)( + AgentOverlay(mode="yolo", extra_rules=_ruleset(allow={"write": [{"pattern": "*"}]})) + ) + child.set_mode("yolo") + assert parent.read_only is True + assert child.read_only is True + assert child.mode == "readonly" + assert child.check("write", {"path": "a"}).decision == Decision.DENY + assert child.check("read", {"path": "a"}).decision == Decision.ALLOW + + +def test_child_path_rules_use_child_cwd_for_all_ancestor_layers(tmp_path): + parent_cwd, child_cwd = tmp_path / "parent", tmp_path / "child" + parent = _checker( + { + "mode": "yolo", + "rules": {"deny": {"read": [{"pattern": str(child_cwd / "secret")}]}}, + }, + cwd=parent_cwd, + ).for_agent( + AgentOverlay( + extra_rules=_ruleset(ask={"read": [{"pattern": str(child_cwd / "review")}]}), + ) + ) + child = parent.for_child(cwd=child_cwd).for_child() + assert child.check("read", {"path": "secret"}).decision == Decision.DENY + assert child.check("read", {"path": "review"}).decision == Decision.ASK + assert parent.check("read", {"path": "secret"}).decision == Decision.ALLOW + assert parent.check("read", {"path": "review"}).decision == Decision.ALLOW + assert child.check("read", {"path": str(parent_cwd / "secret")}).decision == Decision.ALLOW + + +def test_children_have_independent_doom_tracking_without_recording_parent_calls(): + parent = _checker({"mode": "yolo"}).for_agent(AgentOverlay()) + args = {"path": "same"} + assert parent.check("read", args).decision == Decision.ALLOW + assert parent.check("read", args).decision == Decision.ALLOW + child, sibling = parent.for_child(), parent.for_child() + grandchild = child.for_child() + for checker in (child, sibling, grandchild): + assert [checker.check("read", args).decision for _ in range(4)] == [ + Decision.ALLOW, + Decision.ALLOW, + Decision.ASK, + Decision.DENY, + ] + assert parent.check("read", args).decision == Decision.ASK + + +def test_child_uses_supplied_scoped_grants_without_sharing_parent_or_sibling_grants(): + parent_grants = SessionPermissions([("write", "src/*")]) + parent = _checker({"mode": "readonly"}, session_perms=parent_grants) + child_grants = SessionPermissions() + child = parent.for_child(session_perms=child_grants) + sibling = parent.for_child() + assert child.check("write", {"path": "src/before"}).decision == Decision.DENY + child_grants.grant("write", "src/approved") + child_grants.grant("write", "outside/*") + assert child.check("write", {"path": "src/approved"}).decision == Decision.ALLOW + assert child.check("write", {"path": "src/other"}).decision == Decision.DENY + assert child.check("write", {"path": "outside/file"}).decision == Decision.DENY + assert sibling.check("write", {"path": "src/approved"}).decision == Decision.DENY + assert parent.check("write", {"path": "src/other"}).decision == Decision.ALLOW + assert parent.check("write", {"path": "outside/file"}).decision == Decision.DENY + assert parent_grants.grants == [("write", "src/*")] + + +@pytest.mark.parametrize("source", ["global_allow", "global_ask", "grant", "overlay_allow"]) +def test_readonly_with_writable_exceptions_is_not_safe_for_shared_checkout(source): + permissions = {"mode": "readonly"} + grants = SessionPermissions() + overlay = AgentOverlay() + expected = Decision.ALLOW + if source.startswith("global_"): + decision = source.removeprefix("global_") + permissions["rules"] = {decision: {"write": [{"pattern": "allowed/*"}]}} + expected = Decision(decision) + elif source == "grant": + grants.grant("write", "allowed/*") + else: + permissions["mode"] = "yolo" + overlay = AgentOverlay( + mode="readonly", + extra_rules=_ruleset(allow={"write": [{"pattern": "allowed/*"}]}), + ) + checker = _checker(permissions, session_perms=grants).for_agent(overlay) + assert checker.read_only is False + assert checker.check("write", {"path": "allowed/a"}).decision == expected + assert checker.check("write", {"path": "other/a"}).decision == Decision.DENY + strict = checker.for_child(read_only=True) + assert strict.read_only is True + assert strict.check("write", {"path": "allowed/a"}).decision == Decision.DENY + + +@pytest.mark.parametrize("parent_decision", list(Decision)) +@pytest.mark.parametrize("overlay_decision", list(Decision)) +def test_per_call_overlay_composes_with_full_policy(parent_decision, overlay_decision): + parent = _checker( + {"mode": "yolo", "rules": {parent_decision: {"read": [{"pattern": "*"}]}}} + ).for_agent(AgentOverlay(denied_tools=("write",))) + overlay = AgentOverlay(extra_rules=_ruleset(**{overlay_decision: {"read": [{"pattern": "*"}]}})) + if Decision.DENY in (parent_decision, overlay_decision): + expected = Decision.DENY + elif Decision.ASK in (parent_decision, overlay_decision): + expected = Decision.ASK + else: + expected = Decision.ALLOW + assert parent.check("read", {"path": "a"}, overlay).decision == expected + assert parent.check("write", {"path": "a"}, overlay).decision == Decision.DENY + assert parent.check("read", {"path": "b"}).decision == parent_decision diff --git a/tests/test_session_stats.py b/tests/test_session_stats.py index 69bfb3c..16e4d2e 100644 --- a/tests/test_session_stats.py +++ b/tests/test_session_stats.py @@ -40,6 +40,7 @@ def test_counts_and_roles(store, session): assert stats.role_counts == {"user": 1, "assistant": 2} assert stats.input_tokens == 1_001_000 assert stats.output_tokens == 1_000_500 + assert stats.usage_incomplete is False def test_cost_combines_recorded_and_catalog_pricing(store, session): @@ -108,3 +109,52 @@ def test_context_tokens_zero_without_usage(store): s = store.create("plain", cwd="/tmp") store.append_message(s, {"role": "assistant", "content": "no usage"}) assert session_stats(store, s).context_tokens == 0 + + +# -- worker usage ---------------------------------------------------------------- + + +def test_worker_usage_summed_once(store, session): + store.append_event( + session, + "worker_usage", + {"input_tokens": 100, "output_tokens": 50, "cost_usd": 0.002}, + ) + store.append_event( + session, + "worker_usage", + {"usage": {"input_tokens": 200, "output_tokens": 25, "cost_usd": 0.003}}, + ) + stats = session_stats(store, session, catalog=sample_catalog()) + assert stats.input_tokens == 1_001_000 + 300 + assert stats.output_tokens == 1_000_500 + 75 + assert stats.cost_usd == pytest.approx(2.26 + 0.005) + + +def test_worker_usage_includes_failed_and_cancelled_dispatches(store, session): + store.append_event( + session, + "worker_usage", + {"status": "failed", "input_tokens": 10, "output_tokens": 5}, + ) + store.append_event( + session, + "worker_usage", + {"status": "cancelled", "input_tokens": 7, "output_tokens": 3}, + ) + stats = session_stats(store, session, catalog=sample_catalog()) + assert stats.input_tokens == 1_001_000 + 17 + assert stats.output_tokens == 1_000_500 + 8 + + +def test_worker_usage_incomplete_flag(store, session): + store.append_event(session, "worker_usage", {"input_tokens": 1, "output_tokens": 1}) + assert session_stats(store, session).usage_incomplete is False + store.append_event( + session, + "worker_usage", + {"input_tokens": 1, "output_tokens": 1, "incomplete": True}, + ) + stats = session_stats(store, session) + assert stats.usage_incomplete is True + assert stats.input_tokens == 1_001_000 + 2 diff --git a/tests/test_session_storage.py b/tests/test_session_storage.py index 659358c..36ab139 100644 --- a/tests/test_session_storage.py +++ b/tests/test_session_storage.py @@ -333,6 +333,38 @@ def test_agent_runs_respect_tombstones(store, session): assert store.load_agent_runs(session) == [] +def test_load_agent_runs_matches_generic_events(store, session): + store.record_agent_run(session, {"run_id": "one", "agent": "explore", "status": "ok"}) + store.append_event(session, "worker_usage", {"input_tokens": 3}) + assert store.load_agent_runs(session) == store.load_events(session, "agent_run") + assert [run["run_id"] for run in store.load_agent_runs(session)] == ["one"] + + +# -- generic event loading ------------------------------------------------------ + + +def test_load_events_in_append_order(store, session): + store.append_event(session, "worker", {"id": "one"}) + store.append_event(session, "worker_state", {"id": "one", "state": "done"}) + store.append_event(session, "worker", {"id": "two"}) + assert store.load_events(session, "worker") == [{"id": "one"}, {"id": "two"}] + assert store.load_events(session, "worker_state") == [{"id": "one", "state": "done"}] + + +def test_load_events_respects_tombstones(store, session): + first = store.append_event(session, "worker_usage", {"dispatch": 1}) + hidden = store.append_event(session, "worker_usage", {"dispatch": 2}) + store.append_tombstone(session, up_to_seq=hidden.seq - 1) + last = store.append_event(session, "worker_usage", {"dispatch": 3}) + assert store.load_events(session, "worker_usage") == [{"dispatch": 1}, {"dispatch": 3}] + assert first.seq < hidden.seq < last.seq + + +def test_load_events_unknown_kind(store, session): + store.record_agent_run(session, {"run_id": "r", "agent": "explore", "status": "ok"}) + assert store.load_events(session, "no_such_kind") == [] + + # -- attach locking ------------------------------------------------------------- diff --git a/tests/test_slash_features.py b/tests/test_slash_features.py index cecb44a..a5734f5 100644 --- a/tests/test_slash_features.py +++ b/tests/test_slash_features.py @@ -6,6 +6,49 @@ from tests.test_tui_app import make_app from tests.test_worktree import make_repo + +async def test_agent_focus_command_targets_persistent_worker(tmp_path, monkeypatch): + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "done"}]) + manager = app.worker_manager + assert manager is not None + worker = await manager.start( + app.runtime.ctx, + agent="explore", + prompt="inspect", + origin="human", + background=True, + ) + await manager.wait(worker.id) + await app.handle_command("/agent 1 focus") + assert app._focused_worker_id == worker.id + assert "composer focused on @explore" in out.getvalue() + await manager.shutdown() + + +async def test_agent_submit_wakes_root_once(tmp_path, monkeypatch): + app, provider, _ = make_app(tmp_path, monkeypatch, [{"text": "done"}, {"text": "integrated"}]) + manager = app.worker_manager + assert manager is not None + worker = await manager.start( + app.runtime.ctx, + agent="explore", + prompt="inspect", + origin="human", + background=True, + ) + await manager.wait(worker.id) + await app.handle_command("/agent 1 submit") + assert app._worker_wake_task is not None + await app._worker_wake_task + assert app._turn_task is not None + await app._turn_task + assert any( + message["content"] == "Worker updates are available." + for message in provider.requests[-1]["messages"] + ) + await manager.shutdown() + + # -- /init ------------------------------------------------------------------------- diff --git a/tests/test_subagents.py b/tests/test_subagents.py index 3b94dea..38d7a8e 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -578,7 +578,8 @@ async def test_task_result_event_carries_run_id_metadata(tmp_path, monkeypatch): assert task_results run_id = task_results[0].metadata.get("run_id") assert run_id - assert run_id == store.load_agent_runs(session)[0]["run_id"] + assert run_id == task_results[0].metadata["worker_id"] + assert runtime.ctx.extras["workers"].get(run_id).session_id != session.id async def test_roster_panel_visible_while_child_runs(tmp_path, monkeypatch): @@ -682,22 +683,21 @@ async def test_at_agent_runs_directly(tmp_path, monkeypatch): script = [{"text": "42 files", "usage": {"input_tokens": 7, "output_tokens": 3}}] app, provider, out = make_app(tmp_path, monkeypatch, script) await app._submit("@explore count the files") - await app._turn_task + task = app._turn_task + assert task is not None and not task.done() + await task rendered = out.getvalue() assert "> @explore count the files" in rendered - assert "42 files" in rendered request = provider.requests[0] assert "read-only exploration agent" in request["messages"][0]["content"] assert request["messages"][1] == {"role": "user", "content": "count the files"} - # A side query: no message history, but the run itself is persisted. + # A persistent human worker: no root history or legacy agent-run record. assert app.store.load_messages(app.session) == [] - runs = app.store.load_agent_runs(app.session) - assert len(runs) == 1 - assert runs[0]["agent"] == "explore" - assert runs[0]["status"] == "ok" - assert runs[0]["prompt"] == "count the files" - assert runs[0]["answer"] == "42 files" - assert app._last_response == "42 files" + worker = app.worker_manager.list()[0] + assert worker.origin == "human" + assert worker.result is not None and worker.result.final_text == "42 files" + assert app.store.load_agent_runs(app.session) == [] + assert "42 files" not in rendered # completion notifies until /agent submit assert app._status.input_tokens == 7 diff --git a/tests/test_tui_agents.py b/tests/test_tui_agents.py index 0596110..f450e3c 100644 --- a/tests/test_tui_agents.py +++ b/tests/test_tui_agents.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + from lecode.agent.runner import Done, Error, LlmCall, ToolCall, ToolResult from lecode.extras.subagents import SubagentProgress from lecode.tui.agents import ( @@ -102,3 +104,43 @@ def test_detail_lines_show_tools_and_answer(): assert "a.py" in text assert "line one" in text assert "the answer" in text + + +def test_worker_tree_shows_cost_and_selected_context(): + roster = AgentRoster() + root = SimpleNamespace( + id="worker-root", + parent_id=None, + depth=1, + agent="explore", + description="Scan the repository", + origin="delegated", + state="running", + error=None, + result=None, + usage_totals=SimpleNamespace(cost_usd=0.0123, context_tokens=1234), + usage_incomplete=False, + ) + child = SimpleNamespace( + id="worker-child", + parent_id="worker-root", + depth=2, + agent="explore", + description="Inspect tests", + origin="delegated", + state="interrupted", + error=None, + result=None, + usage_totals=SimpleNamespace(cost_usd=0.0, context_tokens=0), + usage_incomplete=True, + ) + roster.sync_worker(root, context_window=200_000) + roster.sync_worker(child, context_window=200_000) + compact = _plain(roster_lines(roster, THEME, width=100)) + detail = _plain(detail_lines(roster.get("worker-root"), THEME, width=100)) + assert "workers · 1 running · 1 done · $0.0123 incomplete" in compact + assert "Inspect tests" in compact and "interrupted" in compact + assert "cost: $0.0123" in detail + assert "subtree: $0.0123 (incomplete)" in detail + assert "ctx: 1.2k/200.0k" in detail + assert "elapsed:" in detail diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index 78c543b..308922e 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -136,6 +136,40 @@ def test_roster_panel_renders_as_prompt_toolkit_text(tmp_path, monkeypatch): assert "Scan src" in rendered +async def test_worker_focus_preserves_drafts_and_routes_composer(tmp_path, monkeypatch): + app, _, _ = make_app( + tmp_path, + monkeypatch, + [{"text": "initial"}, {"text": "follow-up"}], + ) + manager = app.worker_manager + assert manager is not None + with create_pipe_input() as inp: + app._build_app(input=inp, output=DummyOutput()) + worker = await manager.start( + app.runtime.ctx, + agent="explore", + prompt="inspect this", + origin="human", + background=True, + ) + await manager.wait(worker.id) + app._input_area.buffer.text = "main draft" + assert app.focus_worker(worker.id) + app._input_area.buffer.text = "worker draft" + assert app.focus_worker(None) + assert app._input_area.text == "main draft" + assert app.focus_worker(worker.id) + assert app._input_area.text == "worker draft" + await app._submit("follow up") + assert (await manager.wait(worker.id)).final_text == "follow-up" + assert "follow up" in [item["text"] for item in manager.pending(worker.id)] or any( + message["content"] == "follow up" + for message in app.store.load_for_model(worker.session) + ) + await manager.shutdown() + + def _buffer(app): return app._input_area.buffer diff --git a/tests/test_tui_permission.py b/tests/test_tui_permission.py index 77c8a25..e7d07e8 100644 --- a/tests/test_tui_permission.py +++ b/tests/test_tui_permission.py @@ -196,6 +196,101 @@ async def test_approval_prompt_request_resolve_cancel(): await future2 +# -- FIFO queue ------------------------------------------------------------------- + + +async def test_concurrent_requests_resolve_in_fifo_order(): + prompt = ApprovalPrompt() + first = prompt.request("bash", "ls", "r1") + second = prompt.request("bash", "rm", "r2") + assert prompt.pending.tool_name == "bash" and prompt.pending.reason == "r1" + prompt.resolve(AllowOnce()) + assert await first == AllowOnce() + assert not second.done() + assert prompt.pending.reason == "r2" + prompt.resolve(Deny()) + assert await second == Deny() + assert not prompt.is_pending + + +async def test_three_requests_head_grants_only(): + """y/a resolve only the head; a queued request is never granted.""" + prompt = ApprovalPrompt() + futures = [prompt.request("bash", f"c{i}", f"r{i}") for i in range(3)] + prompt.resolve(AllowAlways(pattern="c0")) + assert await futures[0] == AllowAlways(pattern="c0") + assert all(not f.done() for f in futures[1:]) + prompt.resolve(Deny()) + prompt.resolve(Deny()) + results = await asyncio.gather(*futures[1:]) + assert results == [Deny(), Deny()] + assert not prompt.is_pending + + +async def test_cancel_head_promotes_next_without_disturbing_tail(): + prompt = ApprovalPrompt() + head = prompt.request("bash", "ls", "r1") + tail = prompt.request("bash", "rm", "r2") + head.cancel() # awaiter cancelled + with pytest.raises(asyncio.CancelledError): + await head + await asyncio.sleep(0) # let the done-callback remove the entry + assert prompt.pending.reason == "r2" + prompt.resolve(AllowOnce()) + assert await tail == AllowOnce() + assert not prompt.is_pending + + +async def test_cancel_tail_keeps_head(): + prompt = ApprovalPrompt() + head = prompt.request("bash", "ls", "r1") + tail = prompt.request("bash", "rm", "r2") + tail.cancel() + with pytest.raises(asyncio.CancelledError): + await tail + await asyncio.sleep(0) + assert prompt.pending.reason == "r1" + prompt.resolve(Deny()) + assert await head == Deny() + assert not prompt.is_pending + + +async def test_scoped_cancel_future_leaves_others(): + prompt = ApprovalPrompt() + head = prompt.request("bash", "ls", "r1") + middle = prompt.request("bash", "rm", "r2") + tail = prompt.request("bash", "cd", "r3") + prompt.cancel(middle) + assert middle.cancelled() + assert prompt.pending.reason == "r1" + prompt.resolve(AllowOnce()) + assert prompt.pending.reason == "r3" + prompt.cancel(tail) + assert prompt.pending is None + assert await head == AllowOnce() + + +async def test_attribution_passthrough_and_defaults(): + prompt = ApprovalPrompt() + prompt.request("bash", "ls", "r1", worker="w2", conversation="sub-a") + prompt.request("bash", "rm", "r2") + pending = prompt.pending + assert pending.worker == "w2" and pending.conversation == "sub-a" + prompt.resolve(AllowOnce()) + queued = prompt.pending + assert queued.worker is None and queued.conversation == "main" + + +async def test_shutdown_cancels_all_outstanding(): + prompt = ApprovalPrompt() + futures = [prompt.request("bash", f"c{i}", f"r{i}") for i in range(3)] + prompt.cancel() + assert not prompt.is_pending + for future in futures: + with pytest.raises(asyncio.CancelledError): + await future + + # -- app integration --------------------------------------------------------------- @@ -249,6 +344,22 @@ async def test_request_approval_shows_doom_reason(tmp_path, monkeypatch): task.result() # consume +async def test_worker_approval_is_attributed_at_fifo_head(tmp_path, monkeypatch): + app, out = make_app(tmp_path, monkeypatch, []) + first = asyncio.ensure_future( + app._request_approval("bash", {"command": "ls"}, "", worker="worker-1234", conversation="w") + ) + second = asyncio.ensure_future(app._request_approval("bash", {"command": "pwd"}, "")) + await wait_for(lambda: app._approval.pending is not None) + assert app._approval.pending.worker == "worker-1234" + app._resolve_approval(AllowOnce()) + assert await first == AllowOnce() + assert app._approval.pending is not None and app._approval.pending.worker is None + app._resolve_approval(Deny()) + assert await second == Deny() + assert "[worker worker-1 · w] allow bash 'ls'?" in out.getvalue() + + async def test_pipe_approval_y_runs_asked_tool(tmp_path, monkeypatch): """Full flow: model calls bash, user answers 'y', output appears.""" script = [ diff --git a/tests/test_tui_streaming_pty.py b/tests/test_tui_streaming_pty.py index c8d8ef9..0643448 100644 --- a/tests/test_tui_streaming_pty.py +++ b/tests/test_tui_streaming_pty.py @@ -288,6 +288,31 @@ async def test_tool_round_then_answer(tmp_path, monkeypatch): assert "all done here" in dump, "final answer missing:\n" + dump +async def test_focused_worker_uses_the_current_composer(tmp_path, monkeypatch): + """A direct worker focuses in-place; no child TUI or alternate screen opens.""" + + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + await asyncio.sleep(0.8) + os.write(master, b"@explore inspect\r") + await _wait_for(lines, "started for @explore") + os.write(master, b"/agent 1 focus\r") + await _wait_for(lines, "to @explore") + os.write(master, b"follow up\r") + await _wait_for(lines, "sent to @explore") + os.write(master, b"/quit\r") + return lines() + + lines = await _run_pty_app( + tmp_path, + monkeypatch, + [{"text": ["first"]}, {"text": ["second"]}], + drive, + delay=0.05, + ) + assert any("to @explore" in line for line in lines) + + async def test_slash_menu_renders_and_no_match_row(tmp_path, monkeypatch): """Typing '/mod' shows the dropdown on the real terminal (several rows at once, not clipped); an unknown prefix shows the inert 'No matching diff --git a/tests/test_workers.py b/tests/test_workers.py new file mode 100644 index 0000000..9ac2311 --- /dev/null +++ b/tests/test_workers.py @@ -0,0 +1,786 @@ +"""WorkerManager's public persistence and scheduling contract.""" + +import asyncio +import subprocess +from dataclasses import replace +from unittest.mock import Mock + +import pytest +from tests.fakes import FakeProvider + +from lecode.agent.builder import build_runtime +from lecode.agent.runner import AgentRunner, RunResult, UsageTotals +from lecode.config.models import Config +from lecode.context.agents import AgentDefinition, AgentRegistry +from lecode.extras.subagents import SubagentError +from lecode.extras.workers import WORKER_CURRENT_EXTRA, WorkerManager +from lecode.extras.worktree import WorktreeError, WorktreeManager +from lecode.permission import PermissionChecker +from lecode.permission.checker import AgentOverlay +from lecode.session.stats import session_stats +from lecode.session.storage import SessionInUseError, SessionStore + + +class GatedProvider(FakeProvider): + def __init__(self): + super().__init__([]) + self.release = asyncio.Event() + self.started = asyncio.Queue() + self.active = 0 + self.peak = 0 + + async def _stream(self, entry): + self.active += 1 + self.peak = max(self.peak, self.active) + self.started.put_nowait(None) + try: + await self.release.wait() + async for event in super()._stream({"text": "done"}): + yield event + finally: + self.active -= 1 + + +@pytest.fixture +def setup(tmp_path, monkeypatch): + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "config")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + cwd = tmp_path / "project" + cwd.mkdir() + monkeypatch.chdir(cwd) + config = Config() + config.memory.enabled = False + config.lsp.enabled = False + config.pierre.enabled = False + store = SessionStore(tmp_path / "config") + session = store.create("main", cwd) + runtime = build_runtime(config, cwd, session=session, store=store) + provider = FakeProvider([{"text": "first"}, {"text": "second"}]) + runtime.ctx.extras["provider"] = provider + manager = WorkerManager(config, cwd=cwd, root_ctx=runtime.ctx) + runtime.ctx.extras["workers"] = manager + return manager, runtime.ctx, provider, store, session + + +@pytest.fixture +def checker_contract(setup): + """Isolate the sibling-owned checker seam in runner/scheduler tests. + + The real checker integration is tested separately, once for_child lands. + This double does not purport to verify inherited permission restrictions. + """ + _, ctx, _, _, _ = setup + parent = ctx.permission_checker + checker = Mock(wraps=parent) + checker.mode = parent.mode + checker.read_only = parent.read_only + + def derive(overlay=None, **kw): + child = PermissionChecker(ctx.config, mode=parent.mode, _overlay=overlay, **kw) + derived = Mock(wraps=child) + derived.mode = child.mode + derived.read_only = child.read_only + derived.for_child = Mock(side_effect=derive) + return derived + + checker.for_child = Mock(side_effect=derive) + ctx.permission_checker = checker + return checker + + +@pytest.mark.asyncio +async def test_followup_is_persisted_and_replayed(setup, checker_contract): + manager, ctx, provider, store, _ = setup + worker = await manager.start(ctx, agent="explore", prompt="question") + try: + assert (await manager.wait(worker.id)).final_text == "first" + await manager.send(worker.id, "follow up") + assert (await manager.wait(worker.id)).final_text == "second" + messages = store.load_for_model(worker.session) + assert [(m["role"], m["content"]) for m in messages] == [ + ("user", "question"), + ("assistant", "first"), + ("user", "follow up"), + ("assistant", "second"), + ] + assert provider.requests[1]["messages"][1:] == messages[:-1] + assert manager.pending(worker.id) == [] + assert worker.usage_incomplete # Provider omitted usage, not a known zero cost. + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_task_tool_creates_persisted_workers_and_nests(setup, checker_contract): + manager, ctx, _, _, _ = setup + ctx.extras["provider"] = FakeProvider( + [ + {"tool_calls": [{"id": "nested", "name": "task", "arguments": '{"prompt":"child"}'}]}, + {"text": "child result"}, + {"text": "parent result"}, + ] + ) + try: + outer = await manager.start(ctx, agent="explore", prompt="parent") + assert (await manager.wait(outer.id)).final_text == "parent result" + child = manager.children(outer.id)[0] + assert child.depth == 2 + assert child.result is not None and child.result.final_text == "child result" + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_background_task_tool_delivers_worker_notification(setup, checker_contract): + manager, ctx, _, _, _ = setup + try: + _, result = await ctx.extras["registry"].dispatch_result( + "background", "task", '{"prompt":"scan","run_in_background":true}', ctx + ) + assert not result.is_error + worker_id = result.metadata["worker_id"] + await manager.wait(worker_id) + history = [] + assert manager.consume(None, history) + assert worker_id in history[0]["content"] + assert "first" in history[0]["content"] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_parent_cancellation_does_not_cancel_managed_worker(setup, checker_contract): + manager, ctx, _, _, _ = setup + provider = GatedProvider() + ctx.extras["provider"] = provider + call = asyncio.create_task( + ctx.extras["registry"].dispatch_result("task", "task", '{"prompt":"scan"}', ctx) + ) + try: + async with asyncio.timeout(2): + await provider.started.get() + worker = manager.list()[0] + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + assert worker.state == "running" + provider.release.set() + await manager.wait(worker.id) + finally: + if not call.done(): + call.cancel() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_worker_controls_enforce_descendant_hierarchy_and_questions(setup, checker_contract): + from lecode.agent.tools.workers import make_tool + + manager, ctx, _, _, _ = setup + parent = await manager.start(ctx, agent="explore", prompt="parent") + nested = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: parent.id}) + child = await manager.start(nested, agent="explore", prompt="child") + child_ctx = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: child.id}) + tool = make_tool() + try: + listed = await tool.run({"action": "list"}, nested) + assert child.id in listed.content and parent.id not in listed.content + denied = await tool.run({"action": "stop", "id": parent.id}, nested) + assert denied.is_error and "descendants" in denied.content + invalid = await tool.run({"action": "send", "id": child.id, "text": 3}, nested) + assert invalid.is_error + asked = await tool.run({"action": "question", "text": "Need a choice"}, child_ctx) + assert asked.content == "question sent to parent" + history = [] + assert manager.consume(parent.id, history) + assert any("Need a choice" in message["content"] for message in history) + finally: + await manager.shutdown() + + +def git(cwd, *args): + return subprocess.run( + ["git", *args], cwd=cwd, check=True, capture_output=True, text=True + ).stdout.strip() + + +@pytest.mark.asyncio +async def test_write_worktree_pins_parent_head_and_readonly_shares_parent_cwd( + setup, checker_contract +): + manager, ctx, _, _, _ = setup + git(ctx.cwd, "init", "-b", "main") + git( + ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--allow-empty", + "-m", + "initial", + ) + agents = ctx.extras["agents"] + ctx.extras["agents"] = AgentRegistry( + { + "explore": agents.get("explore"), + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + } + ) + try: + writer = await manager.start(ctx, agent="writer", prompt="write") + await manager.wait(writer.id) + assert writer.cwd != ctx.cwd + assert writer.worktree.path == writer.cwd + wm = await WorktreeManager.discover(ctx.cwd) + sidecar = wm.read_sidecar(writer.worktree.name) + assert sidecar["dest_path"] == str(ctx.cwd) + assert sidecar["dest_branch"] == "main" + assert sidecar["base_commit"] == git(ctx.cwd, "rev-parse", "HEAD") + nested = replace( + ctx, cwd=writer.cwd, extras={**ctx.extras, WORKER_CURRENT_EXTRA: writer.id} + ) + reader = await manager.start(nested, agent="explore", prompt="read") + await manager.wait(reader.id) + assert reader.cwd == writer.cwd + assert reader.worktree is None + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_write_worker_refuses_non_git(setup): + manager, ctx, _, _, _ = setup + ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + } + ) + try: + with pytest.raises(WorktreeError, match="git"): + await manager.start(ctx, agent="writer", prompt="write") + assert manager.list() == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_restart_keeps_inbox_and_stopped_intent_without_executing(setup): + manager, ctx, provider, store, session = setup + worker = await manager.start(ctx, agent="explore", prompt="original") + await manager.stop(worker.id) + await manager.send(worker.id, "later") + await manager.shutdown() + restored = WorkerManager(ctx.config, cwd=ctx.cwd, root_ctx=ctx, store=store, session=session) + try: + assert restored.load()[0].id == worker.id + assert restored.get(worker.id).state == "stopped" + assert [m["text"] for m in restored.pending(worker.id)] == ["original", "later"] + assert provider.requests == [] + assert restored.load() == restored.list() + finally: + await restored.shutdown() + + +@pytest.mark.asyncio +async def test_depth_and_agent_eligibility_fail_before_creating_sessions(setup): + manager, ctx, _, store, _ = setup + try: + first = await manager.start(ctx, agent="explore", prompt="one") + nested = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: first.id}) + second = await manager.start(nested, agent="explore", prompt="two") + nested = replace(nested, extras={**nested.extras, WORKER_CURRENT_EXTRA: second.id}) + before = len(store.list_sessions()) + with pytest.raises(SubagentError, match="depth"): + await manager.start(nested, agent="explore", prompt="three") + for name in ("missing", "build"): + with pytest.raises(SubagentError, match="ineligible"): + await manager.start(ctx, agent=name, prompt="no") + assert len(store.list_sessions()) == before + assert manager.children(first.id) == [second] + assert [first.depth, second.depth] == [1, 2] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_strict_cap_and_shielded_wait(setup, checker_contract): + manager, ctx, _, _, _ = setup + provider = GatedProvider() + ctx.extras["provider"] = provider + try: + workers = [await manager.start(ctx, agent="explore", prompt=str(i)) for i in range(12)] + async with asyncio.timeout(2): + for _ in range(10): + await provider.started.get() + await asyncio.sleep(0) + assert provider.active == 10 + assert sum(w.state == "queued" for w in workers) == 2 + waiter = asyncio.create_task(manager.wait(workers[0].id)) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert workers[0].state == "running" + provider.release.set() + await asyncio.gather(*(manager.wait(w.id) for w in workers)) + assert provider.peak == 10 + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_send_does_not_wake_stopped_worker_even_with_interrupt(setup): + manager, ctx, provider, _, _ = setup + try: + worker = await manager.start(ctx, agent="explore", prompt="initial") + await manager.stop(worker.id) + await manager.send(worker.id, "steer", interrupt=True) + await asyncio.sleep(0) + assert worker.state == "stopped" + assert provider.requests == [] + assert [m["text"] for m in manager.pending(worker.id)] == ["initial", "steer"] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_failed_usage_survives_restart_and_child_locks_last_until_shutdown( + setup, checker_contract +): + manager, ctx, _, store, session = setup + ctx.extras["provider"] = FakeProvider( + [ + { + "tool_calls": [{"id": "call", "name": "read", "arguments": '{"path":"missing"}'}], + "usage": {"input_tokens": 7, "output_tokens": 3, "cost_usd": 0.25}, + }, + {"error": RuntimeError("provider broke")}, + ] + ) + worker = await manager.start(ctx, agent="explore", prompt="question") + try: + with pytest.raises(SubagentError, match="provider broke"): + await manager.wait(worker.id) + assert worker.usage_totals.input_tokens == 7 + assert worker.usage_totals.cost_usd == 0.25 + assert worker.usage_incomplete + assert [m["role"] for m in store.load_for_model(worker.session)] == [ + "user", + "assistant", + "tool", + ] + assert store.load_events(worker.session, "worker_usage_checkpoint")[-1]["input_tokens"] == 7 + with pytest.raises(SessionInUseError): + store.acquire_lock(worker.session) + finally: + await manager.shutdown() + lock = store.acquire_lock(worker.session) + lock.release() + assert worker.session.path.with_suffix(".lock").exists() + restored = WorkerManager(ctx.config, cwd=ctx.cwd, root_ctx=ctx, store=store, session=session) + try: + restored.load() + assert restored.get(worker.id).usage_totals == worker.usage_totals + assert restored.get(worker.id).state == "failed" + finally: + await restored.shutdown() + + +@pytest.mark.asyncio +async def test_load_interrupted_uses_child_usage_if_root_snapshot_lagged(setup): + manager, ctx, provider, store, session = setup + worker = await manager.start(ctx, agent="explore", prompt="original") + snapshot = store.load_events(session, "worker")[-1] + await manager.shutdown() + # Simulate a process dying after child usage persisted, before root update. + store.append_event( + worker.session, + "worker_usage_checkpoint", + { + "dispatch_id": worker.dispatch_id, + "input_tokens": 9, + "output_tokens": 2, + "cost_usd": 0.5, + "context_tokens": 9, + }, + ) + store.append_event(session, "worker", snapshot) + restored = WorkerManager(ctx.config, cwd=ctx.cwd, root_ctx=ctx) + try: + loaded = restored.load()[0] + assert loaded.state == "interrupted" + assert loaded.usage_incomplete + assert loaded.usage_totals.input_tokens == 9 + assert provider.requests == [] + finally: + await restored.shutdown() + + +@pytest.mark.asyncio +async def test_interrupt_repairs_unanswered_calls_without_replaying_inputs(setup, checker_contract): + manager, ctx, provider, store, _ = setup + worker = await manager.start(ctx, agent="explore", prompt="original") + await manager.stop(worker.id) + manager.consume(worker.id, []) + store.append_message( + worker.session, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "unanswered", + "type": "function", + "function": {"name": "read", "arguments": "{}"}, + } + ], + }, + ) + try: + message_id = await manager.send(worker.id, "new direction", interrupt=True) + assert manager.pending(worker.id)[0]["id"] == message_id + with pytest.raises(RuntimeError, match="outstanding"): + manager.consume(worker.id, store.load_for_model(worker.session)) + assert [m["role"] for m in store.load_for_model(worker.session)] == ["user", "assistant"] + await manager.resume(worker.id) + await manager.wait(worker.id) + history = provider.requests[0]["messages"][1:] + assert [m["role"] for m in history] == ["user", "assistant", "tool", "user"] + assert history[2]["tool_call_id"] == "unanswered" + assert "unknown" in history[2]["content"] + assert history[3]["content"] == "new direction" + assert manager.pending(worker.id) == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_running_inbox_is_durable_but_only_consumed_after_safe_boundary( + setup, checker_contract +): + manager, ctx, _, store, _ = setup + provider = GatedProvider() + ctx.extras["provider"] = provider + try: + worker = await manager.start(ctx, agent="explore", prompt="initial") + async with asyncio.timeout(2): + await provider.started.get() + await manager.send(worker.id, "follow up") + assert [m["content"] for m in store.load_for_model(worker.session)] == ["initial"] + assert [m["text"] for m in manager.pending(worker.id)] == ["follow up"] + provider.release.set() + await manager.wait(worker.id) + assert len(provider.requests) == 2 + assert manager.pending(worker.id) == [] + assert [m["role"] for m in store.load_for_model(worker.session)] == [ + "user", + "assistant", + "user", + "assistant", + ] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_completion_delivery_background_only_and_human_submit(setup, checker_contract): + manager, ctx, _, _, _ = setup + notifications = [] + manager.notify = notifications.append + try: + foreground = await manager.start(ctx, agent="explore", prompt="foreground") + await manager.wait(foreground.id) + assert manager.drain_notifications() == [] + assert notifications == [] + background = await manager.start(ctx, agent="explore", prompt="background", background=True) + await manager.wait(background.id) + notes = manager.drain_notifications() + assert [n["worker_id"] for n in notes] == [background.id] + assert manager.drain_notifications() == [] + human = await manager.start( + ctx, agent="explore", prompt="human", origin="human", background=True + ) + await manager.wait(human.id) + assert notifications[-1]["worker_id"] == human.id + assert manager.drain_notifications() == [] + await manager.submit(human.id) + assert [n["worker_id"] for n in manager.drain_notifications()] == [human.id] + await manager.submit(human.id) + assert manager.drain_notifications() == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_completed_worker_followup_after_restart(setup, checker_contract): + manager, ctx, provider, store, _ = setup + worker = await manager.start(ctx, agent="explore", prompt="first question") + await manager.wait(worker.id) + await manager.shutdown() + restored = WorkerManager(ctx.config, cwd=ctx.cwd, root_ctx=ctx) + try: + restored.load() + assert (await restored.wait(worker.id)).final_text == "first" + await restored.send(worker.id, "second question") + assert (await restored.wait(worker.id)).final_text == "second" + assert provider.requests[1]["messages"][1:] == store.load_for_model(worker.session)[:-1] + finally: + await restored.shutdown() + + +@pytest.mark.asyncio +async def test_suspended_supervisors_release_capacity_and_reacquire( + setup, checker_contract, monkeypatch +): + manager, ctx, _, _, _ = setup + children_started = asyncio.Queue() + release = asyncio.Event() + + async def run(runner, messages, on_event=None): + id = runner.ctx.extras[WORKER_CURRENT_EXTRA] + assert sum(w.state == "running" for w in manager.list()) <= 10 + if manager.get(id).depth == 1: + child = await manager.start(runner.ctx, agent="explore", prompt="child") + async with manager.suspend(id): + assert manager.get(id).state == "waiting" + await manager.wait(child.id) + assert manager.get(id).state == "running" + assert sum(w.state == "running" for w in manager.list()) <= 10 + else: + children_started.put_nowait(id) + await release.wait() + return RunResult("done", 1, "done", UsageTotals()) + + monkeypatch.setattr(AgentRunner, "run", run) + try: + parents = [await manager.start(ctx, agent="explore", prompt="parent") for _ in range(10)] + async with asyncio.timeout(2): + for _ in range(10): + await children_started.get() + assert sum(w.state == "waiting" for w in manager.list()) == 10 + assert sum(w.state == "running" for w in manager.list()) == 10 + with pytest.raises(RuntimeError, match="supervisor"): + async with manager.suspend(manager.children(parents[0].id)[0].id): + pass + release.set() + async with asyncio.timeout(2): + await asyncio.gather(*(manager.wait(w.id) for w in parents)) + assert all(w.state == "completed" for w in manager.list()) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_stop_is_individual_unless_tree_requested(setup, checker_contract): + manager, ctx, _, _, _ = setup + provider = GatedProvider() + ctx.extras["provider"] = provider + try: + parent = await manager.start(ctx, agent="explore", prompt="parent") + nested = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: parent.id}) + child = await manager.start(nested, agent="explore", prompt="child") + other = await manager.start(ctx, agent="explore", prompt="other") + async with asyncio.timeout(2): + for _ in range(3): + await provider.started.get() + await manager.stop(parent.id) + assert parent.state == "stopped" + assert child.state == other.state == "running" + await manager.stop(parent.id, tree=True) + assert child.state == "stopped" + assert other.state == "running" + provider.release.set() + await manager.wait(other.id) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_runtime_is_rebuilt_with_fresh_grants_and_cwd_bound_extras( + setup, checker_contract, monkeypatch +): + manager, ctx, _, _, _ = setup + ctx.session_perms.grant("bash", "*") + ctx.extras["registry"].unregister("bash") + captured = [] + original_run = AgentRunner.run + + async def run(runner, *args, **kwargs): + captured.append(runner.ctx) + return await original_run(runner, *args, **kwargs) + + monkeypatch.setattr(AgentRunner, "run", run) + try: + worker = await manager.start(ctx, agent="explore", prompt="read") + await manager.wait(worker.id) + child = captured[0] + assert child.cwd == ctx.cwd + assert child.session is worker.session + assert child.session_perms is not ctx.session_perms + assert child.session_perms.grants == [] + assert child.extras["background"] is not ctx.extras["background"] + assert child.extras["registry"] is not ctx.extras["registry"] + assert "bash" not in child.extras["registry"].names() + assert checker_contract.for_child.call_args_list[0].args == ( + ctx.extras["agents"].get("explore").overlay, + ) + assert checker_contract.for_child.call_args_list[0].kwargs == {"cwd": ctx.cwd} + assert checker_contract.for_child.call_args_list[-1].args == ( + ctx.extras["agents"].get("explore").overlay, + ) + assert checker_contract.for_child.call_args_list[-1].kwargs == { + "cwd": child.cwd, + "session_perms": child.session_perms, + "read_only": True, + } + finally: + await manager.shutdown() + + +@pytest.mark.skipif( + not hasattr(PermissionChecker, "for_child"), + reason="sibling-owned PermissionChecker.for_child has not landed", +) +@pytest.mark.asyncio +async def test_real_parent_checker_restrictions_reach_child_dispatch(setup): + manager, ctx, _, _, _ = setup + ctx.permission_checker = ctx.permission_checker.for_agent(AgentOverlay(denied_tools=("read",))) + ctx.extras["provider"] = FakeProvider( + [ + {"tool_calls": [{"id": "call", "name": "read", "arguments": '{"path":"file"}'}]}, + {"text": "denied as expected"}, + ] + ) + try: + worker = await manager.start(ctx, agent="explore", prompt="read") + await manager.wait(worker.id) + tool = ctx.extras["provider"].requests[1]["messages"][-1] + assert tool["role"] == "tool" + assert "denied" in tool["content"] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_usage_records_do_not_double_count_child_transcript(setup, checker_contract): + manager, ctx, _, store, session = setup + ctx.extras["provider"] = FakeProvider( + [ + { + "tool_calls": [{"id": "call", "name": "read", "arguments": '{"path":"missing"}'}], + "usage": {"input_tokens": 7, "output_tokens": 3, "cost_usd": 0.25}, + }, + {"text": "answer", "usage": {"input_tokens": 11, "output_tokens": 5, "cost_usd": 0.5}}, + ] + ) + try: + worker = await manager.start(ctx, agent="explore", prompt="question") + await manager.wait(worker.id) + assert worker.usage_totals.input_tokens == 18 + assert session_stats(store, session).input_tokens == 18 + assert session_stats(store, worker.session).input_tokens == 18 + assert session_stats(store, session).cost_usd == 0.75 + assert not worker.usage_incomplete + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_stopped_worker_does_not_return_a_stale_result(setup, checker_contract): + manager, ctx, _, _, _ = setup + try: + worker = await manager.start(ctx, agent="explore", prompt="question") + await manager.wait(worker.id) + await manager.stop(worker.id) + with pytest.raises(SubagentError, match="stopped"): + await manager.wait(worker.id) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_rejects_followup_before_writing_without_child_lock(setup): + manager, ctx, _, _, _ = setup + worker = await manager.start(ctx, agent="explore", prompt="question") + await manager.shutdown() + before = worker.session.path.read_text() + with pytest.raises(RuntimeError, match="shut down"): + await manager.send(worker.id, "late") + with pytest.raises(RuntimeError, match="shut down"): + await manager.resume(worker.id, "late") + assert worker.session.path.read_text() == before + + +@pytest.mark.asyncio +async def test_dirty_root_requires_human_confirmation_and_detached_is_refused(setup): + manager, ctx, _, _, _ = setup + git(ctx.cwd, "init", "-b", "main") + git( + ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--allow-empty", + "-m", + "initial", + ) + ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + } + ) + ctx.auto_approve = True + dirty = ctx.cwd / "uncommitted.txt" + dirty.write_text("not committed") + try: + with pytest.raises(WorktreeError, match="human confirmation"): + await manager.start(ctx, agent="writer", prompt="write") + confirmations = [] + + async def confirm(question): + confirmations.append(question) + return True + + manager.confirm = confirm + worker = await manager.start(ctx, agent="writer", prompt="write") + assert len(confirmations) == 1 + assert not (worker.cwd / dirty.name).exists() + assert dirty.read_text() == "not committed" + git(ctx.cwd, "checkout", "--detach") + with pytest.raises(WorktreeError, match="detached"): + await manager.start(ctx, agent="writer", prompt="write") + finally: + await manager.shutdown() + + +@pytest.mark.parametrize( + "agent_model,subagent_model,expected", + [ + ("agent-model", "subagent-model", "agent-model"), + (None, "subagent-model", "subagent-model"), + (None, None, "main-model"), + ], +) +@pytest.mark.asyncio +async def test_worker_model_precedence_is_recorded( + setup, checker_contract, agent_model, subagent_model, expected +): + manager, ctx, provider, _, _ = setup + ctx.config.llm.model = "main-model" + ctx.config.agent.subagent_model = subagent_model + ctx.extras["agents"] = AgentRegistry( + { + "explore": replace(ctx.extras["agents"].get("explore"), model=agent_model), + } + ) + try: + worker = await manager.start(ctx, agent="explore", prompt="question") + await manager.wait(worker.id) + assert provider.requests[0]["model"] == expected + assert worker.session.meta.model == expected + finally: + await manager.shutdown() diff --git a/tests/test_worktree.py b/tests/test_worktree.py index 169dd55..1e6f47a 100644 --- a/tests/test_worktree.py +++ b/tests/test_worktree.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio import subprocess from os.path import realpath @@ -34,6 +35,12 @@ async def commit_all(cwd, message): await git(cwd, *_COMMIT, "-m", message) +async def shell_validation(cmd, cwd): + """An approved test caller that intentionally runs validation commands.""" + result = await run_proc(["sh", "-lc", cmd], cwd=cwd, timeout=30) + return result.exit_code, result.stdout + result.stderr + + async def make_repo(path): """A git repo at ``path`` with one commit on ``main``.""" path.mkdir(parents=True, exist_ok=True) @@ -202,6 +209,387 @@ async def test_exit_deletes_branch(tmp_path): assert result.exit_code != 0 +# -- worker worktrees: create/attach/inspect/integrate/cleanup ----------------------- + + +async def test_discover_from_linked_worktree_returns_main_root(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + info = await manager.create("feat") + subdir = info.path / "a" / "b" + subdir.mkdir(parents=True) + found = await WorktreeManager.discover(subdir) + assert found.repo_root == repo + + +async def test_create_worker_pins_base_and_writes_sidecar(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + (repo / "later.txt").write_text("later\n", encoding="utf-8") + await commit_all(repo, "later") + destination_subdir = repo / "nested" + destination_subdir.mkdir() + info = await manager.create_worker( + "worker", base_commit=base, dest_path=destination_subdir, dest_branch="main" + ) + assert info.branch == "lecode/worker" + assert await git(info.path, "rev-parse", "HEAD") == base + assert not (info.path / "later.txt").exists() + assert manager.read_sidecar("worker") == { + "name": "worker", + "path": str(info.path), + "branch": "lecode/worker", + "base_commit": base, + "dest_path": str(repo), + "dest_branch": "main", + "dest_common_dir": str(repo / ".git"), + "integrated_at": None, + "integrated_head": None, + } + assert (repo / ".lecode" / "worktrees" / "worker.json").is_file() + + +async def test_attach_existing_worktree(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + created = await manager.create("feat") + assert await manager.attach("feat") == created + + +async def test_attach_missing_worktree_raises(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + with pytest.raises(WorktreeError, match="no such worktree"): + await manager.attach("ghost") + + +async def test_inspect_reports_state(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + info = await manager.create("feat") + state = await manager.inspect("feat") + assert state.present is True + assert state.dirty is False + assert state.merge_in_progress is False + assert state.sidecar is None + + (info.path / "dirty.txt").write_text("x\n", encoding="utf-8") + assert (await manager.inspect("feat")).dirty is True + + (info.path / "dirty.txt").unlink() + (info.path / "file.txt").write_text("worker\n", encoding="utf-8") + await commit_all(info.path, "worker edit") + (repo / "file.txt").write_text("main\n", encoding="utf-8") + await commit_all(repo, "main edit") + merge = await run_proc(["git", "merge", "--no-edit", "main"], cwd=info.path, timeout=30) + assert merge.exit_code != 0 + assert (await manager.inspect("feat")).merge_in_progress is True + + absent = await manager.inspect("ghost") + assert absent.present is False + assert absent.info.name == "ghost" + assert absent.dirty is False + assert absent.merge_in_progress is False + + +async def test_integrate_happy_path(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + + result = await manager.integrate( + "worker", validation=["test -f feature.txt"], validation_runner=shell_validation + ) + assert result.status == "integrated" + assert result.conflicts == [] + assert (repo / "feature.txt").read_text() == "feature\n" + sidecar = manager.read_sidecar("worker") + assert sidecar["integrated_at"] + assert sidecar["integrated_head"] == await git(repo, "rev-parse", "HEAD") + + +async def test_integrate_requires_validation_unless_explicitly_allowed(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + + denied = await manager.integrate("worker", validation=[]) + assert denied.status == "blocked" + assert denied.detail == "validation required" + allowed = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert allowed.status == "integrated" + + +async def test_integrate_refuses_worker_not_at_reviewed_head(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + + with pytest.raises(WorktreeError, match="differs from reviewed"): + await manager.integrate( + "worker", + validation=[], + allow_unvalidated=True, + reviewed_head=base, + ) + assert not (repo / "feature.txt").exists() + + +async def test_integrate_rejects_destination_changed_during_validation(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + + async def change_destination(_cmd, _cwd): + (repo / "moved.txt").write_text("moved\n", encoding="utf-8") + await commit_all(repo, "move destination") + return 0, "" + + result = await manager.integrate( + "worker", validation=["approved"], validation_runner=change_destination + ) + assert result.status == "paused" + assert result.detail == "destination changed during validation" + assert not (repo / "feature.txt").exists() + + +async def test_integrate_without_sidecar_pauses(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + await manager.create("plain") + result = await manager.integrate("plain", validation=[], allow_unvalidated=True) + assert result.status == "paused" + assert result.detail == "no destination recorded" + + +async def test_integrate_validation_failure_blocks_merge(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + + result = await manager.integrate( + "worker", validation=["echo boom >&2; exit 1"], validation_runner=shell_validation + ) + assert result.status == "validation_failed" + assert "echo boom" in result.detail + assert "boom" in result.validation_output + assert not (repo / "feature.txt").exists() + assert await git(repo, "rev-parse", "HEAD") == base + assert manager.read_sidecar("worker")["integrated_at"] is None + + +async def test_integrate_pauses_and_blocks(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + + (info.path / "dirty.txt").write_text("x\n", encoding="utf-8") + result = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert result.status == "blocked" + assert result.detail == "uncommitted changes" + (info.path / "dirty.txt").unlink() + + (repo / "uncommitted.txt").write_text("x\n", encoding="utf-8") + result = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert result.status == "paused" + assert result.detail == "destination has uncommitted changes" + (repo / "uncommitted.txt").unlink() + + await git(repo, "checkout", "-b", "side") + result = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert result.status == "paused" + assert result.detail == "destination is on 'side', expected 'main'" + + +async def test_integrate_replays_destination_before_validation(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + (repo / "later.txt").write_text("later\n", encoding="utf-8") + await commit_all(repo, "later") + dest_commit = await git(repo, "rev-parse", "HEAD") + + result = await manager.integrate( + "worker", validation=["test -f later.txt"], validation_runner=shell_validation + ) + assert result.status == "integrated" + ancestor = await run_proc( + ["git", "merge-base", "--is-ancestor", dest_commit, "lecode/worker"], + cwd=repo, + timeout=30, + ) + assert ancestor.exit_code == 0 + + +async def test_integrate_conflict_leaves_merge_in_progress(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "file.txt").write_text("worker\n", encoding="utf-8") + await commit_all(info.path, "worker edit") + (repo / "file.txt").write_text("main\n", encoding="utf-8") + await commit_all(repo, "main edit") + + result = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert result.status == "conflict" + assert result.conflicts == ["file.txt"] + assert (await manager.inspect("worker")).merge_in_progress is True + assert manager.read_sidecar("worker")["integrated_at"] is None + + +async def test_cleanup_worker_guards(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + with pytest.raises(WorktreeError, match="not integrated"): + await manager.cleanup_worker("worker") + + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + assert ( + await manager.integrate("worker", validation=[], allow_unvalidated=True) + ).status == "integrated" + (info.path / "dirty.txt").write_text("x\n", encoding="utf-8") + with pytest.raises(WorktreeError, match="uncommitted changes"): + await manager.cleanup_worker("worker") + + other = await manager.create_worker( + "other", base_commit=base, dest_path=repo, dest_branch="main" + ) + removed = await manager.cleanup_worker("other", discard=True) + assert removed == other + assert not other.path.exists() + branch = await run_proc( + ["git", "show-ref", "--verify", "refs/heads/lecode/other"], cwd=repo, timeout=30 + ) + assert branch.exit_code != 0 + + await manager.cleanup_worker("worker", discard=True) + assert not info.path.exists() + + +async def test_cleanup_worker_refuses_commit_after_integration(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + assert ( + await manager.integrate("worker", validation=[], allow_unvalidated=True) + ).status == "integrated" + (info.path / "after.txt").write_text("after\n", encoding="utf-8") + await commit_all(info.path, "post-integration commit") + + with pytest.raises(WorktreeError, match="not fully integrated"): + await manager.cleanup_worker("worker") + await manager.cleanup_worker("worker", discard=True) + + +async def test_concurrent_integrations_are_serialized(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + first = await manager.create_worker( + "first", base_commit=base, dest_path=repo, dest_branch="main" + ) + second = await manager.create_worker( + "second", base_commit=base, dest_path=repo, dest_branch="main" + ) + (first.path / "first.txt").write_text("first\n", encoding="utf-8") + await commit_all(first.path, "first") + (second.path / "second.txt").write_text("second\n", encoding="utf-8") + await commit_all(second.path, "second") + first_started = asyncio.Event() + release_first = asyncio.Event() + second_started = asyncio.Event() + + async def validate(cmd, _cwd): + if cmd == "first": + first_started.set() + await release_first.wait() + else: + second_started.set() + return 0, "" + + first_task = asyncio.create_task( + manager.integrate("first", validation=["first"], validation_runner=validate) + ) + await first_started.wait() + second_task = asyncio.create_task( + manager.integrate("second", validation=["second"], validation_runner=validate) + ) + await asyncio.sleep(0) + assert not second_started.is_set() + release_first.set() + assert (await first_task).status == "integrated" + assert (await second_task).status == "integrated" + assert (repo / "first.txt").is_file() + assert (repo / "second.txt").is_file() + + +async def test_paused_integrate_is_non_destructive(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + base = await git(repo, "rev-parse", "HEAD") + info = await manager.create_worker( + "worker", base_commit=base, dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") + await commit_all(info.path, "add feature") + (repo / "uncommitted.txt").write_text("keep me\n", encoding="utf-8") + + result = await manager.integrate("worker", validation=[], allow_unvalidated=True) + assert result.status == "paused" + assert (repo / "uncommitted.txt").read_text(encoding="utf-8") == "keep me\n" + assert await git(repo, "status", "--porcelain") == "?? uncommitted.txt" + assert await git(repo, "rev-parse", "HEAD") == base + assert not (repo / ".git" / "MERGE_HEAD").exists() + assert not (repo / "feature.txt").exists() + + # -- /worktree /wt-merge /wt-exit commands ------------------------------------------ From 9de35591a8735a7801557651ddc84c39b71b2a1e Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Tue, 15 Sep 2026 14:51:30 +0400 Subject: [PATCH 5/8] fix: close worker review gaps --- docs/configuration.md | 16 -- src/lecode/agent/tools/task.py | 7 +- src/lecode/agent/tools/workers.py | 6 +- src/lecode/config/models.py | 9 - src/lecode/extras/workers.py | 9 +- src/lecode/extras/worktree.py | 270 +------------------------- src/lecode/permission/checker.py | 36 ++-- src/lecode/tui/agents.py | 35 +++- src/lecode/tui/app.py | 70 ++++++- src/lecode/tui/permission.py | 10 +- tests/test_agent_builder.py | 11 ++ tests/test_config_loader.py | 13 -- tests/test_permission_checker.py | 18 ++ tests/test_subagents.py | 17 +- tests/test_tui_app.py | 43 +++++ tests/test_tui_permission.py | 25 +++ tests/test_workers.py | 43 ++++- tests/test_worktree.py | 306 +----------------------------- 18 files changed, 301 insertions(+), 643 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 0e0a60d..9c8f4e8 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -180,22 +180,6 @@ trouble never blocks the agent. See [memory.md](memory.md). -## `[worktree]` - -| field | default | meaning | -|---|---|---| -| `validation` | `[]` | commands a human-approved caller runs in the worker tree before fast-forward integration | - -For example: - -```toml -[worktree] -validation = ["uv run ruff check", "uv run python -m pytest"] -``` - -The worktree helper does not execute these commands itself. Its caller supplies -the approved command runner; an empty list requires explicit human authorization. - ## `[pierre]` Post-task reviewer: after every completed task, a second model compares the diff --git a/src/lecode/agent/tools/task.py b/src/lecode/agent/tools/task.py index 6df2c10..2c3b895 100644 --- a/src/lecode/agent/tools/task.py +++ b/src/lecode/agent/tools/task.py @@ -77,12 +77,7 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: agents = ctx.extras.get(AGENTS_EXTRA) if registry is None or agents is None: return ToolResult("error: subagents are unavailable in this context", is_error=True) - # The TUI still owns its transient roster through run_subagent. - manager = ( - None - if ctx.extras.get(SUBAGENT_EVENTS_EXTRA) is not None - else ctx.extras.get(WORKER_EXTRA) - ) + manager = ctx.extras.get(WORKER_EXTRA) if manager is not None: return await self._start_worker(args, ctx, manager, prompt) if args.get("run_in_background"): diff --git a/src/lecode/agent/tools/workers.py b/src/lecode/agent/tools/workers.py index dd37372..35e7a21 100644 --- a/src/lecode/agent/tools/workers.py +++ b/src/lecode/agent/tools/workers.py @@ -8,7 +8,7 @@ from lecode.extras.subagents import SubagentError from lecode.extras.workers import WORKER_CURRENT_EXTRA, WORKER_EXTRA -_ACTIONS = ("list", "send", "stop", "resume", "submit", "question", "integrate", "cleanup") +_ACTIONS = ("list", "send", "stop", "resume", "submit", "question") class WorkersTool(Tool): @@ -50,8 +50,6 @@ def _validate(self, args: dict[str, Any]) -> str | None: "resume": {"action", "id", "text"}, "submit": {"action", "id"}, "question": {"action", "text"}, - "integrate": {"action"}, - "cleanup": {"action"}, } if action not in _ACTIONS or set(args) - allowed[action]: return "invalid workers action or arguments" @@ -73,8 +71,6 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: if manager is None: return ToolResult("error: workers are unavailable in this context", is_error=True) action = args["action"] - if action in {"integrate", "cleanup"}: - return ToolResult(f"error: workers {action} is unavailable", is_error=True) current = ctx.extras.get(WORKER_CURRENT_EXTRA) if action == "list": workers = ( diff --git a/src/lecode/config/models.py b/src/lecode/config/models.py index 67ef166..31897eb 100644 --- a/src/lecode/config/models.py +++ b/src/lecode/config/models.py @@ -213,14 +213,6 @@ class MemoryConfig(BaseModel): max_bytes: int = 32768 -class WorktreeConfig(BaseModel): - """``[worktree]`` -- worker worktree integration checks.""" - - model_config = ConfigDict(extra="ignore") - - validation: list[str] = Field(default_factory=list) - - class PierreConfig(BaseModel): """``[pierre]`` — post-task reviewer: a second model compares the request with the result and gives feedback after every completed task.""" @@ -284,7 +276,6 @@ class Config(BaseModel): mcp: McpConfig = Field(default_factory=McpConfig) lsp: LspConfig = Field(default_factory=LspConfig) memory: MemoryConfig = Field(default_factory=MemoryConfig) - worktree: WorktreeConfig = Field(default_factory=WorktreeConfig) pierre: PierreConfig = Field(default_factory=PierreConfig) telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) hooks: dict[str, list[str]] = Field(default_factory=dict) diff --git a/src/lecode/extras/workers.py b/src/lecode/extras/workers.py index cfe7baf..8c816b5 100644 --- a/src/lecode/extras/workers.py +++ b/src/lecode/extras/workers.py @@ -20,7 +20,7 @@ from lecode.agent.runner import AgentRunner, LlmResponse, RunResult, UsageTotals from lecode.agent.tools.base import ToolContext from lecode.config.models import Config -from lecode.extras.subagents import SubagentError +from lecode.extras.subagents import SUBAGENT_EVENTS_EXTRA, SubagentError, SubagentProgress from lecode.extras.worktree import WorktreeError, WorktreeInfo, WorktreeManager from lecode.session.model import EventRecord, MessageRecord from lecode.session.storage import Session, SessionStore @@ -599,7 +599,7 @@ async def suspend(self, worker_id: str): worker.state = "running" self._record(worker) - def _event(self, worker, event): + async def _event(self, worker, event): if isinstance(event, LlmResponse): old = worker.usage_totals worker.usage_totals = UsageTotals( @@ -617,6 +617,11 @@ def _event(self, worker, event): }, ) self._record(worker) + callback = self.root_ctx.extras.get(SUBAGENT_EVENTS_EXTRA) + if callback is not None: + result = callback(SubagentProgress(worker.id, worker.agent, worker.description, event)) + if inspect.isawaitable(result): + await result async def wait(self, id: str) -> RunResult: worker = self.get(id) diff --git a/src/lecode/extras/worktree.py b/src/lecode/extras/worktree.py index fbe030a..b66bbad 100644 --- a/src/lecode/extras/worktree.py +++ b/src/lecode/extras/worktree.py @@ -13,26 +13,18 @@ reported — no auto-resolution, ``git merge --abort`` stays available. Worker worktrees: ``create_worker`` pins a base commit and a destination -(path + branch) in a sidecar under ``.lecode/worktrees/.json``; -``integrate`` validates the worker tree and merges it into that pinned -destination. ``discover`` returns the *main* repository root even when -called from inside a linked worktree, so worktrees and sidecars agree -across restarts. +(path + branch) in a sidecar under ``.lecode/worktrees/.json`` for a +future explicit integration workflow. ``discover`` returns the *main* +repository root even when called from inside a linked worktree, so worktrees +and sidecars agree across restarts. """ from __future__ import annotations -import asyncio -import fcntl -import hashlib -import inspect import json import os import re -from collections.abc import Callable -from contextlib import asynccontextmanager from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path from typing import Any @@ -41,9 +33,6 @@ #: Per-git-command timeout. GIT_TIMEOUT_S = 30.0 -#: Validation output kept in the returned :class:`IntegrationResult`. -_VALIDATION_OUTPUT_LIMIT = 4000 - #: Worktree directory, relative to the repo root. WORKTREE_ROOT = ".lecode/worktrees" @@ -53,13 +42,6 @@ _NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") -def _clip_validation(text: str) -> str: - """Keep validation output bounded; the tail usually holds the failure.""" - if len(text) <= _VALIDATION_OUTPUT_LIMIT: - return text - return "[… output clipped …]\n" + text[-_VALIDATION_OUTPUT_LIMIT:] - - class WorktreeError(Exception): """A git or worktree operation failed cleanly.""" @@ -103,20 +85,6 @@ class WorktreeInspection: sidecar: dict[str, Any] | None -@dataclass(frozen=True) -class IntegrationResult: - """The outcome of integrating a worker into its pinned destination. - - ``status`` is one of ``integrated``, ``paused``, ``blocked``, - ``conflict``, or ``validation_failed``. - """ - - status: str - detail: str - conflicts: list[str] - validation_output: str - - class WorktreeManager: """Worktree operations over one repository root.""" @@ -374,7 +342,7 @@ async def attach(self, name: str) -> WorktreeInfo: async def inspect(self, name: str) -> WorktreeInspection: """Best-effort worker state; absent worktrees never raise.""" info = self._info(name) - sidecar = self.read_sidecar(name) + sidecar = self._validated_sidecar(name, info) if not info.path.is_dir(): return WorktreeInspection( present=False, info=info, dirty=False, merge_in_progress=False, sidecar=sidecar @@ -440,234 +408,6 @@ async def _conflicted_files(self, cwd: Path | None = None) -> list[str]: ) return [line for line in result.stdout.splitlines() if line.strip()] - async def _destination_problem( - self, dest_path: Path, dest_branch: str, common_dir: Path - ) -> str | None: - """Check the pinned checkout without mutating it.""" - if not dest_path.is_dir(): - raise WorktreeError(f"destination missing: {dest_path}") - root = await run_proc( - ["git", "rev-parse", "--show-toplevel"], cwd=dest_path, timeout=GIT_TIMEOUT_S - ) - if root.exit_code != 0: - raise WorktreeError(f"destination is not a git repository: {dest_path}") - if Path(root.stdout.strip()).resolve() != dest_path: - raise WorktreeError(f"destination path was replaced: {dest_path}") - if await self._common_dir(dest_path) != common_dir: - raise WorktreeError(f"destination repository was replaced: {dest_path}") - ref = await run_proc( - ["git", "show-ref", "--verify", f"refs/heads/{dest_branch}"], - cwd=dest_path, - timeout=GIT_TIMEOUT_S, - ) - if ref.exit_code != 0: - raise WorktreeError(f"destination branch missing: {dest_branch}") - branch = await self._git("rev-parse", "--abbrev-ref", "HEAD", cwd=dest_path) - if branch != dest_branch: - return f"destination is on '{branch}', expected '{dest_branch}'" - if await self._git("status", "--porcelain", cwd=dest_path): - return "destination has uncommitted changes" - return None - - async def _worker_state(self, info: WorktreeInfo, common_dir: Path) -> tuple[str, bool]: - """Return the worker HEAD and dirty state after proving its identity.""" - if not info.path.is_dir(): - raise WorktreeError(f"worker path missing: {info.path}") - if await self._common_dir(info.path) != common_dir: - raise WorktreeError(f"worker repository was replaced: {info.path}") - ref = await run_proc( - ["git", "show-ref", "--verify", f"refs/heads/{info.branch}"], - cwd=info.path, - timeout=GIT_TIMEOUT_S, - ) - if ref.exit_code != 0: - raise WorktreeError(f"worker branch missing: {info.branch}") - branch = await self._git("rev-parse", "--abbrev-ref", "HEAD", cwd=info.path) - if branch != info.branch: - raise WorktreeError( - f"worker branch changed: expected '{info.branch}', found '{branch}'" - ) - return ( - await self._commit("HEAD", cwd=info.path), - bool(await self._git("status", "--porcelain", cwd=info.path)), - ) - - @asynccontextmanager - async def _integration_lock(self, common_dir: Path, dest_path: Path, dest_branch: str): - """Hold a persistent per-destination flock for the whole integration.""" - if not common_dir.is_dir(): - raise WorktreeError(f"destination git directory missing: {common_dir}") - digest = hashlib.sha256(f"{dest_path}\0{dest_branch}".encode()).hexdigest() - lock = (common_dir / f"lecode-integrate-{digest}.lock").open("a+") - try: - await asyncio.to_thread(fcntl.flock, lock.fileno(), fcntl.LOCK_EX) - yield - finally: - await asyncio.to_thread(fcntl.flock, lock.fileno(), fcntl.LOCK_UN) - lock.close() - - async def _is_ancestor(self, ancestor: str, descendant: str, *, cwd: Path) -> bool: - result = await run_proc( - ["git", "merge-base", "--is-ancestor", ancestor, descendant], - cwd=cwd, - timeout=GIT_TIMEOUT_S, - ) - return result.exit_code == 0 - - async def integrate( - self, - name: str, - *, - validation: list[str], - validation_runner: Callable[[str, Path], Any] | None = None, - allow_unvalidated: bool = False, - reviewed_head: str | None = None, - ) -> IntegrationResult: - """Validate a pinned candidate, then fast-forward its destination only.""" - if not validation and not allow_unvalidated: - return IntegrationResult("blocked", "validation required", [], "") - if validation and validation_runner is None: - raise WorktreeError("validation runner required") - if any(not isinstance(cmd, str) or not cmd for cmd in validation): - raise WorktreeError("validation commands must be nonempty strings") - - info = self._require(name) - sidecar = self._validated_sidecar(name, info) - if sidecar is None: - return IntegrationResult("paused", "no destination recorded", [], "") - dest_path = Path(sidecar["dest_path"]) - dest_branch = sidecar["dest_branch"] - common_dir = Path(sidecar["dest_common_dir"]) - if common_dir != await self._common_dir(self.repo_root): - raise WorktreeError("worker sidecar belongs to a different repository") - - async with self._integration_lock(common_dir, dest_path, dest_branch): - # Re-read after acquiring the stable lock so a replaced sidecar cannot retarget us. - if self._validated_sidecar(name, info) != sidecar: - raise WorktreeError(f"sidecar changed while integrating '{name}'") - problem = await self._destination_problem(dest_path, dest_branch, common_dir) - if problem is not None: - return IntegrationResult("paused", problem, [], "") - worker_head, worker_dirty = await self._worker_state(info, common_dir) - if worker_dirty: - return IntegrationResult("blocked", "uncommitted changes", [], "") - if reviewed_head is not None and worker_head != reviewed_head: - raise WorktreeError("worker head differs from reviewed head") - base_commit = await self._commit(sidecar["base_commit"], cwd=info.path) - if not await self._is_ancestor(base_commit, worker_head, cwd=info.path): - raise WorktreeError("worker branch no longer descends from its pinned base") - dest_head = await self._commit("HEAD", cwd=dest_path) - - if not await self._is_ancestor(dest_head, worker_head, cwd=info.path): - merge = await run_proc( - ["git", "merge", "--no-edit", dest_head], cwd=info.path, timeout=GIT_TIMEOUT_S - ) - if merge.exit_code != 0: - conflicts = await self._conflicted_files(cwd=info.path) - if conflicts: - return IntegrationResult( - "conflict", - f"conflicts merging {dest_branch} into {info.branch}", - conflicts, - "", - ) - return IntegrationResult( - "blocked", merge.stderr.strip() or "merge failed", [], "" - ) - candidate, worker_dirty = await self._worker_state(info, common_dir) - if worker_dirty: - return IntegrationResult("blocked", "uncommitted changes", [], "") - - validation_output = "" - for cmd in validation: - result = validation_runner(cmd, info.path) # type: ignore[misc] - if inspect.isawaitable(result): - result = await result - try: - exit_code, output = result - except (TypeError, ValueError) as error: - raise WorktreeError( - "validation runner must return (exit_code, output)" - ) from error - if not isinstance(exit_code, int) or not isinstance(output, str): - raise WorktreeError("validation runner must return (int, str)") - validation_output += output - if exit_code != 0: - return IntegrationResult( - "validation_failed", - f"validation failed: {cmd}", - [], - _clip_validation(validation_output), - ) - validation_output = _clip_validation(validation_output) - - problem = await self._destination_problem(dest_path, dest_branch, common_dir) - if problem is not None: - return IntegrationResult("paused", problem, [], validation_output) - if await self._commit("HEAD", cwd=dest_path) != dest_head: - return IntegrationResult( - "paused", "destination changed during validation", [], validation_output - ) - current_worker, worker_dirty = await self._worker_state(info, common_dir) - if worker_dirty or current_worker != candidate: - return IntegrationResult( - "blocked", "worker changed during validation", [], validation_output - ) - if not await self._is_ancestor(dest_head, candidate, cwd=info.path): - raise WorktreeError("validated candidate does not contain the pinned destination") - - merge = await run_proc( - ["git", "merge", "--ff-only", candidate], cwd=dest_path, timeout=GIT_TIMEOUT_S - ) - if merge.exit_code != 0: - return IntegrationResult( - "blocked", merge.stderr.strip() or "fast-forward failed", [], validation_output - ) - integrated_head = await self._commit("HEAD", cwd=dest_path) - if integrated_head != candidate: - return IntegrationResult( - "blocked", "destination changed while fast-forwarding", [], validation_output - ) - sidecar["integrated_at"] = datetime.now(UTC).isoformat() - sidecar["integrated_head"] = integrated_head - self.write_sidecar(name, sidecar) - return IntegrationResult( - "integrated", f"fast-forwarded {dest_branch} to {candidate}", [], validation_output - ) - - async def cleanup_worker(self, name: str, *, discard: bool = False) -> WorktreeInfo: - """Remove an integrated worker, or explicitly discard one.""" - info = self._require(name) - if discard: - await self._git("worktree", "remove", "--force", str(info.path)) - await self._git("branch", "-D", info.branch) - return info - - sidecar = self._validated_sidecar(name, info) - if sidecar is None or sidecar.get("integrated_at") is None: - raise WorktreeError( - f"worktree '{name}' is not integrated (use discard=True to discard)" - ) - dest_path = Path(sidecar["dest_path"]) - dest_branch = sidecar["dest_branch"] - common_dir = Path(sidecar["dest_common_dir"]) - problem = await self._destination_problem(dest_path, dest_branch, common_dir) - if problem is not None: - raise WorktreeError(problem) - _worker_head, worker_dirty = await self._worker_state(info, common_dir) - if worker_dirty: - raise WorktreeError( - f"worktree '{name}' has uncommitted changes (use discard=True to discard)" - ) - dest_head = await self._commit("HEAD", cwd=dest_path) - if not await self._is_ancestor(info.branch, dest_head, cwd=dest_path): - raise WorktreeError( - f"worktree '{name}' is not fully integrated (use discard=True to discard)" - ) - await self._git("worktree", "remove", str(info.path)) - await self._git("branch", "-d", info.branch, cwd=dest_path) - return info - async def exit_worktree( self, name: str, *, delete_branch: bool = False, force: bool = False ) -> WorktreeInfo: diff --git a/src/lecode/permission/checker.py b/src/lecode/permission/checker.py index eba8892..8650cd1 100644 --- a/src/lecode/permission/checker.py +++ b/src/lecode/permission/checker.py @@ -251,23 +251,29 @@ def check( ): absolute = os.path.abspath(self._cwd / target) targets = (target, absolute, os.path.relpath(absolute, self._cwd)) - result = self._policy_decision(tool_name, targets) + result = self._policy_decision(tool_name, args, targets) return self._apply_doom_loop(tool_name, args, result) # -- pipeline steps ------------------------------------------------------ - def _policy_decision(self, tool_name: str, targets: tuple[str, ...]) -> CheckResult: + def _policy_decision( + self, tool_name: str, args: dict[str, Any], targets: tuple[str, ...] + ) -> CheckResult: """Intersect every policy layer without recording an ancestor call.""" - results = [self._base_decision(tool_name, targets, None)] + results = [self._base_decision(tool_name, args, targets, None)] if self._overlay is not None: - results.append(self._base_decision(tool_name, targets, self._overlay)) + results.append(self._base_decision(tool_name, args, targets, self._overlay)) if self._parent is not None: - results.append(self._parent._policy_decision(tool_name, targets)) + results.append(self._parent._policy_decision(tool_name, args, targets)) priority = {Decision.ALLOW: 0, Decision.ASK: 1, Decision.DENY: 2} return max(results, key=lambda result: priority[result.decision]) def _base_decision( - self, tool_name: str, targets: tuple[str, ...], overlay: AgentOverlay | None + self, + tool_name: str, + args: dict[str, Any], + targets: tuple[str, ...], + overlay: AgentOverlay | None, ) -> CheckResult: # Deny rules are unbypassable: global table + overlay extras. deny = self._last_match(self._rules.deny, tool_name, targets) @@ -282,7 +288,7 @@ def _base_decision( # Read-only checkers can never be widened by overlay rules, global # rules, session grants, or the mode fallback. - if self._read_only and not self._is_read_class(tool_name): + if self._read_only and not self._is_read_class(tool_name, args): return CheckResult(Decision.DENY, f"read-only: {tool_name} is not a read-class tool") # Overlay extra allow/ask rules first (last match wins within them). @@ -304,7 +310,7 @@ def _base_decision( # Mode fallback. mode = overlay.mode if overlay and overlay.mode else self._mode - return self._mode_fallback(mode, tool_name) + return self._mode_fallback(mode, tool_name, args) def _last_match( self, table: dict[str, list[PermissionRule]], tool_name: str, targets: tuple[str, ...] @@ -345,13 +351,19 @@ def _apply_doom_loop( # -- mode fallback --------------------------------------------------------- - def _is_read_class(self, tool_name: str) -> bool: - return tool_name in READ_TOOLS or is_read_equiv_mcp(tool_name) + def _is_read_class(self, tool_name: str, args: dict[str, Any] | None = None) -> bool: + return ( + tool_name in READ_TOOLS + or is_read_equiv_mcp(tool_name) + or (tool_name == "workers" and args is not None and args.get("action") == "question") + ) - def _mode_fallback(self, mode: PermissionMode, tool_name: str) -> CheckResult: + def _mode_fallback( + self, mode: PermissionMode, tool_name: str, args: dict[str, Any] + ) -> CheckResult: reason = f"mode: {mode}" if mode == "yolo": return CheckResult(Decision.ALLOW, reason) # readonly: read-class tools are allowed, everything else is denied. - decision = Decision.ALLOW if self._is_read_class(tool_name) else Decision.DENY + decision = Decision.ALLOW if self._is_read_class(tool_name, args) else Decision.DENY return CheckResult(decision, reason) diff --git a/src/lecode/tui/agents.py b/src/lecode/tui/agents.py index 60217cc..2dc6d28 100644 --- a/src/lecode/tui/agents.py +++ b/src/lecode/tui/agents.py @@ -104,6 +104,7 @@ class AgentRun: elapsed_s: float = 0.0 subtree_cost_usd: float = 0.0 subtree_usage_incomplete: bool = False + session_id: str | None = None @property def current(self) -> str: @@ -175,12 +176,17 @@ def sync_worker(self, worker, *, context_window: int = 0) -> AgentRun: self._runs[run.run_id] = run self._order.append(run.run_id) run.status = {"completed": "done", "failed": "error"}.get(worker.state, worker.state) + run.parent_id = worker.parent_id + run.depth = worker.depth + run.worker = True + run.origin = worker.origin run.error = worker.error or "" run.answer = worker.result.final_text if worker.result is not None else "" run.cost_usd = worker.usage_totals.cost_usd run.usage_incomplete = worker.usage_incomplete run.context_used = worker.usage_totals.context_tokens run.context_window = context_window + run.session_id = getattr(worker, "session_id", None) started_at = getattr(worker, "started_at", "") now = datetime.now(UTC) with suppress(TypeError, ValueError): @@ -335,7 +341,29 @@ def _preview(text: str, width: int, limit: int = RESULT_PREVIEW_LINES) -> list[s return shown -def detail_lines(run: AgentRun | None, theme: Theme, width: int) -> list[Text]: +def _message_lines(message: dict, theme: Theme) -> list[Text]: + role = str(message.get("role", "unknown")) + name = str(message.get("name") or "") + label = f" [{role}{f' {name}' if name else ''}]" + lines = [Text(label, style=theme.muted)] + content = message.get("content") + if content not in (None, ""): + text = content if isinstance(content, str) else json.dumps(content, ensure_ascii=True) + lines.extend(Text(f" {row}", style=theme.text) for row in text.splitlines() or [""]) + for call in message.get("tool_calls") or []: + function = call.get("function", call) + name = str(function.get("name", "tool")) + args = str(function.get("arguments", "")) + lines.append(Text(f" ⚙ {name} {args}", style=theme.tool)) + return lines + + +def detail_lines( + run: AgentRun | None, + theme: Theme, + width: int, + transcript: list[dict] | None = None, +) -> list[Text]: """The per-run detail panel: identity, tool trail, answer/error.""" if run is None: return [] @@ -362,6 +390,11 @@ def detail_lines(run: AgentRun | None, theme: Theme, width: int) -> list[Text]: style=theme.muted, ) ) + if transcript is not None: + lines.append(Text(" transcript:", style=theme.muted)) + for message in transcript: + lines.extend(_message_lines(message, theme)) + return lines if not run.activity: lines.append(Text(" (no tool calls yet)", style=theme.muted)) for entry in run.activity: diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index 080ffd4..6c50cd2 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -593,7 +593,10 @@ async def _on_worker_notification(self, note: dict[str, Any]) -> None: if worker is None: return run = self._roster.sync_worker(worker, context_window=self._status.context_window) - self._feed.agent_summary(run) + if note.get("kind") == "human_message": + self._feed.info(f"[human → @{worker.agent} {worker.id[:8]}] {note['content']}") + else: + self._feed.agent_summary(run) if ( (note["origin"] == "delegated" or note["deliver"]) and note["parent_id"] is None @@ -745,6 +748,8 @@ def _approve_once(event: Any) -> None: @kb.add("a", filter=approval_pending) def _approve_always(event: Any) -> None: pending = self._approval.pending + if pending is not None and not pending.allow_always: + return pattern = pending.target if pending is not None else "*" self._resolve_approval(AllowAlways(pattern=pattern)) @@ -1231,6 +1236,8 @@ async def run(self, *, input: Input | None = None, output: Output | None = None) self._app = self._build_app(input=input, output=output) self._runtime.ctx.approval_callback = self._request_approval self._runtime.ctx.question_callback = self._request_question + if self._worker_manager is not None: + self._worker_manager.confirm = self._confirm_worker_worktree await self._fire_hook(SESSION_START) # MCP attaches in the background so the chat opens immediately; # per-server status lands in the feed when the connect finishes. @@ -1254,6 +1261,8 @@ async def run(self, *, input: Input | None = None, output: Output | None = None) self._spinner_task.cancel() self._approval.cancel() self._runtime.ctx.approval_callback = None + if self._worker_manager is not None: + self._worker_manager.confirm = None self._question.cancel() self._runtime.ctx.question_callback = None self.cancel_turn() @@ -1593,7 +1602,12 @@ def _show_approval_head(self) -> None: attribution = ( f"[worker {pending.worker[:8]} · {pending.conversation}] " if pending.worker else "" ) - self._feed.permission(attribution + approval_prompt_text(pending.tool_name, pending.target)) + self._feed.permission( + attribution + + approval_prompt_text( + pending.tool_name, pending.target, allow_always=pending.allow_always + ) + ) def _resolve_approval(self, decision: ApprovalDecision) -> None: self._approval.resolve(decision) @@ -1635,6 +1649,48 @@ async def _request_approval( ) self._invalidate() + async def _confirm_worker_worktree(self, *args: Any, **kwargs: Any) -> bool: + """Ask the TUI user before a dirty write worker gets a separate worktree.""" + question = str( + kwargs.get("question") or next((arg for arg in args if isinstance(arg, str)), "") + ) + worker = kwargs.get("worker") + worktree = kwargs.get("worktree") or kwargs.get("cwd") + for arg in args: + if not isinstance(arg, str) and worker is None: + worker = arg + elif not isinstance(arg, str) and worktree is None: + worktree = arg + worker_id = getattr(worker, "id", None) or getattr(worker, "worker_id", None) or "new" + agent = getattr(worker, "agent", None) or kwargs.get("agent") or "write" + path = getattr(worktree, "path", worktree) or self._cwd + target = f"@{agent} worker {str(worker_id)[:8]} in {path}" + future = self._approval.request( + "dirty worktree", + target, + question, + worker=str(worker_id), + conversation=str(path), + allow_always=False, + ) + self._show_approval_head() + if question: + self._feed.info(question) + self._status.state = StatusLineState.AWAITING_APPROVAL + self._invalidate() + try: + return isinstance(await future, AllowOnce) + finally: + self._approval.cancel(future) + self._shown_approval = None + self._show_approval_head() + self._status.state = ( + StatusLineState.AWAITING_APPROVAL + if self._approval.is_pending + else StatusLineState.RUNNING + ) + self._invalidate() + def _render_question(self) -> None: """Record the current question in the feed; its options live in the picker panel.""" question = self._question.current() @@ -2046,7 +2102,15 @@ def _roster_text(self) -> ANSI: fragments a ``FormattedTextControl`` caches.""" width = self._term_width() if self._detail_run_id is not None: - lines = detail_lines(self._roster.get(self._detail_run_id), self._theme, width) + run = self._roster.get(self._detail_run_id) + transcript = None + if run is not None and run.session_id is not None: + with contextlib.suppress(Exception): + child = self._store.open(run.session_id) + transcript = [ + record.message for record in self._store.load_messages(child) + ] or None + lines = detail_lines(run, self._theme, width, transcript) else: lines = roster_lines(self._roster, self._theme, width) with self._console.capture() as capture: diff --git a/src/lecode/tui/permission.py b/src/lecode/tui/permission.py index 30fc8d7..840f0b0 100644 --- a/src/lecode/tui/permission.py +++ b/src/lecode/tui/permission.py @@ -31,11 +31,14 @@ class PendingApproval: #: Attribution for "[w2 bash]" style prompts; None = the main turn. worker: str | None = None conversation: str = "main" + allow_always: bool = True -def approval_prompt_text(tool_name: str, target: str) -> str: +def approval_prompt_text(tool_name: str, target: str, *, allow_always: bool = True) -> str: """The one-line ask: ``allow bash 'ls'? (y)once (a)lways (n)deny — ESC denies``.""" shown = target if len(target) <= _TARGET_MAX_LEN else target[: _TARGET_MAX_LEN - 1] + "…" + if not allow_always: + return f"confirm {tool_name} '{shown}'? (y)es (n)o — ESC denies" return f"allow {tool_name} '{shown}'? (y)once (a)lways (n)deny — ESC denies" @@ -61,9 +64,12 @@ def request( *, worker: str | None = None, conversation: str = "main", + allow_always: bool = True, ) -> asyncio.Future[ApprovalDecision]: future: asyncio.Future[ApprovalDecision] = asyncio.get_running_loop().create_future() - entry = PendingApproval(tool_name, target, reason, future, worker, conversation) + entry = PendingApproval( + tool_name, target, reason, future, worker, conversation, allow_always + ) self._queue.append(entry) future.add_done_callback(partial(self._on_future_done, entry)) return future diff --git a/tests/test_agent_builder.py b/tests/test_agent_builder.py index 14c3ba5..4ec947d 100644 --- a/tests/test_agent_builder.py +++ b/tests/test_agent_builder.py @@ -78,6 +78,17 @@ def test_session_runtime_installs_workers(cwd, tmp_path): assert runtime.ctx.extras["workers"].session is session +def test_workers_schema_omits_unimplemented_integration_actions(cwd, tmp_path): + from lecode.session.storage import SessionStore + + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("workers", cwd) + runtime = build_runtime(Config(), cwd, session=session, store=store) + actions = runtime.registry.get("workers").parameters["properties"]["action"]["enum"] + assert "integrate" not in actions + assert "cleanup" not in actions + + def test_agent_name_applies_overlay(cwd): runtime = build_runtime(Config(), cwd, agent_name="plan") checker = runtime.ctx.permission_checker diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 483a9f2..5bd70d0 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -5,7 +5,6 @@ import json import pytest -from pydantic import ValidationError from lecode.config.loader import deep_merge, load_config from lecode.config.migrations import MIGRATIONS @@ -48,7 +47,6 @@ def test_defaults_validate_from_empty(): assert config.mcp.enable_exa is True assert config.mcp.enable_context7 is False assert config.memory.max_bytes == 32768 - assert config.worktree.validation == [] assert config.telemetry.enabled is False assert config.telemetry.sentry_dsn is None assert config.telemetry.otlp_endpoint is None @@ -74,17 +72,6 @@ def test_telemetry_section_parses(global_dir, tmp_path): assert result.warnings == [] -def test_worktree_validation_parses_and_requires_a_list(global_dir, tmp_path): - (global_dir / "config.toml").write_text( - '[worktree]\nvalidation = ["uv run pytest tests/test_worktree.py"]\n' - ) - assert load_config(cwd=tmp_path).config.worktree.validation == [ - "uv run pytest tests/test_worktree.py" - ] - with pytest.raises(ValidationError): - Config.model_validate({"worktree": {"validation": "pytest"}}) - - def test_first_run_creates_default_config(global_dir, tmp_path): result = load_config(cwd=tmp_path) created = global_dir / "config.toml" diff --git a/tests/test_permission_checker.py b/tests/test_permission_checker.py index 3a26fcd..41183e3 100644 --- a/tests/test_permission_checker.py +++ b/tests/test_permission_checker.py @@ -289,6 +289,24 @@ def test_read_only_denies_writes_even_in_yolo(): assert "read-only" in checker.check("bash", {"command": "ls"}).reason +@pytest.mark.parametrize( + ("args", "expected"), + [ + ({"action": "question", "text": "Need a choice"}, Decision.ALLOW), + ({"action": "list"}, Decision.DENY), + ({"action": "send", "id": "w", "text": "continue"}, Decision.DENY), + ({"action": "stop", "id": "w"}, Decision.DENY), + ({"action": "resume", "id": "w"}, Decision.DENY), + ({"action": "submit", "id": "w"}, Decision.DENY), + ({"action": "integrate"}, Decision.DENY), + ({"action": "cleanup"}, Decision.DENY), + ], +) +def test_strict_readonly_allows_only_workers_question(args, expected): + checker = _checker({"mode": "yolo"}, read_only=True) + assert checker.check("workers", args).decision == expected + + def test_read_only_not_widened_by_overlay_allow_rule(): overlay = AgentOverlay(extra_rules=_ruleset(allow={"bash": [{"pattern": "*"}]})) checker = _checker({"mode": "yolo"}, read_only=True).for_agent(overlay) diff --git a/tests/test_subagents.py b/tests/test_subagents.py index 38d7a8e..cfb797f 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -31,6 +31,7 @@ def make_runtime(tmp_path, monkeypatch, provider, config=None): """A built runtime with the provider seam installed on the context.""" monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) config = config or Config() runtime = build_runtime(config, tmp_path) runtime.ctx.extras["provider"] = provider @@ -105,7 +106,9 @@ def __init__(self, parent_entry) -> None: def stream_chat(self, messages, model, tools=None, **kwargs): self.requests.append({"messages": [dict(m) for m in messages], "model": model}) - is_child = not any(t["function"]["name"] == "task" for t in tools or []) + is_child = not any(t["function"]["name"] == "task" for t in tools or []) or ( + getattr(self, "worker_requests_are_children", False) and len(self.requests) > 1 + ) if is_child: return self._child_stream() return _scripted_stream(self.parent_entry) @@ -536,7 +539,8 @@ async def test_child_events_feed_roster_and_one_summary_line(tmp_path, monkeypat assert len(runs) == 1 run = runs[0] assert run.description == "Scan repo" - assert run.status == "ok" + assert run.status == "done" + assert run.worker assert run.answer == "scan result" assert [entry.name for entry in run.activity] == ["list_dir"] @@ -595,6 +599,7 @@ async def test_roster_panel_visible_while_child_runs(tmp_path, monkeypatch): ] } ) + provider.worker_requests_are_children = True app, _, _ = make_app(tmp_path, monkeypatch, []) app._runner.provider = provider app._runtime.ctx.extras["provider"] = provider @@ -680,7 +685,10 @@ async def test_detail_panel_preserves_draft(tmp_path, monkeypatch): async def test_at_agent_runs_directly(tmp_path, monkeypatch): - script = [{"text": "42 files", "usage": {"input_tokens": 7, "output_tokens": 3}}] + script = [ + {"text": "42 files", "usage": {"input_tokens": 7, "output_tokens": 3}}, + {"text": "parent answer"}, + ] app, provider, out = make_app(tmp_path, monkeypatch, script) await app._submit("@explore count the files") task = app._turn_task @@ -699,6 +707,9 @@ async def test_at_agent_runs_directly(tmp_path, monkeypatch): assert app.store.load_agent_runs(app.session) == [] assert "42 files" not in rendered # completion notifies until /agent submit assert app._status.input_tokens == 7 + await app._submit("continue as parent") + await app._turn_task + assert "parent answer" in out.getvalue() async def test_at_primary_mention_degrades_to_note(tmp_path, monkeypatch): diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index 308922e..cd1cc52 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -7,6 +7,7 @@ from typing import Any, ClassVar import pytest +from prompt_toolkit.formatted_text import to_formatted_text from prompt_toolkit.input import create_pipe_input from prompt_toolkit.output import DummyOutput from rich.console import Console @@ -26,6 +27,8 @@ def make_app(tmp_path, monkeypatch, script, config=None): """A TuiApp over a FakeProvider with a recorded console.""" monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) config = config or Config() config.notifications.enabled = False # never play sounds in tests store = SessionStore() @@ -167,6 +170,44 @@ async def test_worker_focus_preserves_drafts_and_routes_composer(tmp_path, monke message["content"] == "follow up" for message in app.store.load_for_model(worker.session) ) + notes = app.store.load_events(app.session, "worker_notification") + assert any( + note.get("kind") == "human_message" + and note["worker_id"] == worker.id + and note["content"] == "follow up" + for note in notes + ) + await manager.shutdown() + + +async def test_worker_detail_renders_persisted_child_transcript(tmp_path, monkeypatch): + """Worker detail reads the child session, not just the bounded live trail.""" + (tmp_path / "note.txt").write_text("persisted tool result", encoding="utf-8") + app, _, _ = make_app( + tmp_path, + monkeypatch, + [ + {"tool_calls": [{"name": "read", "arguments": '{"path": "note.txt"}'}]}, + {"text": "persisted child answer"}, + ], + ) + manager = app.worker_manager + assert manager is not None + worker = await manager.start( + app.runtime.ctx, + agent="explore", + prompt="inspect the child transcript", + origin="human", + background=True, + ) + await manager.wait(worker.id) + + assert app.open_agent_run(worker.id) + detail = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) + assert "inspect the child transcript" in detail + assert "read" in detail + assert "persisted tool result" in detail + assert "persisted child answer" in detail await manager.shutdown() @@ -482,6 +523,8 @@ async def test_slash_menu_mid_message_does_not_open(tmp_path, monkeypatch): async def test_resume_restores_status_usage(tmp_path, monkeypatch): """A session with stored usage opens with the statusline pre-filled.""" monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) config = Config() config.notifications.enabled = False store = SessionStore() diff --git a/tests/test_tui_permission.py b/tests/test_tui_permission.py index e7d07e8..0753f86 100644 --- a/tests/test_tui_permission.py +++ b/tests/test_tui_permission.py @@ -47,6 +47,8 @@ def make_ctx(tmp_path, monkeypatch, callback=None, mode="yolo", tool_name="bash" rule for ``tool_name`` to drive the approval-prompt flows. """ monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) config = Config() config.notifications.enabled = False # never play sounds in tests if ask: @@ -296,6 +298,8 @@ async def test_shutdown_cancels_all_outstanding(): def make_app(tmp_path, monkeypatch, script): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) config = Config() config.notifications.enabled = False # never play sounds in tests # the two modes never Ask by themselves; gate bash with an ask rule @@ -360,6 +364,27 @@ async def test_worker_approval_is_attributed_at_fifo_head(tmp_path, monkeypatch) assert "[worker worker-1 · w] allow bash 'ls'?" in out.getvalue() +async def test_dirty_worker_confirmation_is_one_shot(tmp_path, monkeypatch): + from types import SimpleNamespace + + app, out = make_app(tmp_path, monkeypatch, []) + task = asyncio.ensure_future( + app._confirm_worker_worktree( + "Uncommitted changes will not enter the worker's committed-HEAD worktree. Continue?", + worker=SimpleNamespace(id="worker-1234", agent="build"), + worktree=tmp_path / "worker-tree", + ) + ) + await wait_for(lambda: app._approval.pending is not None) + pending = app._approval.pending + assert pending is not None and pending.allow_always is False + assert "@build worker worker-1" in out.getvalue() + assert "worker-tree" in out.getvalue() + assert "(y)es (n)o" in out.getvalue() + app._resolve_approval(AllowAlways(pattern="*")) + assert await task is False + + async def test_pipe_approval_y_runs_asked_tool(tmp_path, monkeypatch): """Full flow: model calls bash, user answers 'y', output appears.""" script = [ diff --git a/tests/test_workers.py b/tests/test_workers.py index 9ac2311..bf59d0a 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -9,7 +9,15 @@ from tests.fakes import FakeProvider from lecode.agent.builder import build_runtime -from lecode.agent.runner import AgentRunner, RunResult, UsageTotals +from lecode.agent.runner import ( + AgentRunner, + LlmCall, + LlmResponse, + RunResult, + Token, + UsageTotals, +) +from lecode.agent.runner import Done as RunnerDone from lecode.config.models import Config from lecode.context.agents import AgentDefinition, AgentRegistry from lecode.extras.subagents import SubagentError @@ -148,6 +156,39 @@ async def test_background_task_tool_delivers_worker_notification(setup, checker_ await manager.shutdown() +@pytest.mark.asyncio +async def test_task_uses_worker_manager_when_tui_events_are_installed(setup, checker_contract): + manager, ctx, _, _, _ = setup + ctx.extras["subagent_events"] = lambda event: None + try: + _, result = await ctx.extras["registry"].dispatch_result( + "task", "task", '{"prompt":"scan"}', ctx + ) + assert not result.is_error + assert result.metadata["worker_id"] == manager.list()[0].id + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_worker_forwards_all_runner_events_to_tui_callback(setup, checker_contract): + manager, ctx, _, _, _ = setup + seen = [] + ctx.extras["subagent_events"] = seen.append + try: + worker = await manager.start(ctx, agent="explore", prompt="scan") + await manager.wait(worker.id) + assert {type(progress.event) for progress in seen} == { + LlmCall, + Token, + LlmResponse, + RunnerDone, + } + assert {progress.run_id for progress in seen} == {worker.id} + finally: + await manager.shutdown() + + @pytest.mark.asyncio async def test_parent_cancellation_does_not_cancel_managed_worker(setup, checker_contract): manager, ctx, _, _, _ = setup diff --git a/tests/test_worktree.py b/tests/test_worktree.py index 1e6f47a..5abd936 100644 --- a/tests/test_worktree.py +++ b/tests/test_worktree.py @@ -7,7 +7,6 @@ from __future__ import annotations -import asyncio import subprocess from os.path import realpath @@ -35,12 +34,6 @@ async def commit_all(cwd, message): await git(cwd, *_COMMIT, "-m", message) -async def shell_validation(cmd, cwd): - """An approved test caller that intentionally runs validation commands.""" - result = await run_proc(["sh", "-lc", cmd], cwd=cwd, timeout=30) - return result.exit_code, result.stdout + result.stderr - - async def make_repo(path): """A git repo at ``path`` with one commit on ``main``.""" path.mkdir(parents=True, exist_ok=True) @@ -209,7 +202,7 @@ async def test_exit_deletes_branch(tmp_path): assert result.exit_code != 0 -# -- worker worktrees: create/attach/inspect/integrate/cleanup ----------------------- +# -- worker worktrees: create/attach/inspect ------------------------------------------ async def test_discover_from_linked_worktree_returns_main_root(tmp_path): @@ -293,303 +286,6 @@ async def test_inspect_reports_state(tmp_path): assert absent.merge_in_progress is False -async def test_integrate_happy_path(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - - result = await manager.integrate( - "worker", validation=["test -f feature.txt"], validation_runner=shell_validation - ) - assert result.status == "integrated" - assert result.conflicts == [] - assert (repo / "feature.txt").read_text() == "feature\n" - sidecar = manager.read_sidecar("worker") - assert sidecar["integrated_at"] - assert sidecar["integrated_head"] == await git(repo, "rev-parse", "HEAD") - - -async def test_integrate_requires_validation_unless_explicitly_allowed(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - - denied = await manager.integrate("worker", validation=[]) - assert denied.status == "blocked" - assert denied.detail == "validation required" - allowed = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert allowed.status == "integrated" - - -async def test_integrate_refuses_worker_not_at_reviewed_head(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - - with pytest.raises(WorktreeError, match="differs from reviewed"): - await manager.integrate( - "worker", - validation=[], - allow_unvalidated=True, - reviewed_head=base, - ) - assert not (repo / "feature.txt").exists() - - -async def test_integrate_rejects_destination_changed_during_validation(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - - async def change_destination(_cmd, _cwd): - (repo / "moved.txt").write_text("moved\n", encoding="utf-8") - await commit_all(repo, "move destination") - return 0, "" - - result = await manager.integrate( - "worker", validation=["approved"], validation_runner=change_destination - ) - assert result.status == "paused" - assert result.detail == "destination changed during validation" - assert not (repo / "feature.txt").exists() - - -async def test_integrate_without_sidecar_pauses(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - await manager.create("plain") - result = await manager.integrate("plain", validation=[], allow_unvalidated=True) - assert result.status == "paused" - assert result.detail == "no destination recorded" - - -async def test_integrate_validation_failure_blocks_merge(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - - result = await manager.integrate( - "worker", validation=["echo boom >&2; exit 1"], validation_runner=shell_validation - ) - assert result.status == "validation_failed" - assert "echo boom" in result.detail - assert "boom" in result.validation_output - assert not (repo / "feature.txt").exists() - assert await git(repo, "rev-parse", "HEAD") == base - assert manager.read_sidecar("worker")["integrated_at"] is None - - -async def test_integrate_pauses_and_blocks(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - - (info.path / "dirty.txt").write_text("x\n", encoding="utf-8") - result = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert result.status == "blocked" - assert result.detail == "uncommitted changes" - (info.path / "dirty.txt").unlink() - - (repo / "uncommitted.txt").write_text("x\n", encoding="utf-8") - result = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert result.status == "paused" - assert result.detail == "destination has uncommitted changes" - (repo / "uncommitted.txt").unlink() - - await git(repo, "checkout", "-b", "side") - result = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert result.status == "paused" - assert result.detail == "destination is on 'side', expected 'main'" - - -async def test_integrate_replays_destination_before_validation(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - (repo / "later.txt").write_text("later\n", encoding="utf-8") - await commit_all(repo, "later") - dest_commit = await git(repo, "rev-parse", "HEAD") - - result = await manager.integrate( - "worker", validation=["test -f later.txt"], validation_runner=shell_validation - ) - assert result.status == "integrated" - ancestor = await run_proc( - ["git", "merge-base", "--is-ancestor", dest_commit, "lecode/worker"], - cwd=repo, - timeout=30, - ) - assert ancestor.exit_code == 0 - - -async def test_integrate_conflict_leaves_merge_in_progress(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "file.txt").write_text("worker\n", encoding="utf-8") - await commit_all(info.path, "worker edit") - (repo / "file.txt").write_text("main\n", encoding="utf-8") - await commit_all(repo, "main edit") - - result = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert result.status == "conflict" - assert result.conflicts == ["file.txt"] - assert (await manager.inspect("worker")).merge_in_progress is True - assert manager.read_sidecar("worker")["integrated_at"] is None - - -async def test_cleanup_worker_guards(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - with pytest.raises(WorktreeError, match="not integrated"): - await manager.cleanup_worker("worker") - - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - assert ( - await manager.integrate("worker", validation=[], allow_unvalidated=True) - ).status == "integrated" - (info.path / "dirty.txt").write_text("x\n", encoding="utf-8") - with pytest.raises(WorktreeError, match="uncommitted changes"): - await manager.cleanup_worker("worker") - - other = await manager.create_worker( - "other", base_commit=base, dest_path=repo, dest_branch="main" - ) - removed = await manager.cleanup_worker("other", discard=True) - assert removed == other - assert not other.path.exists() - branch = await run_proc( - ["git", "show-ref", "--verify", "refs/heads/lecode/other"], cwd=repo, timeout=30 - ) - assert branch.exit_code != 0 - - await manager.cleanup_worker("worker", discard=True) - assert not info.path.exists() - - -async def test_cleanup_worker_refuses_commit_after_integration(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - assert ( - await manager.integrate("worker", validation=[], allow_unvalidated=True) - ).status == "integrated" - (info.path / "after.txt").write_text("after\n", encoding="utf-8") - await commit_all(info.path, "post-integration commit") - - with pytest.raises(WorktreeError, match="not fully integrated"): - await manager.cleanup_worker("worker") - await manager.cleanup_worker("worker", discard=True) - - -async def test_concurrent_integrations_are_serialized(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - first = await manager.create_worker( - "first", base_commit=base, dest_path=repo, dest_branch="main" - ) - second = await manager.create_worker( - "second", base_commit=base, dest_path=repo, dest_branch="main" - ) - (first.path / "first.txt").write_text("first\n", encoding="utf-8") - await commit_all(first.path, "first") - (second.path / "second.txt").write_text("second\n", encoding="utf-8") - await commit_all(second.path, "second") - first_started = asyncio.Event() - release_first = asyncio.Event() - second_started = asyncio.Event() - - async def validate(cmd, _cwd): - if cmd == "first": - first_started.set() - await release_first.wait() - else: - second_started.set() - return 0, "" - - first_task = asyncio.create_task( - manager.integrate("first", validation=["first"], validation_runner=validate) - ) - await first_started.wait() - second_task = asyncio.create_task( - manager.integrate("second", validation=["second"], validation_runner=validate) - ) - await asyncio.sleep(0) - assert not second_started.is_set() - release_first.set() - assert (await first_task).status == "integrated" - assert (await second_task).status == "integrated" - assert (repo / "first.txt").is_file() - assert (repo / "second.txt").is_file() - - -async def test_paused_integrate_is_non_destructive(tmp_path): - repo = await make_repo(tmp_path / "repo") - manager = WorktreeManager(repo) - base = await git(repo, "rev-parse", "HEAD") - info = await manager.create_worker( - "worker", base_commit=base, dest_path=repo, dest_branch="main" - ) - (info.path / "feature.txt").write_text("feature\n", encoding="utf-8") - await commit_all(info.path, "add feature") - (repo / "uncommitted.txt").write_text("keep me\n", encoding="utf-8") - - result = await manager.integrate("worker", validation=[], allow_unvalidated=True) - assert result.status == "paused" - assert (repo / "uncommitted.txt").read_text(encoding="utf-8") == "keep me\n" - assert await git(repo, "status", "--porcelain") == "?? uncommitted.txt" - assert await git(repo, "rev-parse", "HEAD") == base - assert not (repo / ".git" / "MERGE_HEAD").exists() - assert not (repo / "feature.txt").exists() - - # -- /worktree /wt-merge /wt-exit commands ------------------------------------------ From 6b5ca03b5cec90ef2d089bd274f2785794a06d95 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Tue, 15 Sep 2026 15:00:58 +0400 Subject: [PATCH 6/8] fix: make worker submission idempotent --- src/lecode/extras/workers.py | 7 +++++-- src/lecode/tui/app.py | 3 ++- tests/test_slash_features.py | 4 ++++ tests/test_workers.py | 16 ++++++++++++++-- 4 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/lecode/extras/workers.py b/src/lecode/extras/workers.py index 8c816b5..f608f4b 100644 --- a/src/lecode/extras/workers.py +++ b/src/lecode/extras/workers.py @@ -716,13 +716,16 @@ def _notification(self, worker, *, submitted): "deliver": submitted or (worker.origin == "delegated" and worker.background), } existing = self.store.load_events(self.session, "worker_notification") - if not any(item["id"] == note["id"] for item in existing): + is_new = not any(item["id"] == note["id"] for item in existing) + if is_new: self.store.append_event(self.session, "worker_notification", note) - return note + return {**note, "new": is_new} async def submit(self, id: str) -> dict[str, Any]: """Explicitly make a human-origin result available to its parent.""" worker = self.get(id) + if worker.origin != "human": + raise SubagentError("only human workers can be submitted") if worker.state != "completed" or worker.result is None: raise SubagentError("only completed workers can be submitted") return self._notification(worker, submitted=True) diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index 6c50cd2..d352094 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -2066,7 +2066,8 @@ async def submit_worker(self, worker_id: str) -> dict[str, Any]: if self._worker_manager is None: raise RuntimeError("workers are unavailable") note = await self._worker_manager.submit(worker_id) - await self._on_worker_notification(note) + if note["new"] and note["deliver"]: + await self._on_worker_notification(note) return note @property diff --git a/tests/test_slash_features.py b/tests/test_slash_features.py index a5734f5..c4ca124 100644 --- a/tests/test_slash_features.py +++ b/tests/test_slash_features.py @@ -46,6 +46,10 @@ async def test_agent_submit_wakes_root_once(tmp_path, monkeypatch): message["content"] == "Worker updates are available." for message in provider.requests[-1]["messages"] ) + wake_task = app._worker_wake_task + await app.handle_command("/agent 1 submit") + assert app._worker_wake_task is wake_task + assert len(provider.requests) == 2 await manager.shutdown() diff --git a/tests/test_workers.py b/tests/test_workers.py index bf59d0a..35a3757 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -546,14 +546,26 @@ async def test_completion_delivery_background_only_and_human_submit(setup, check await manager.wait(human.id) assert notifications[-1]["worker_id"] == human.id assert manager.drain_notifications() == [] - await manager.submit(human.id) + assert (await manager.submit(human.id))["new"] assert [n["worker_id"] for n in manager.drain_notifications()] == [human.id] - await manager.submit(human.id) + assert not (await manager.submit(human.id))["new"] assert manager.drain_notifications() == [] finally: await manager.shutdown() +@pytest.mark.asyncio +async def test_submit_rejects_delegated_worker(setup, checker_contract): + manager, ctx, _, _, _ = setup + try: + worker = await manager.start(ctx, agent="explore", prompt="delegated") + await manager.wait(worker.id) + with pytest.raises(SubagentError, match="only human workers"): + await manager.submit(worker.id) + finally: + await manager.shutdown() + + @pytest.mark.asyncio async def test_completed_worker_followup_after_restart(setup, checker_contract): manager, ctx, provider, store, _ = setup From d48eacc64879fcc75487102d218c79fcac526500 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Wed, 16 Sep 2026 15:12:15 +0400 Subject: [PATCH 7/8] feat: complete persistent worker supervision --- docs/agents-and-skills.md | 107 ++- docs/configuration.md | 57 ++ src/lecode/agent/builder.py | 5 + src/lecode/agent/runner.py | 160 ++- src/lecode/agent/tools/base.py | 78 +- src/lecode/agent/tools/bash.py | 10 + src/lecode/agent/tools/task.py | 14 +- src/lecode/agent/tools/workers.py | 317 +++++- src/lecode/config/models.py | 9 + src/lecode/context/agents.py | 10 +- src/lecode/extras/workers.py | 617 +++++++++--- src/lecode/extras/worktree.py | 425 +++++++- src/lecode/hooks/decorator.py | 27 +- src/lecode/permission/checker.py | 21 +- src/lecode/session/stats.py | 20 +- src/lecode/slash/catalog.py | 2 +- src/lecode/slash/handlers.py | 83 +- src/lecode/tui/agents.py | 4 +- src/lecode/tui/app.py | 338 +++++-- src/lecode/tui/permission.py | 4 +- src/lecode/tui/statusline.py | 21 +- tests/test_agent_builder.py | 21 +- tests/test_agent_runner.py | 447 ++++++++- tests/test_agents.py | 14 +- tests/test_config_loader.py | 22 + tests/test_hooks_decorator.py | 57 ++ tests/test_permission_checker.py | 96 +- tests/test_session_stats.py | 117 ++- tests/test_slash_features.py | 175 ++++ tests/test_subagents.py | 51 +- tests/test_tool_bash.py | 8 + tests/test_tui_app.py | 411 ++++++++ tests/test_tui_permission.py | 40 +- tests/test_tui_statusline.py | 17 +- tests/test_tui_streaming_pty.py | 161 ++++ tests/test_worker_controls.py | 488 ++++++++++ tests/test_workers.py | 1494 ++++++++++++++++++++++++++--- tests/test_worktree.py | 661 ++++++++++++- 38 files changed, 6019 insertions(+), 590 deletions(-) create mode 100644 tests/test_worker_controls.py diff --git a/docs/agents-and-skills.md b/docs/agents-and-skills.md index e5d0eed..c8abc2d 100644 --- a/docs/agents-and-skills.md +++ b/docs/agents-and-skills.md @@ -12,7 +12,7 @@ Agents are markdown files with YAML frontmatter: - project: `.lecode/agents/*.md` (nearest from the cwd up to the git root) The project layer wins on name collisions, and user files may override the -built-ins (`build`, `plan`, `explore`) by name. `/agents` lists them; Tab +built-ins (`build`, `plan`, `explore`, `general`) by name. `/agents` lists them; Tab cycles the primary agents in the TUI. ```markdown @@ -42,6 +42,111 @@ You are a careful code reviewer. … `denied_tools` always deny, and the overlay's rules/mode can never grant more than the global config. +## Persistent workers: review, validate, integrate + +`build` is the primary coding agent and `plan` is the read-only primary. +`explore` is a read-only subagent; `general` is a general-purpose coding subagent +that can write, subject to all inherited permissions. It does not grant access +denied by the parent, including a read-only parent. Global/project agent files +still override built-ins by name, and hidden/primary-only agents are not eligible +for delegation. The system prompt lists eligible subagents at runtime creation. + +Create a coding worker with +`task(agent='general', prompt='Implement and verify ...', run_in_background=True)`. +The result returns the actual `worker_id`. Then use +`workers(action='send', id=, text='Follow-up ...')` to send +feedback, or `workers(action='list')` to discover existing IDs. `workers` controls +existing workers only: sending to an invented ID or an agent name does not spawn +one. Omit `run_in_background` to wait for the task's answer. + +Writable workers require a Git repository with a committed HEAD and an attached +branch. If the session started outside such a repository, restart it from the +repository: a shell `cd` does not change the session's runtime cwd. + +The `task` tool starts persistent workers. Read-only workers share their parent's +cwd. Write workers get an isolated branch and checkout based on their immediate +parent's committed HEAD. Their sidecar pins the parent checkout, branch, and +base commit, including for nested workers. Dirty parent changes require human +confirmation because they are not copied into the child. + +The supervising model can carry out the following workflow without asking for +routine review/integration approval, subject to the normal tool permissions: + +1. Let the worker finish, or stop it and wait for it and its descendants to become + idle. The worker may checkpoint changes with `bash` on **its own branch**. +2. Call `workers` with `{"action":"inspect","id":"WORKER_ID"}` (`review` is + an alias). This returns the assignment, pinned parent/base, current parent and + worker HEADs, and the actual committed diff. Review that exact diff against + the assignment. A worker's completion message alone is not a review. +3. If changes are needed, use `send` to request them and review again. Inspection + of dirty workers also returns their tracked uncommitted diff and untracked + paths. Commit/checkpoint in the worker, resolve any merge, and inspect again + before integration. These controls never commit the user's root checkout. +4. After accepting the diff, the **immediate supervisor** explicitly calls + `{"action":"integrate","id":"WORKER_ID","reviewed_head":"WORKER_HASH","reviewed_parent_head":"PARENT_HASH"}`. + Both must be the exact full hashes returned by the accepted review, not fresh + lookups at integration time. + General controls can address descendants, but a grandparent model cannot + integrate a grandchild directly. Integrate the grandchild into its parent, + then review and integrate that parent separately. +5. Configured validation commands run in the child's cwd through the registered + `bash` permission gate. Fresh worker-bound permission and hook contexts remain + constrained by the supervisor. A denial, hook rewrite, execution error, + missing process exit result, timeout, or failed check blocks integration. + The integration primitive rechecks both checkouts before fast-forwarding the + pinned parent. Any changed worker or parent HEAD, including a parent rewind, + demands another review; hashes are never automatically refreshed. It never + blindly merges after a validation failure. +6. Once integrated, call `{"action":"cleanup","id":"WORKER_ID"}` if the + workspace is no longer needed. Cleanup refuses uncommitted or unmerged work. + Clean up nested child workspaces before their parents. Worker records, + transcripts, and sidecars remain available; there is no automatic discard + or push. + +Configure project validation in TOML (see [configuration](configuration.md)): + +```toml +[worktree] +validation = ["uv run ruff check", "uv run ruff format --check", "uv run python -m pytest"] +``` + +Only this configured command list is accepted. Models cannot supply replacement +checks or an `allow_unvalidated` flag. An empty list requires an explicit human +confirmation even with tool auto-approval enabled. Without a human callback, +unvalidated integration is unavailable. Confirmations identify the worker, +agent, and cwd. + +### Human controls + +`/agent` lists workers; use either a worker ID or its roster number: + +| Command | Effect | +| --- | --- | +| `/agent ID` | Open the retained transcript/detail view | +| `/agent ID inspect` | Show assignment, pinned destination, both HEADs, and diff | +| `/agent ID integrate WORKER_HASH PARENT_HASH` | Validate and integrate using both full reviewed hashes | +| `/agent ID cleanup` | Remove a clean, integrated workspace; retain its transcript | +| `/agent ID recover` | Ask before recreating a missing checkout from committed history | +| `/agent ID send TEXT` | Queue feedback or a follow-up | +| `/agent ID stop [tree]` | Stop the worker, optionally its descendants | +| `/agent ID resume [TEXT]` | Resume a retained worker | +| `/agent ID submit` | Submit a human-started worker's result to its parent | +| `/agent ID focus` | Focus the composer on that worker | + +Mutating slash controls use the same `workers` dispatch permissions as model +calls. Nested integration first passes the current root's permission gate, then +runs in its immediate supervisor's context with the complete ancestor policy. +Validation also respects current root restrictions after a root agent switch; +neither human controls nor child grants widen ancestor Deny or Ask decisions. +Cleanup has no force/discard argument. Recovery requires human confirmation and +warns that old uncommitted data is unrecoverable. + +Before a retained write worker follows up or resumes, its workspace guard +reconciles committed parent progress into a clean checkout. Dirty work is +preserved; conflicts remain in the worker for resolution. Missing checkouts +fail until a human confirms recovery. Workspace maintenance reserves the idle +worker against concurrent follow-ups, resumes, and nested worker creation. + ## Skills Skills are `SKILL.md` packs: diff --git a/docs/configuration.md b/docs/configuration.md index 9c8f4e8..501524c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -74,6 +74,63 @@ default context window and zeroed pricing. | `enabled` | `{}` | per-tool on/off, e.g. `enabled = { bash = false }` | | `allowlist` | `[]` | when non-empty, only these tools are registered | +## `[worktree]` + +| field | default | meaning | +|---|---|---| +| `validation` | `[]` | ordered list of validation commands run in the worker checkout before integration | + +```toml +[worktree] +validation = ["uv run ruff check", "uv run python -m pytest"] +``` + +Lists replace inherited lists, so a project can explicitly set `validation = []`. +**Empty means an explicit human gate**, not automatic approval or successful +validation. Integrating without checks requires a human decision for that operation, +passed as `allow_unvalidated=True`. Tool auto-approval never supplies that decision. +Blank commands are rejected. Configuring checks does not grant permission to run them. + +### Worker lifecycle integration API + +These `WorktreeManager` primitives are for the worker/tool caller to wire at an +idle boundary, after stopping or joining the worker: + +- `await reconcile(name, *, recreate=False) -> WorktreeInspection` verifies the + checkout against its sidecar and its pinned immediate parent. It merges the + parent's latest committed HEAD into a clean worker for follow-up work. Dirty + or in-progress work is returned unchanged. A conflicting merge raises + `WorktreeError` and remains visible in the worker. It never resets or rebases. + Missing checkouts raise a clean error unless `recreate=True` was explicitly + requested. Recreation uses the surviving worker branch, or the recorded base + commit if the branch is gone, then merges current parent progress. + **Missing uncommitted content cannot be recovered.** +- `await integrate(name, *, reviewed_head, validation, validation_runner, + allow_unvalidated=False) -> MergeResult` requires the parent model to review + the actual diff against the pinned destination and approve the exact full + commit hash. The mechanical gate cannot judge whether that review was honest. + `ValidationRunner` is `Callable[[str, Path], Awaitable[ProcResult]]`, using + `lecode.extras.proc.ProcResult`. The caller must dispatch every command through + `ToolRegistry` and map denials/failures to nonzero results. There is no built-in + shell runner. If merging destination progress changes the candidate HEAD, + integration stops with `WorktreeError` requiring a fresh diff review and hash. + Checks validate that exact candidate. Changed HEADs, dirty checkouts, in-progress + Git operations, and replaced or switched destinations prevent integration. + Only a fast-forward of the pinned parent is performed. No push occurs. +- `await cleanup_worker(name, *, discard=False) -> WorktreeInfo` proves that the + clean worker's **current** HEAD is an ancestor of its current pinned destination, + then removes the checkout and uses safe branch deletion (`-d`). Extra commits + after integration prevent cleanup. Only explicit human `discard=True` permits + deleting dirty or unmerged work and force-deleting its branch. Checkout identity + checks still apply. Sidecars and session data are retained. + +These operations serialize by canonical common Git directory and destination +branch using a cancellable, nonblocking `flock`. Lock sidecars are never unlinked. +The caller must also keep workers idle and prevent concurrent tool writes to the +worker or parent during the operation. Destination path, repository identity and +branch are pinned; the helper never silently switches branches. A nested worker +integrates into its immediate parent, not directly into the repository's main branch. + ## `[ui]` | field | default | meaning | diff --git a/src/lecode/agent/builder.py b/src/lecode/agent/builder.py index 648f262..15846cd 100644 --- a/src/lecode/agent/builder.py +++ b/src/lecode/agent/builder.py @@ -135,6 +135,11 @@ def build_runtime( extra_parts: list[str] = [] if agent is not None and agent.body: extra_parts.append(agent.body) + if "task" in registry.names(): + extra_parts.append( + "Available subagents for task(agent=..., prompt=...):\n" + + "\n".join(f"- {a.name}: {a.description}" for a in agents.subagents()) + ) listing = skills.render_listing() if listing: extra_parts.append(listing) diff --git a/src/lecode/agent/runner.py b/src/lecode/agent/runner.py index 3cbf193..1452c71 100644 --- a/src/lecode/agent/runner.py +++ b/src/lecode/agent/runner.py @@ -113,6 +113,7 @@ class LlmResponse: cost_usd: float #: Characters sent in this call's prompt — calibrates live token estimates. prompt_chars: int = 0 + usage_incomplete: bool = False @dataclass(frozen=True) @@ -216,6 +217,7 @@ class UsageTotals: #: Prompt size of the last API call — the real context fill (the #: accumulated ``input_tokens`` double-counts across tool-call rounds). context_tokens: int = 0 + usage_incomplete: bool = False @dataclass(frozen=True) @@ -280,12 +282,16 @@ async def run( history: list[ChatMessage] = list(messages) # The live conversation, visible through ctx (subagents, hooks). self.ctx.extras["conversation"] = history + manager = self.ctx.extras.get("workers") + if manager is not None and self.session is not None: + manager.repair_interrupted_tools(self.session, history) # Background tasks finished between runs surface at the start. await self._drain_background(history, on_event) await self._consume_workers(history, on_event) input_tokens = 0 output_tokens = 0 cost_usd = 0.0 + usage_incomplete = False context_tokens = 0 turns = 0 tool_calls = 0 @@ -327,7 +333,8 @@ async def run( completed = await self._stream_turn(history, on_event) turns += 1 - in_tok, out_tok, cost = self._turn_cost(completed) + in_tok, out_tok, cost, incomplete = self._turn_cost(completed) + usage_incomplete |= incomplete await self._emit( on_event, LlmResponse( @@ -337,6 +344,7 @@ async def run( output_tokens=out_tok, cost_usd=cost, prompt_chars=prompt_chars, + usage_incomplete=incomplete, ), ) input_tokens += in_tok @@ -344,7 +352,7 @@ async def run( cost_usd += cost context_tokens = in_tok or context_tokens history.append(completed.as_message()) - self._persist_assistant(completed, in_tok, out_tok, cost) + self._persist_assistant(completed, in_tok, out_tok, cost, incomplete) if completed.tool_calls: final_text = "" @@ -369,7 +377,29 @@ async def run( continuing = True continue final_text = (final_text if continuing else "") + completed.content - if await self._consume_workers(history, on_event): + manager = self.ctx.extras.get("workers") + worker_id = self.ctx.extras.get("worker_id") + if ( + turns >= max_turns + and manager is not None + and ( + any(w.is_active for w in manager.descendants(worker_id)) + or manager.pending_notifications(worker_id) + or ( + worker_id is not None + and (manager.pending(worker_id) or manager.questions(worker_id)) + ) + or any( + q is not None and not q.empty() + for q in (self.steer_queue, self.input_queue) + ) + ) + ): + stop_reason = "max_turns" + break + if turns < max_turns and await self._consume_workers( + history, on_event, completing=True + ): continue stop_reason = "done" break @@ -377,6 +407,23 @@ async def run( self._persist_partial(history) raise + if stop_reason != "done": + manager = self.ctx.extras.get("workers") + message = ( + manager.stop_message(self.ctx.extras.get("worker_id"), stop_reason) + if manager is not None + else f"Run stopped: {stop_reason}." + ) + if self.session is not None and self.store is not None: + self.store.append_event( + self.session, + "run_stopped", + { + "reason": stop_reason, + "message": message, + }, + ) + await self._emit(on_event, Error(message)) await self._emit(on_event, Done(stop_reason=stop_reason, turns=turns)) hooks = self.ctx.extras.get("hooks") if hooks is not None and hooks.handlers.get(STOP): @@ -395,7 +442,8 @@ async def run( ) if outcome is not None: review_text = outcome.feedback - in_tok, out_tok, cost = self._usage_cost(outcome.model, outcome.usage or {}) + in_tok, out_tok, cost, incomplete = self._usage_cost(outcome.model, outcome.usage) + usage_incomplete |= incomplete input_tokens += in_tok output_tokens += out_tok cost_usd += cost @@ -410,6 +458,7 @@ async def run( "input_tokens": in_tok, "output_tokens": out_tok, "cost_usd": cost, + "incomplete": incomplete, }, }, ) @@ -434,6 +483,7 @@ async def run( output_tokens=output_tokens, cost_usd=cost_usd, context_tokens=context_tokens, + usage_incomplete=usage_incomplete, ), tool_calls=tool_calls, elapsed_s=elapsed_s, @@ -545,29 +595,25 @@ async def _run_tools( ] manager = self.ctx.extras.get("workers") worker_id = self.ctx.extras.get("worker_id") - suspend = ( - manager is not None - and worker_id is not None - and all(call["function"]["name"] == "task" for call in completed.tool_calls) - ) try: - if suspend: - async with manager.suspend(worker_id): - pairs = await asyncio.gather(*tasks) + if manager is not None: + pairs = await manager.await_tools(worker_id, tasks) else: pairs = await asyncio.gather(*tasks) except asyncio.CancelledError: # Cancel in-flight tools; persist the results that did complete. for task in tasks: task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) for task in tasks: + usage = manager.result_usage(task) if manager else None if task.done() and not task.cancelled(): - self._persist_message(task.result()[0]) + self._persist_message(task.result()[0], usage) raise messages: list[ChatMessage] = [] - for call, (message, result) in zip(completed.tool_calls, pairs, strict=True): + for call, task, (message, result) in zip(completed.tool_calls, tasks, pairs, strict=True): messages.append(message) - self._persist_message(message) + self._persist_message(message, manager.result_usage(task) if manager else None) await self._emit( on_event, ToolResult( @@ -582,26 +628,34 @@ async def _run_tools( # -- queues ------------------------------------------------------------------- - async def _drain_queues(self, history: list[ChatMessage], on_event: OnEvent | None) -> None: + async def _drain_queues( + self, history: list[ChatMessage], on_event: OnEvent | None, *, prefetched=None + ) -> bool: """Drain the steer queue first (priority), then the input queue. Drained items are appended to the history as user messages, with a ``QueuedMessage`` event each so the TUI echoes them when the model actually sees them. """ + messages = [] + prefetched = dict(prefetched or {}) for queue in (self.steer_queue, self.input_queue): if queue is None: continue while True: try: - item = queue.get_nowait() + item = prefetched.pop(queue) if queue in prefetched else queue.get_nowait() except asyncio.QueueEmpty: break message: ChatMessage = {"role": "user", "content": item} history.append(message) self._persist_message(message) - await self._emit(on_event, QueuedMessage(content=item)) + messages.append(item) + # Persist every fetched item before yielding, including cancellation races. + for item in messages: + await self._emit(on_event, QueuedMessage(content=item)) await self._drain_background(history, on_event) + return bool(messages) async def _drain_background(self, history: list[ChatMessage], on_event: OnEvent | None) -> None: """Feed background-task completions in as synthetic user messages. @@ -618,15 +672,41 @@ async def _drain_background(self, history: list[ChatMessage], on_event: OnEvent self._persist_message(message) await self._emit(on_event, QueuedMessage(content=note)) - async def _consume_workers(self, history: list[ChatMessage], on_event: OnEvent | None) -> bool: + async def _consume_workers( + self, history: list[ChatMessage], on_event: OnEvent | None, *, completing=False + ) -> bool: """Deliver worker inboxes only between model turns, never mid tool batch.""" manager = self.ctx.extras.get("workers") if manager is None: return False - items = manager.consume(self.ctx.extras.get("worker_id"), history) + worker_id = self.ctx.extras.get("worker_id") + getters = {} + if completing and worker_id is None: + for queue in (self.steer_queue, self.input_queue): + if queue is not None and queue not in getters: + getters[queue] = asyncio.create_task(queue.get()) + getters[queue].add_done_callback(manager._signal) + queued = False + try: + items = await manager.boundary( + worker_id, + history, + completing=completing, + input_ready=lambda: any(task.done() for task in getters.values()), + ) + finally: + for task in getters.values(): + if not task.done(): + task.cancel() + if getters: + await asyncio.gather(*getters.values(), return_exceptions=True) + prefetched = { + q: task.result() for q, task in getters.items() if not task.cancelled() + } + queued = await self._drain_queues(history, on_event, prefetched=prefetched) for item in items: await self._emit(on_event, QueuedMessage(content=item["text"])) - return bool(items) + return bool(items or queued) # -- automatic compaction ----------------------------------------------------- @@ -681,33 +761,35 @@ async def _maybe_compact( # -- usage / cost --------------------------------------------------------------- - def _usage_cost(self, model: str, usage: dict[str, Any]) -> tuple[int, int, float]: - """(input tokens, output tokens, cost in USD) for one usage dict.""" + def _usage_cost(self, model: str, usage: dict[str, Any] | None) -> tuple[int, int, float, bool]: + """Input/output tokens, known cost in USD, and whether usage is incomplete.""" + if not usage: + return 0, 0, 0.0, True in_tok, out_tok = _usage_tokens(usage) + incomplete = bool(usage.get("incomplete")) if usage.get("cost_usd") is not None: - return in_tok, out_tok, float(usage["cost_usd"]) - if not (in_tok or out_tok): - return in_tok, out_tok, 0.0 + return in_tok, out_tok, float(usage["cost_usd"]), incomplete if self._catalog is None: - self._catalog = Catalog.default() + self._catalog = self.ctx.catalog or Catalog.default() try: pricing = self._catalog.get(model).pricing - except ModelNotFoundError: - return in_tok, out_tok, 0.0 - return in_tok, out_tok, (in_tok * pricing.prompt + out_tok * pricing.completion) / 1e6 + except (ModelNotFoundError, AmbiguousModelError): + return in_tok, out_tok, 0.0, True + cost = (in_tok * pricing.prompt + out_tok * pricing.completion) / 1e6 + return in_tok, out_tok, cost, incomplete - def _turn_cost(self, completed: CompletedMessage) -> tuple[int, int, float]: - """(input tokens, output tokens, cost in USD) for one turn.""" - return self._usage_cost(self.model, completed.usage or {}) + def _turn_cost(self, completed: CompletedMessage) -> tuple[int, int, float, bool]: + """Input/output tokens, known cost, and completeness for one turn.""" + return self._usage_cost(self.model, completed.usage) # -- persistence ------------------------------------------------------------------ def _persist_assistant( - self, completed: CompletedMessage, in_tok: int, out_tok: int, cost: float + self, completed: CompletedMessage, in_tok: int, out_tok: int, cost: float, incomplete: bool ) -> None: - usage = None - if completed.usage is not None: - usage = {"input_tokens": in_tok, "output_tokens": out_tok, "cost_usd": cost} + usage = {"input_tokens": in_tok, "output_tokens": out_tok, "cost_usd": cost} + if incomplete: + usage["incomplete"] = True self._persist_message(completed.as_message(), usage) def _persist_message(self, message: ChatMessage, usage: dict[str, Any] | None = None) -> None: @@ -718,10 +800,10 @@ def _persist_partial(self, history: list[ChatMessage]) -> None: """On cancellation, keep whatever partial assistant turn exists.""" partial = self._partial self._partial = None - if partial is None or not (partial.content or partial.tool_calls): + if partial is None: return history.append(partial.as_message()) - self._persist_message(partial.as_message()) + self._persist_assistant(partial, *self._turn_cost(partial)) # -- events ------------------------------------------------------------------------- diff --git a/src/lecode/agent/tools/base.py b/src/lecode/agent/tools/base.py index ba3f77d..5f0c6aa 100644 --- a/src/lecode/agent/tools/base.py +++ b/src/lecode/agent/tools/base.py @@ -18,6 +18,7 @@ from lecode.config.models import Config from lecode.permission import AllowAlways, Decision, Deny, PermissionChecker +from lecode.permission.checker import CheckResult from lecode.telemetry import capture_exception, record_tool_call #: Cap on tool-argument JSON size (guard against runaway payloads). @@ -93,6 +94,48 @@ async def _fire_permission_hook( await dispatcher.fire(event, tool_name=tool_name, tool_args=args, decision=decision) +async def authorize_tool( + name: str, args: dict[str, Any], ctx: ToolContext, check: CheckResult +) -> ToolResult | None: + """Enforce a verdict for the exact arguments about to execute. + + Shared by initial dispatch and hook-rewritten inputs; approvals apply to + these arguments, never to a previously approved command or action. + """ + if check.decision == Decision.DENY: + return ToolResult(f"denied: {check.reason}", is_error=True) + if check.decision != Decision.ASK: + return None + # Deferred import: hooks.decorator wraps this module's tools. + from lecode.hooks import PERMISSION_REQUEST, PERMISSION_RESULT + + if ctx.auto_approve: + await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="auto") + elif ctx.approval_callback is None: + await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="deny") + return ToolResult( + f"denied: requires approval ({check.reason})", + is_error=True, + metadata={"needs_approval": True}, + ) + else: + await _fire_permission_hook(ctx, PERMISSION_REQUEST, name, args) + approval = await ctx.approval_callback(name, args, check.reason) + if isinstance(approval, Deny): + await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="deny") + return ToolResult(f"denied by user ({check.reason})", is_error=True) + if isinstance(approval, AllowAlways): + grant_always(ctx, name, approval.pattern) + await _fire_permission_hook( + ctx, + PERMISSION_RESULT, + name, + args, + decision="allow_always" if isinstance(approval, AllowAlways) else "allow_once", + ) + return None + + class ToolRegistry: def __init__(self, tools: list[Tool] | None = None) -> None: self._tools: dict[str, Tool] = {} @@ -169,38 +212,9 @@ async def _execute(self, name: str, args_json: str, ctx: ToolContext) -> ToolRes if not isinstance(args, dict): return ToolResult("error: tool arguments must be a JSON object", is_error=True) - check = ctx.permission_checker.check(name, args) - if check.decision == Decision.DENY: - return ToolResult(f"denied: {check.reason}", is_error=True) - if check.decision == Decision.ASK: - # Deferred import: hooks.decorator wraps this module's tools. - from lecode.hooks import PERMISSION_REQUEST, PERMISSION_RESULT - - if ctx.auto_approve: - await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="auto") - elif ctx.approval_callback is None: - await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="deny") - return ToolResult( - f"denied: requires approval ({check.reason})", - is_error=True, - metadata={"needs_approval": True}, - ) - else: - await _fire_permission_hook(ctx, PERMISSION_REQUEST, name, args) - approval = await ctx.approval_callback(name, args, check.reason) - if isinstance(approval, Deny): - await _fire_permission_hook(ctx, PERMISSION_RESULT, name, args, decision="deny") - return ToolResult(f"denied by user ({check.reason})", is_error=True) - if isinstance(approval, AllowAlways): - grant_always(ctx, name, approval.pattern) - await _fire_permission_hook( - ctx, - PERMISSION_RESULT, - name, - args, - decision="allow_always" if isinstance(approval, AllowAlways) else "allow_once", - ) - # AllowOnce / AllowAlways fall through to running the tool. + denied = await authorize_tool(name, args, ctx, ctx.permission_checker.check(name, args)) + if denied is not None: + return denied try: started = time.monotonic() diff --git a/src/lecode/agent/tools/bash.py b/src/lecode/agent/tools/bash.py index accd415..d71e34a 100644 --- a/src/lecode/agent/tools/bash.py +++ b/src/lecode/agent/tools/bash.py @@ -19,6 +19,7 @@ from typing import Any from lecode.agent.tools.base import Tool, ToolContext, ToolResult +from lecode.extras.proc import ProcResult from lecode.extras.rtk import rewrite_command DEFAULT_TIMEOUT_S = 120.0 @@ -204,6 +205,15 @@ async def run(self, args: dict, ctx: ToolContext) -> ToolResult: return ToolResult( (text.rstrip("\n") or "(no output)") + suffix, is_error=timed_out or idle_killed or exit_code != 0, + metadata={ + "proc_result": ProcResult( + exit_code=exit_code, + stdout=text, + stderr="", # The shell merges stderr into stdout. + timed_out=timed_out or idle_killed, + truncated=len(output) > MAX_OUTPUT_BYTES, + ) + }, ) diff --git a/src/lecode/agent/tools/task.py b/src/lecode/agent/tools/task.py index 2c3b895..1902c4d 100644 --- a/src/lecode/agent/tools/task.py +++ b/src/lecode/agent/tools/task.py @@ -38,7 +38,14 @@ def __init__(self) -> None: super().__init__( name="task", description=( - "Run a subagent on a self-contained task and get its final answer. " + "Create a new subagent worker on a self-contained task. " + "Use agent='general' for coding/writes within inherited permissions; " + "use agent='explore' for read-only research. Available subagents are listed " + "in the system prompt; build and plan are primary agents, not subagents. " + "task(agent='general', prompt=..., run_in_background=True) creates a worker " + "and returns its actual worker_id. Use workers(action='send', id=that_id, " + "text=...) for follow-ups; never invent an id. Without background, wait for " + "the final answer. Writable persistent workers use isolated Git worktrees. " "Independent tasks can be dispatched in parallel in one turn. " f"Default agent: {DEFAULT_AGENT} (read-only codebase search)." ), @@ -60,8 +67,9 @@ def __init__(self) -> None: "run_in_background": { "type": "boolean", "description": ( - "Run detached and return a task id immediately; " - "track with the tasks_* tools" + "Create a worker and return its worker_id immediately; " + "track/control it with workers. In contexts without persistent " + "workers, returns a background task id for tasks_* instead." ), }, }, diff --git a/src/lecode/agent/tools/workers.py b/src/lecode/agent/tools/workers.py index 35e7a21..f2e79c2 100644 --- a/src/lecode/agent/tools/workers.py +++ b/src/lecode/agent/tools/workers.py @@ -2,29 +2,72 @@ from __future__ import annotations +import inspect +import json +import re +from dataclasses import replace from typing import Any -from lecode.agent.tools.base import Tool, ToolContext, ToolResult +from lecode.agent.tools.base import Tool, ToolContext, ToolRegistry, ToolResult +from lecode.agent.tools.bash import BashTool +from lecode.extras.proc import ProcResult, run_proc from lecode.extras.subagents import SubagentError from lecode.extras.workers import WORKER_CURRENT_EXTRA, WORKER_EXTRA +from lecode.extras.worktree import WorktreeError, WorktreeManager +from lecode.hooks import apply_hooks, dispatcher_from_config +from lecode.permission import SessionPermissions -_ACTIONS = ("list", "send", "stop", "resume", "submit", "question") +HUMAN_CONTROL_EXTRA = "worker_human_control" +_HUMAN_ROOT_EXTRA = "worker_human_root_context" + +_ACTIONS = ( + "list", + "send", + "stop", + "resume", + "submit", + "question", + "inspect", + "review", + "integrate", + "cleanup", + "recover", +) class WorkersTool(Tool): def __init__(self) -> None: super().__init__( name="workers", - description="List and control delegated workers. Workers can manage descendants only.", + description=( + "List and control existing delegated workers; this tool does not create workers. " + "First call task(agent='general', prompt=..., run_in_background=True) to create " + "a coding worker. Use the actual returned worker_id (or an id from action='list') " + "for action='send' follow-ups; never invent an id or use an agent name as an id. " + "Workers can manage descendants only. " + "After a write worker completes, inspect/review its exact diff against the " + "assignment, then explicitly integrate its reviewed_head and reviewed_parent_head. " + "Only the immediate " + "supervisor integrates; configured validation runs before merging. " + "Cleanup removes only clean integrated workspaces, retaining transcripts. " + "Recover recreates missing checkouts only with human confirmation. " + "Run inspect, review, integrate, cleanup, and recover separately after sibling " + "tools complete; task delegation can still run concurrently." + ), parameters={ "type": "object", "additionalProperties": False, "properties": { "action": {"type": "string", "enum": list(_ACTIONS)}, - "id": {"type": "string"}, + "id": { + "type": "string", + "description": "Existing worker_id returned by task or workers list.", + }, "text": {"type": "string"}, "interrupt": {"type": "boolean"}, "tree": {"type": "boolean"}, + "reviewed_head": {"type": "string"}, + "reviewed_parent_head": {"type": "string"}, }, "required": ["action"], }, @@ -50,10 +93,15 @@ def _validate(self, args: dict[str, Any]) -> str | None: "resume": {"action", "id", "text"}, "submit": {"action", "id"}, "question": {"action", "text"}, + "inspect": {"action", "id"}, + "review": {"action", "id"}, + "integrate": {"action", "id", "reviewed_head", "reviewed_parent_head"}, + "cleanup": {"action", "id"}, + "recover": {"action", "id"}, } if action not in _ACTIONS or set(args) - allowed[action]: return "invalid workers action or arguments" - if action in {"send", "stop", "resume", "submit"} and not isinstance(args.get("id"), str): + if action not in {"list", "question"} and not isinstance(args.get("id"), str): return f"workers {action} needs an id" if action in {"send", "question"} and not isinstance(args.get("text"), str): return f"workers {action} needs text" @@ -61,6 +109,14 @@ def _validate(self, args: dict[str, Any]) -> str | None: args.get("tree", False), bool ): return "interrupt and tree must be booleans" + if "text" in args and not isinstance(args["text"], str): + return "text must be a string" + if action == "integrate": + for field in ("reviewed_head", "reviewed_parent_head"): + if not isinstance(args.get(field), str) or not re.fullmatch( + r"[0-9a-f]{40}|[0-9a-f]{64}", args[field] + ): + return f"{field} must be the exact full reviewed commit hash" return None async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: @@ -83,7 +139,10 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: ] ) if not workers: - return ToolResult("(no workers)") + return ToolResult( + "(no workers) Create one with task(agent='general', prompt=..., " + "run_in_background=True), then use its returned worker_id." + ) return ToolResult( "\n".join( f"{worker.id} {worker.state} {worker.agent} {worker.description}" @@ -101,9 +160,102 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult("question sent to parent") if current is not None and id not in self._descendants(manager, current): return ToolResult("error: workers can manage descendants only", is_error=True) + if not any(worker.id == id for worker in manager.list()): + return ToolResult( + f"error: unknown worker id: {id}. Use workers(action='list') to find existing " + "ids, or task(agent='general', prompt=..., run_in_background=True) to create " + "a worker. Send to its actual returned worker_id, not an invented id.", + is_error=True, + ) try: + if ( + action == "integrate" + and current is None + and ctx.extras.get(HUMAN_CONTROL_EXTRA) is True + ): + worker = manager.get(id) + if worker.parent_id is not None: + # Root dispatch already approved this human control. Preserve the + # live supervisor policy, including every ancestor and its grants. + supervisor = manager._runtime(manager.get(worker.parent_id)).ctx + supervisor = replace( + supervisor, + extras={ + **supervisor.extras, + HUMAN_CONTROL_EXTRA: True, + _HUMAN_ROOT_EXTRA: ctx, + }, + ) + _, result = await supervisor.extras["registry"].dispatch_result( + "human-worker-control", "workers", json.dumps(args), supervisor + ) + return result + if action in {"inspect", "review", "integrate", "cleanup", "recover"}: + async with manager.maintain_workspace(id) as worker: + if action == "integrate" and worker.parent_id != current: + raise WorktreeError("only the immediate supervisor can integrate a worker") + if worker.worktree is None: + raise WorktreeError("worker has no isolated write workspace") + worktrees = await WorktreeManager.discover(manager.cwd) + state = await worktrees.inspect(worker.worktree.name) + if state.info != worker.worktree or worker.cwd != state.info.path: + raise WorktreeError("worker workspace identity changed") + if action == "cleanup": + if any( + child.worktree is not None and child.worktree.path.exists() + for child in manager.children(id) + ): + raise WorktreeError("cleanup retained child workspaces first") + await worktrees.cleanup_worker(worker.worktree.name) + return ToolResult(f"worker {id} workspace cleaned; transcript retained") + if action == "recover": + await self._confirm( + manager, + worker, + "Recreate the missing checkout from retained committed history? " + "Old uncommitted data is unrecoverable.", + ) + await worktrees.reconcile(worker.worktree.name, recreate=True) + return ToolResult( + f"worker {id} checkout recovered; inspect before resuming" + ) + if action == "integrate": + checks = list(manager.config.worktree.validation) + if not checks: + await self._confirm( + manager, + worker, + "No validation checks are configured. Integrate this reviewed " + "commit without validation?", + ) + + async def validate(command, path): + return await self._validation_command( + manager, worker, ctx, command, path + ) + + result = await worktrees.integrate( + worker.worktree.name, + reviewed_head=args["reviewed_head"], + reviewed_parent_head=args["reviewed_parent_head"], + validation=checks, + validation_runner=validate, + allow_unvalidated=not checks, + ) + return ToolResult(result.message, is_error=not result.merged) + if not state.present: + return ToolResult( + f"Worker {id} (@{worker.agent}) checkout missing: {worker.cwd}. " + "Use /agent recover with human confirmation." + ) + return await self._review(manager, worker, worktrees, state) if action == "send": - message_id = await manager.send(id, args["text"], bool(args.get("interrupt"))) + message_id = await manager.send( + id, + args["text"], + bool(args.get("interrupt")), + from_human=ctx.extras.get(HUMAN_CONTROL_EXTRA) is True, + ) return ToolResult(f"worker {id} message {message_id} queued") if action == "stop": await manager.stop(id, bool(args.get("tree"))) @@ -112,10 +264,157 @@ async def run(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult: await manager.resume(id, args.get("text")) return ToolResult(f"worker {id} resumed") note = await manager.submit(id) - return ToolResult(f"worker {id} submitted", metadata={"notification_id": note["id"]}) - except (KeyError, RuntimeError, SubagentError) as e: + return ToolResult( + f"worker {id} submitted", + metadata={"notification_id": note["id"], "notification": note}, + ) + except (KeyError, RuntimeError, SubagentError, WorktreeError) as e: return ToolResult(f"error: {e}", is_error=True) + @staticmethod + async def _confirm(manager, worker, question): + if manager.confirm is None: + raise WorktreeError("human confirmation unavailable in this context") + answer = manager.confirm( + f"Worker {worker.id} (@{worker.agent}), cwd={worker.cwd}: {question}" + ) + if inspect.isawaitable(answer): + answer = await answer + if answer is not True: + raise WorktreeError("human confirmation declined") + + @staticmethod + async def _validation_command(manager, worker, supervisor, command, path): + """Fresh child policy and hooks, with no shell path outside tool dispatch.""" + if "bash" not in supervisor.extras["registry"].names(): + raise WorktreeError("validation bash is unavailable to the supervisor") + definition = supervisor.extras["agents"].get(worker.agent) + if definition is None: + raise WorktreeError("worker agent definition is unavailable") + grants = SessionPermissions(manager.store.load_grants(worker.session)) + hooks, _ = dispatcher_from_config(manager.config, path, session=worker.session) + ctx = replace( + supervisor, + cwd=path, + config=manager.config, + session=worker.session, + session_store=manager.store, + session_perms=grants, + permission_checker=supervisor.permission_checker.for_child( + definition.overlay, + cwd=path, + session_perms=grants, + ), + extras={**supervisor.extras, "hooks": hooks, WORKER_CURRENT_EXTRA: worker.id}, + ) + if supervisor.approval_callback is not None: + + async def approve(name, args, reason): + answer = supervisor.approval_callback( + name, args, f"Worker {worker.id} (@{worker.agent}), cwd={path}: {reason}" + ) + return await answer if inspect.isawaitable(answer) else answer + + ctx.approval_callback = approve + tool = BashTool() + run = tool.run + + async def exact_command(args, context): + if args != {"command": command}: + return ToolResult("validation command rewritten by hook; refused", is_error=True) + root = supervisor.extras.get(_HUMAN_ROOT_EXTRA) + if root is not None: + if "bash" not in root.extras["registry"].names(): + raise WorktreeError("validation bash is unavailable to the current root") + # A root agent switch can replace its checker. Gate the live root + # separately without rebuilding or widening the supervisor chain. + root_tool = BashTool() + + async def execute(approved_args, _): + return await run(approved_args, context) + + async def approve_root(name, args, reason): + answer = root.approval_callback( + name, args, f"Worker {worker.id} (@{worker.agent}), cwd={path}: {reason}" + ) + return await answer if inspect.isawaitable(answer) else answer + + root_tool.run = execute + _, result = await ToolRegistry([root_tool]).dispatch_result( + "human-worker-validation", + "bash", + json.dumps(args), + replace( + root, + approval_callback=approve_root if root.approval_callback else None, + ), + ) + return result + return await run(args, context) + + tool.run = exact_command + registry = ToolRegistry([tool]) + ctx.extras["registry"] = registry + if hooks is not None: + apply_hooks(registry, hooks) + _, result = await registry.dispatch_result( + "worker-validation", "bash", json.dumps({"command": command}), ctx + ) + proc = result.metadata.get("proc_result") + if result.is_error and ( + not isinstance(proc, ProcResult) or (proc.exit_code == 0 and not proc.timed_out) + ): + raise WorktreeError(result.content) + if not isinstance(proc, ProcResult): + raise WorktreeError("validation bash returned no structured process exit result") + return proc + + async def _review(self, manager, worker, worktrees, state) -> ToolResult: + data = state.sidecar + if data is None: + raise WorktreeError("worker has no pinned parent sidecar") + head = await worktrees._git("rev-parse", "HEAD", cwd=worker.cwd) + parent = f"refs/heads/{data['dest_branch']}" + target = await worktrees._git("rev-parse", "--verify", parent, cwd=worker.cwd) + diff = await run_proc( + ["git", "diff", "--no-ext-diff", "--no-textconv", target, head, "--"], cwd=worker.cwd + ) + if diff.truncated: + raise WorktreeError( + "review diff truncated; inspect the full diff before approving HEAD" + ) + if diff.exit_code != 0 or diff.timed_out: + raise WorktreeError(f"cannot read review diff: {diff.stderr}") + assignment = "\n".join( + item["text"] for item in manager.store.load_events(worker.session, "worker_inbox") + ) + content = ( + f"Worker {worker.id} (@{worker.agent}), cwd={worker.cwd}\n" + f"Assignment:\n{assignment}\n" + f"Pinned parent: {data['dest_branch']} at {data['dest_path']}\n" + f"Base: {data['base_commit']}\nParent HEAD: {target}\nWorker HEAD: {head}\n" + f"Committed diff (parent HEAD to worker HEAD):\n{diff.stdout or '(no changes)'}\n" + ) + if state.dirty or state.merge_in_progress: + pending = await worktrees._git( + "diff", "--no-ext-diff", "--no-textconv", "HEAD", "--", cwd=worker.cwd + ) + untracked = await worktrees._git( + "ls-files", "--others", "--exclude-standard", cwd=worker.cwd + ) + return ToolResult( + content + f"Uncommitted diff:\n{pending}\nUntracked paths:\n{untracked}\n" + "Commit/checkpoint changes on the worker's OWN branch using bash, resolve any " + "merge, then review again before integration. Never commit the user root." + ) + return ToolResult( + content + "review this exact diff against assignment then integrate reviewed_head " + f"{head} and reviewed_parent_head {target}. " + "If either HEAD changes, review again; never refresh a hash without review. " + "If changes are needed, send feedback to the worker and review again.", + metadata={"reviewed_head": head, "reviewed_parent_head": target}, + ) + def make_tool() -> Tool: return WorkersTool() diff --git a/src/lecode/config/models.py b/src/lecode/config/models.py index 31897eb..d9cf081 100644 --- a/src/lecode/config/models.py +++ b/src/lecode/config/models.py @@ -82,6 +82,14 @@ class ToolsConfig(BaseModel): allowlist: list[str] = Field(default_factory=list) +class WorktreeConfig(BaseModel): + """``[worktree]``: integration checks; empty requires explicit human approval.""" + + model_config = ConfigDict(extra="ignore") + + validation: list[str] = Field(default_factory=list) + + class UiConfig(BaseModel): """``[ui]`` — display preferences.""" @@ -269,6 +277,7 @@ class Config(BaseModel): compaction: CompactionConfig = Field(default_factory=CompactionConfig) agent: AgentConfig = Field(default_factory=AgentConfig) tools: ToolsConfig = Field(default_factory=ToolsConfig) + worktree: WorktreeConfig = Field(default_factory=WorktreeConfig) ui: UiConfig = Field(default_factory=UiConfig) permissions: PermissionsConfig = Field(default_factory=PermissionsConfig) notifications: NotificationsConfig = Field(default_factory=NotificationsConfig) diff --git a/src/lecode/context/agents.py b/src/lecode/context/agents.py index c93974d..31108e6 100644 --- a/src/lecode/context/agents.py +++ b/src/lecode/context/agents.py @@ -4,7 +4,7 @@ ``~/.config/lecode/agents/*.md`` (global, ``LECODE_CONFIG_DIR`` aware) and ``.lecode/agents/*.md`` (project, nearest from the cwd up to the git root). The project layer wins on name collisions; user files may also override the -built-in agents (``build`` / ``plan`` / ``explore``) by name. +built-in agents (``build`` / ``plan`` / ``explore`` / ``general``) by name. Frontmatter: ``description`` (required), ``mode: primary|subagent|all`` (default ``all``), ``model``, ``temperature``, ``permission`` (overlay mapping @@ -77,6 +77,14 @@ def _builtin_agents() -> dict[str, AgentDefinition]: overlay=BUILTIN_AGENT_OVERLAYS["plan"], builtin=True, ), + "general": AgentDefinition( + name="general", + description="General-purpose coding subagent; can write within inherited permissions.", + body="Complete the delegated task, implement and verify changes as needed, " + "and report the result to your supervisor. Respect inherited permissions.", + mode="subagent", + builtin=True, + ), "explore": AgentDefinition( name="explore", description="Fast read-only codebase explorer (subagent).", diff --git a/src/lecode/extras/workers.py b/src/lecode/extras/workers.py index f608f4b..24e6f9f 100644 --- a/src/lecode/extras/workers.py +++ b/src/lecode/extras/workers.py @@ -1,8 +1,8 @@ """Persistent child runners. Runner boundary wiring is deliberately external. -``consume`` may only be called at a safe conversation boundary. ``suspend`` -belongs around the supervisor's whole tool batch, never around individual -concurrent tool calls. Inbox durability does not make external effects atomic. +``consume`` may only be called at a safe conversation boundary. Only the +supervisor releases its lease, once all unfinished tools are blocked on workers. +Inbox durability does not make external effects atomic. """ from __future__ import annotations @@ -10,7 +10,8 @@ import asyncio import inspect import uuid -from contextlib import asynccontextmanager +from collections.abc import Awaitable, Callable +from contextlib import asynccontextmanager, nullcontext from dataclasses import asdict, dataclass, field from datetime import UTC, datetime from pathlib import Path @@ -22,6 +23,7 @@ from lecode.config.models import Config from lecode.extras.subagents import SUBAGENT_EVENTS_EXTRA, SubagentError, SubagentProgress from lecode.extras.worktree import WorktreeError, WorktreeInfo, WorktreeManager +from lecode.hooks import SUBAGENT_END, SUBAGENT_START, build_envelope, dispatch_event from lecode.session.model import EventRecord, MessageRecord from lecode.session.storage import Session, SessionStore @@ -50,18 +52,26 @@ class Worker: result: RunResult | None = None error: str | None = None started_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + # Includes terminal hooks and the done callback's follow-up scheduling. + _dispatch_active: bool = field(default=False, repr=False, compare=False) @property def session_id(self) -> str: return self.session.id + @property + def is_active(self) -> bool: + """Authoritative activity predicate for scheduling and UI consumers.""" + return self._dispatch_active or self.state in {"queued", "running", "waiting"} + class WorkerManager: """Own child sessions, tasks, and locks until shutdown. ``confirm(question: str) -> bool`` and ``notify(note: dict) -> None`` may be synchronous or asynchronous. Notify is observational; parent delivery - happens only through ``drain_notifications(parent_id)`` at a safe boundary. + happens through ``boundary``/``consume`` (live) or ``drain_notifications`` + (idle), always after outstanding tool calls have been paired. """ def __init__( @@ -74,6 +84,7 @@ def __init__( store: SessionStore | None = None, confirm=None, notify=None, + workspace_guard: Callable[[Worker], Awaitable[None] | None] | None = None, ) -> None: self.config = config self.cwd = Path(cwd) @@ -82,6 +93,11 @@ def __init__( self.session = session or root_ctx.session or self.store.create("workers", self.cwd) self.confirm = confirm self.notify = notify + #: Called with Worker before reused dispatches, may await or raise. + #: Required for write worktrees; must validate/reconcile, never relocate. + self.workspace_guard = workspace_guard or self.reconcile_workspace + self._maintenance: dict[str, asyncio.Task] = {} + self._workspace_starts: dict[str, int] = {} #: UI observer for durable state/usage changes; never affects execution. self.progress = None self._workers: dict[str, Worker] = {} @@ -93,10 +109,124 @@ def __init__( self._closed = False self._slots = asyncio.Semaphore(MAX_EXECUTING) self._leases: set[str] = set() + self._changed = asyncio.Event() + self._blocked: set[asyncio.Task] = set() + self._tool_gates: dict[asyncio.Task, asyncio.Event] = {} + self._tool_owners: dict[asyncio.Task, str | None] = {} + self._wait_results: dict[asyncio.Task, str] = {} + + def _signal(self, *_): + self._changed.set() + self._changed = asyncio.Event() + + @asynccontextmanager + async def _blocking(self): + task = asyncio.current_task() + self._blocked.add(task) + self._signal() + try: + yield + finally: + self._blocked.discard(task) + self._signal() + gate = self._tool_gates.get(task) + if gate is not None and not task.cancelling(): + await gate.wait() + + async def await_tools(self, worker_id: str | None, tasks: list[asyncio.Task]): + """Keep ordinary tool work leased; yield only when every tool is waiting. + + Wait/stop register actual waits, irrespective of tool name or rewritten + arguments. Returning tools cannot execute their post-hooks until leased. + """ + gate = asyncio.Event() + gate.set() + for task in tasks: + self._tool_gates[task] = gate + self._tool_owners[task] = worker_id + task.add_done_callback(self._signal) + try: + while pending := {task for task in tasks if not task.done()}: + changed = self._changed + if worker_id is not None and pending <= self._blocked: + gate.clear() + async with self.suspend(worker_id): + while pending <= self._blocked: + await changed.wait() + changed = self._changed + pending = {task for task in tasks if not task.done()} + if not pending: + break + gate.set() + else: + await changed.wait() + return await asyncio.gather(*tasks) + finally: + for task in tasks: + self._tool_gates.pop(task, None) + self._tool_owners.pop(task, None) def get(self, id: str) -> Worker: return self._workers[id] + def _workspace_available(self, id: str | None) -> None: + while id is not None: + if id in self._maintenance: + raise WorktreeError(f"worker {id} workspace maintenance is in progress") + id = self.get(id).parent_id + + @asynccontextmanager + async def maintain_workspace(self, id: str, *, include_parent: bool = True): + """Reserve the workspace subtree and its integration destination. + + Only the destination's sole unfinished tool may maintain its child. + Reconciliation runs under the worker supervisor's existing reservation. + """ + worker = self.get(id) + task = asyncio.current_task() + if include_parent and task in self._tool_owners: + owner = self._tool_owners[task] + if any( + other is not task and not other.done() and other_owner == owner + for other, other_owner in self._tool_owners.items() + ): + raise WorktreeError( + "workspace maintenance must run separately after sibling tools complete" + ) + if not include_parent and self._maintenance.get(id) is task: + yield worker + return + self._workspace_available(id) + root = self.get(worker.parent_id) if include_parent and worker.parent_id else worker + pending = [root] + while pending: + item = pending.pop() + own_execution = ( + item.id == worker.parent_id and self._tool_owners.get(task) == item.id + ) or (not include_parent and item.id == id and self._tasks.get(id) is task) + if ( + (item.is_active and not own_execution) + or item.id in self._maintenance + or self._workspace_starts.get(item.id, 0) + ): + raise WorktreeError("workspace maintenance requires idle worker and descendants") + pending.extend(self.children(item.id)) + self._maintenance[root.id] = task + try: + yield worker + finally: + del self._maintenance[root.id] + + async def reconcile_workspace(self, worker: Worker) -> None: + """Default retained write-worker guard; never recreate missing data.""" + if worker.worktree is not None: + async with self.maintain_workspace(worker.id, include_parent=False): + manager = await WorktreeManager.discover(self.cwd) + state = await manager.inspect(worker.worktree.name) + if state.info != worker.worktree or worker.cwd != state.info.path: + raise WorktreeError("worker workspace identity changed") + await manager.reconcile(worker.worktree.name) + def list(self) -> list[Worker]: return list(self._workers.values()) @@ -119,10 +249,20 @@ def load(self) -> list[Worker]: data["cwd"] = Path(data["cwd"]) usages = self.store.load_events(child, "worker_usage_checkpoint") if usages: + checkpoint = dict(usages[-1]) + incomplete = checkpoint.pop("incomplete", False) + data["usage_incomplete"] = bool( + data.get("usage_incomplete") + or incomplete + or checkpoint.get("usage_incomplete", False) + ) data["usage_totals"] = { - key: value for key, value in usages[-1].items() if key != "dispatch_id" + key: value for key, value in checkpoint.items() if key != "dispatch_id" } data["usage_totals"] = UsageTotals(**data["usage_totals"]) + data["usage_incomplete"] = bool( + data.get("usage_incomplete") or data["usage_totals"].usage_incomplete + ) if data.get("worktree"): info = data["worktree"] data["worktree"] = WorktreeInfo( @@ -132,10 +272,10 @@ def load(self) -> list[Worker]: result = dict(data["result"]) result["usage_totals"] = UsageTotals(**result["usage_totals"]) data["result"] = RunResult(**result) - if data["state"] in {"queued", "running", "waiting"}: - data["state"] = "interrupted" - data["usage_incomplete"] = True worker = Worker(**data) + if worker.is_active: + worker.state = "interrupted" + worker.usage_incomplete = True except BaseException: if lock is not None: lock.release() @@ -147,7 +287,7 @@ def load(self) -> list[Worker]: def attach(self, session: Session) -> list[Worker]: """Move this root manager to an idle session and hydrate its workers.""" - if any(worker.state in {"queued", "running", "waiting"} for worker in self.list()): + if any(worker.is_active for worker in self.list()): raise RuntimeError("workers are still active") for lock in self._locks.values(): if lock is not None: @@ -194,6 +334,7 @@ def _record(self, worker: Worker) -> None: if worker.worktree is not None: data["worktree"] = {**asdict(worker.worktree), "path": str(worker.worktree.path)} self.store.append_event(self.session, "worker", data) + self._signal() if self.progress is not None: result = self.progress(worker) if inspect.isawaitable(result): @@ -235,8 +376,13 @@ def _record_usage(self, worker): def _agent(self, ctx, name): agents = ctx.extras.get("agents") - if agents is None or name not in {a.name for a in agents.subagents()}: - raise SubagentError(f"unknown or ineligible subagent: {name}") + available = [a.name for a in agents.subagents()] if agents is not None else [] + if name not in available: + raise SubagentError( + f"unknown or ineligible subagent: {name} " + f"(available: {', '.join(available) or 'none'}). " + "Create a worker with task(agent=, prompt=...)." + ) return agents.get(name) @staticmethod @@ -246,16 +392,24 @@ def _read_only(ctx, definition): async def _workspace(self, ctx, definition, id): if self._read_only(ctx, definition): return Path(ctx.cwd), None - manager = await WorktreeManager.discover(ctx.cwd) - branch = await manager._git("rev-parse", "--abbrev-ref", "HEAD", cwd=ctx.cwd) + try: + manager = await WorktreeManager.discover(ctx.cwd) + base = await manager._commit("HEAD", cwd=ctx.cwd) + branch = await manager._git("rev-parse", "--abbrev-ref", "HEAD", cwd=ctx.cwd) + except WorktreeError as error: + raise WorktreeError( + f"Cannot create a writable worker at runtime cwd {ctx.cwd}: {error}. " + "Restart the session from a Git repository with a committed HEAD. " + "A shell cd does not change the session's runtime cwd." + ) from error if branch == "HEAD": raise WorktreeError("write workers require an attached branch, not detached HEAD") destination = Path(await manager._git("rev-parse", "--show-toplevel", cwd=ctx.cwd)) - base = await manager._git("rev-parse", "HEAD", cwd=ctx.cwd) dirty_root = await manager._git("status", "--porcelain", cwd=self.cwd) dirty_parent = await manager._git("status", "--porcelain", cwd=ctx.cwd) if dirty_root or dirty_parent: question = ( + f"Worker {id} (@{definition.name}), parent cwd={ctx.cwd}: " "Uncommitted changes will not enter the worker's committed-HEAD worktree. Continue?" ) approved = self.confirm(question) if self.confirm is not None else False @@ -284,11 +438,18 @@ async def start( if ctx.extras.get("provider") is None: raise SubagentError("no provider available for workers") parent_id = ctx.extras.get(WORKER_CURRENT_EXTRA) + self._workspace_available(parent_id) depth = self.get(parent_id).depth + 1 if parent_id else 1 if depth > MAX_DEPTH: raise SubagentError(f"worker depth exceeds {MAX_DEPTH}") id = uuid.uuid4().hex - cwd, worktree = await self._workspace(ctx, definition, id) + if parent_id is not None: + self._workspace_starts[parent_id] = self._workspace_starts.get(parent_id, 0) + 1 + try: + cwd, worktree = await self._workspace(ctx, definition, id) + finally: + if parent_id is not None: + self._workspace_starts[parent_id] -= 1 if self._closed: raise RuntimeError("worker manager is shut down") session = self.store.create( @@ -323,7 +484,10 @@ def _runtime(self, worker): parent = self._parent_context(worker) definition = self._agent(parent, worker.agent) config = self.config.model_copy( - update={"pierre": self.config.pierre.model_copy(update={"enabled": False})} + update={ + "pierre": self.config.pierre.model_copy(update={"enabled": False}), + "llm": self.config.llm.model_copy(update={"model": worker.session.meta.model}), + } ) runtime = build_runtime( config, @@ -346,7 +510,14 @@ def _runtime(self, worker): callback = parent.approval_callback if callback is not None: - async def approval_callback(tool_name, args, reason): + async def approval_callback( + tool_name, + args, + reason, + *, + worker=worker.id, + conversation=worker.session.name, + ): """Keep worker identity with the root approval FIFO.""" params = inspect.signature(callback).parameters if "worker" in params or any( @@ -356,8 +527,8 @@ async def approval_callback(tool_name, args, reason): tool_name, args, reason, - worker=worker.id, - conversation=worker.session.name, + worker=worker, + conversation=conversation, ) else: result = callback(tool_name, args, reason) @@ -378,8 +549,13 @@ async def approval_callback(tool_name, args, reason): def _enqueue(self, worker, text): if self._closed: raise RuntimeError("worker manager is shut down") - item = {"id": uuid.uuid4().hex, "text": text} + item = { + "id": uuid.uuid4().hex, + "text": text, + "answers": [note["id"] for note in self.questions(worker.id)], + } self.store.append_event(worker.session, "worker_inbox", item) + self._signal() return item["id"] def pending(self, id: str) -> list[dict[str, str]]: @@ -405,6 +581,23 @@ def _outstanding(history): outstanding.pop(message.get("tool_call_id"), None) return outstanding + def repair_interrupted_tools(self, session: Session, history: list[dict] | None = None) -> None: + """Pair persisted calls left by an interrupted run without replaying them.""" + outstanding = self._outstanding(self.store.load_for_model(session)) + live = self._outstanding(history) if history is not None else {} + for call in outstanding.values(): + message = { + "role": "tool", + "tool_call_id": call["id"], + "name": call["function"]["name"], + "content": ( + "Worker interrupted; tool outcome unknown. Inspect state before retrying." + ), + } + self.store.append_message(session, message) + if history is not None and call["id"] in live: + history.append(message) + @staticmethod def _notification_text(note: dict[str, Any]) -> str: if note.get("kind") == "question": @@ -426,14 +619,7 @@ def consume(self, id: str | None, history: list[dict]) -> list[dict[str, str]]: message = {"role": "user", "content": item["text"]} self.store.append_message(session, message, usage={"worker_inbox_id": item["id"]}) history.append(message) - acknowledged = { - item["id"] for item in self.store.load_events(self.session, "worker_notification_ack") - } - notes = [ - note - for note in self.store.load_events(self.session, "worker_notification") - if note["deliver"] and note["parent_id"] == id and note["id"] not in acknowledged - ] + notes = self.pending_notifications(id) for note in notes: text = self._notification_text(note) self.store.append_message( @@ -441,53 +627,167 @@ def consume(self, id: str | None, history: list[dict]) -> list[dict[str, str]]: {"role": "user", "content": text}, usage={"worker_notification_id": note["id"]}, ) - self.store.append_event(self.session, "worker_notification_ack", {"id": note["id"]}) history.append({"role": "user", "content": text}) + self.store.append_event(self.session, "worker_notification_ack", {"id": note["id"]}) items.append({"id": note["id"], "text": text}) return items + def pending_notifications(self, parent_id: str | None = None) -> list[dict[str, Any]]: + """Undelivered notes; the durable message is authoritative, not its ack.""" + # ponytail: scan durable recipients; index delivery IDs if histories grow. + delivered = { + id: { + record.usage["worker_notification_id"] + for record in self.store.read_records(session) + if isinstance(record, MessageRecord) + and record.usage + and "worker_notification_id" in record.usage + } + for id, session in [(None, self.session), *((w.id, w.session) for w in self.list())] + } + delivered_anywhere = set().union(*delivered.values()) + notes = [] + for note in self.store.load_events(self.session, "worker_notification"): + if not note["deliver"] or self._recipient(note) != parent_id: + continue + if note.get("kind") == "question": + if ( + note in self.questions(note["worker_id"]) + and note["id"] not in delivered[parent_id] + ): + notes.append(note) + elif note["id"] not in delivered_anywhere: + notes.append(note) + return notes + + def _recipient(self, note): + parent_id = note["parent_id"] + while parent_id is not None and not self.get(parent_id).is_active: + parent_id = self.get(parent_id).parent_id + return parent_id + + def result_usage(self, task: asyncio.Task) -> dict | None: + """Mark a foreground result delivered in the same durable tool record.""" + if notification_id := self._wait_results.pop(task, None): + return {"worker_notification_id": notification_id} + return None + + async def boundary( + self, + id: str | None, + history: list[dict], + *, + completing: bool = False, + input_ready: Callable[[], bool] | None = None, + ) -> list[dict[str, str]]: + """Deliver input and review delegated descendants before a final answer. + + Human-origin branches are independent unless explicitly submitted. + Wake on each state/input change, so questions can preempt completion. + """ + while True: + changed = self._changed + if id is not None and self.questions(id): + async with self.suspend(id): + while self.questions(id): + await changed.wait() + changed = self._changed + items = self.consume(id, history) + if items or not completing or (input_ready is not None and input_ready()): + return items + if not any(worker.is_active for worker in self.descendants(id)): + return [] + if id is None: + await changed.wait() + else: + async with self.suspend(id): + while ( + any(worker.is_active for worker in self.descendants(id)) + and not self.pending(id) + and not self.pending_notifications(id) + ): + changed = self._changed + await changed.wait() + + def descendants(self, id: str | None) -> list[Worker]: + """Delegated descendants, excluding independent human-origin branches.""" + pending = [id] + descendants = [] + while pending: + children = [w for w in self.children(pending.pop()) if w.origin == "delegated"] + descendants.extend(children) + pending.extend(child.id for child in children) + return descendants + + def stop_message(self, id: str | None, reason: str) -> str: + message = f"Run stopped: {reason}." + unresolved = [ + f"{worker.id} ({worker.state})" + for worker in self.descendants(id) + if worker.is_active or worker.state != "completed" + ] + if unresolved: + message += " Unresolved delegated workers: " + ", ".join(unresolved) + return message + def _launch(self, worker): + self._workspace_available(worker.id) if self._closed: raise RuntimeError("worker manager is shut down") existing = self._tasks.get(worker.id) if existing is not None and not existing.done(): raise RuntimeError("worker is already active") + resuming = worker.dispatch_id is not None worker.state = "queued" worker.error = None worker.result = None worker.dispatch_id = uuid.uuid4().hex + worker._dispatch_active = True self._record(worker) - self._tasks[worker.id] = asyncio.create_task(self._execute(worker)) + self._tasks[worker.id] = asyncio.create_task(self._execute(worker, resuming=resuming)) self._tasks[worker.id].add_done_callback(lambda task: self._finished(worker, task)) def _finished(self, worker, task): - if ( - not self._closed - and self._tasks.get(worker.id) is task - and worker.state == "completed" - and self.pending(worker.id) - ): + if self._tasks.get(worker.id) is not task: + return + worker._dispatch_active = False + if not self._closed and worker.state == "completed" and self.pending(worker.id): self._launch(worker) + else: + self._record(worker) - def _background_descendants(self, id: str) -> list[Worker]: - pending = [id] - descendants: list[Worker] = [] - while pending: - parent = pending.pop() - children = self.children(parent) - descendants.extend(children) - pending.extend(child.id for child in children) - return [worker for worker in descendants if worker.background] - - async def _execute(self, worker): + async def _execute(self, worker, *, resuming=False): try: + start = asyncio.create_task(self._hook(worker, SUBAGENT_START)) + try: + await asyncio.shield(start) + except asyncio.CancelledError: + await start + raise + if worker.state == "stopped": + raise asyncio.CancelledError await self._slots.acquire() self._leases.add(worker.id) worker.state = "running" self._record(worker) + if not worker.cwd.is_dir(): + raise WorktreeError(f"worker checkout missing: {worker.cwd}") + if resuming: + if worker.worktree is not None and self.workspace_guard is None: + raise WorktreeError("write worker resume requires a workspace guard") + if self.workspace_guard is not None: + async with ( + self.maintain_workspace(worker.id, include_parent=False) + if worker.worktree is not None + else nullcontext() + ): + guarded = self.workspace_guard(worker) + if inspect.isawaitable(guarded): + await guarded + runtime = self._runtime(worker) - definition = self._agent(self._parent_context(worker), worker.agent) + runtime.ctx.config.llm.model = worker.session.meta.model runner = AgentRunner( runtime.ctx.extras["provider"], runtime.registry, @@ -497,42 +797,15 @@ async def _execute(self, worker): config=runtime.ctx.config, catalog=runtime.ctx.catalog, ) - runner.model = ( - definition.model or self.config.agent.subagent_model or self.config.llm.model + history = self.store.load_for_model(worker.session) + self.consume(worker.id, history) + worker.result = await runner.run( + [{"role": "system", "content": runtime.system_prompt}, *history], + on_event=lambda event: self._event(worker, event), ) - while True: - history = self.store.load_for_model(worker.session) - self.consume(worker.id, history) - worker.result = await runner.run( - [{"role": "system", "content": runtime.system_prompt}, *history], - on_event=lambda event: self._event(worker, event), - ) - active = [ - child - for child in self._background_descendants(worker.id) - if child.state in {"queued", "running", "waiting"} - ] - if active: - # Let queued descendants use this supervisor's lease. - async with self.suspend(worker.id): - await asyncio.gather( - *(self.wait(child.id) for child in active), return_exceptions=True - ) - if not self.pending(worker.id): - # Background completions are durable notes, consumed by the - # next child run rather than racing a final response. - acknowledged = { - item["id"] - for item in self.store.load_events(self.session, "worker_notification_ack") - } - if not any( - note["deliver"] - and note["parent_id"] == worker.id - and note["id"] not in acknowledged - for note in self.store.load_events(self.session, "worker_notification") - ): - break - worker.state = "completed" + worker.state = "completed" if worker.result.stop_reason == "done" else "failed" + if worker.state == "failed": + worker.error = self.stop_message(worker.id, worker.result.stop_reason) except asyncio.CancelledError: worker.state = "stopped" worker.usage_incomplete = True @@ -552,27 +825,55 @@ async def _execute(self, worker): ) self._record_usage(worker) self._record(worker) + await self._hook(worker, SUBAGENT_END) - if worker.background or worker.origin == "human": - note = self._notification(worker, submitted=False) - if self.notify is not None: - try: - result = self.notify(note) - if inspect.isawaitable(result): - await result - except Exception as error: - self.store.append_event( - self.session, - "worker_notify_error", - { - "worker_id": worker.id, - "error": str(error), - }, - ) + note = self._notification(worker, submitted=False) + if (worker.background or worker.origin == "human") and self.notify is not None: + try: + result = self.notify(note) + if inspect.isawaitable(result): + await result + except Exception as error: + self.store.append_event( + self.session, + "worker_notify_error", + {"worker_id": worker.id, "error": str(error)}, + ) + + async def _hook(self, worker: Worker, event: str) -> None: + parent = self._parent_context(worker) + hooks = parent.extras.get("hooks") + if hooks is None or not hooks.handlers.get(event): + return + envelope = build_envelope( + event, + parent.cwd, + session=parent.session, + agent=worker.agent, + prompt="\n".join(item["text"] for item in self.pending(worker.id)) + if event == SUBAGENT_START + else None, + result={ + "content": ( + worker.error or (worker.result.final_text if worker.result else worker.state) + )[:2000], + "is_error": worker.state != "completed", + } + if event == SUBAGENT_END + else None, + ) + envelope["worker"] = { + "id": worker.id, + "parent_id": worker.parent_id, + "dispatch_id": worker.dispatch_id, + "session_id": worker.session_id, + "cwd": str(worker.cwd), + } + await dispatch_event(event, envelope, hooks.handlers[event]) @asynccontextmanager async def suspend(self, worker_id: str): - """Yield the supervisor lease while awaiting its entire tool batch. + """Yield a supervisor lease only while all of its work is blocked. Only the worker runner task may suspend itself. Concurrent tool tasks cannot release a supervisor that is still executing sibling tools. @@ -588,25 +889,25 @@ async def suspend(self, worker_id: str): self._record(worker) try: yield - finally: - acquire = asyncio.create_task(self._slots.acquire()) - try: - await asyncio.shield(acquire) - except asyncio.CancelledError: - await asyncio.shield(acquire) - raise + except BaseException: + # Cancellation must not queue behind the children we were awaiting. + raise + else: + await self._slots.acquire() self._leases.add(worker_id) worker.state = "running" self._record(worker) async def _event(self, worker, event): if isinstance(event, LlmResponse): + worker.usage_incomplete |= event.usage_incomplete old = worker.usage_totals worker.usage_totals = UsageTotals( old.input_tokens + event.input_tokens, old.output_tokens + event.output_tokens, old.cost_usd + event.cost_usd, event.input_tokens or old.context_tokens, + usage_incomplete=worker.usage_incomplete, ) self.store.append_event( worker.session, @@ -625,20 +926,38 @@ async def _event(self, worker, event): async def wait(self, id: str) -> RunResult: worker = self.get(id) - task = self._tasks.get(id) - while task is not None: - await asyncio.shield(task) - current = self._tasks.get(id) - if current is task: - break - task = current - if worker.result is None or worker.state != "completed": - raise SubagentError(worker.error or f"worker is {worker.state}") - return worker.result + waiting = False + async with self._blocking(): + while worker.is_active: + changed = self._changed + if asyncio.current_task() in self._tool_gates and any( + note.get("kind") == "question" + for note in self.pending_notifications(worker.parent_id) + ): + waiting = True + break + await changed.wait() + result, state, error = worker.result, worker.state, worker.error + caller = asyncio.current_task() + if not waiting and caller in self._tool_gates and not worker.background: + self._wait_results[caller] = f"{worker.dispatch_id}:completed" + if waiting: + # Close this tool pair before delivering the question. Completion + # arrives separately; the child keeps its session and dispatch. + return RunResult( + f"Worker {id} needs a parent answer; use workers send to reply.", + 0, + "waiting", + UsageTotals(), + ) + if result is None or state != "completed": + raise SubagentError(error or f"worker is {state}") + return result async def send( self, id: str, text: str, interrupt: bool = False, *, from_human: bool = False ) -> str: + self._workspace_available(id) worker = self.get(id) message_id = self._enqueue(worker, text) if from_human: @@ -658,7 +977,7 @@ async def send( "deliver": True, }, ) - if interrupt and worker.state in {"queued", "running", "waiting"}: + if interrupt and worker.is_active: await self.stop(id) await self.resume(id) elif worker.state == "completed" and (id not in self._tasks or self._tasks[id].done()): @@ -667,17 +986,23 @@ async def send( async def stop(self, id: str, tree: bool = False) -> None: worker = self.get(id) + finalizing = worker.state in {"completed", "failed", "stopped"} worker.state = "stopped" self._record(worker) task = self._tasks.get(id) if task is not None and not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) + # Let a newly queued coroutine enter its lifecycle try/finally. + await asyncio.sleep(0) + if not finalizing and not task.cancelling(): + task.cancel() + async with self._blocking(): + await asyncio.shield(asyncio.gather(task, return_exceptions=True)) if tree: for child in self.children(id): await self.stop(child.id, tree=True) async def resume(self, id: str, text: str | None = None) -> Worker: + self._workspace_available(id) if self._closed: raise RuntimeError("worker manager is shut down") worker = self.get(id) @@ -686,18 +1011,7 @@ async def resume(self, id: str, text: str | None = None) -> Worker: raise RuntimeError("worker is already active") # Finish protocol pairs, not the interrupted external actions. Replaying # those actions could duplicate an effect whose outcome is unknown. - for call in self._outstanding(self.store.load_for_model(worker.session)).values(): - self.store.append_message( - worker.session, - { - "role": "tool", - "tool_call_id": call["id"], - "name": call["function"]["name"], - "content": ( - "Worker interrupted; tool outcome unknown. Inspect state before retrying." - ), - }, - ) + self.repair_interrupted_tools(worker.session) if text is not None: self._enqueue(worker, text) self._launch(worker) @@ -712,13 +1026,15 @@ def _notification(self, worker, *, submitted): "agent": worker.agent, "origin": worker.origin, "state": worker.state, - "content": worker.result.final_text if worker.result else worker.error or worker.state, - "deliver": submitted or (worker.origin == "delegated" and worker.background), + "content": worker.error + or (worker.result.final_text if worker.result else worker.state), + "deliver": submitted or worker.origin == "delegated", } existing = self.store.load_events(self.session, "worker_notification") is_new = not any(item["id"] == note["id"] for item in existing) if is_new: self.store.append_event(self.session, "worker_notification", note) + self._signal() return {**note, "new": is_new} async def submit(self, id: str) -> dict[str, Any]: @@ -733,8 +1049,6 @@ async def submit(self, id: str) -> dict[str, Any]: def ask_parent(self, id: str, text: str) -> dict[str, Any]: """Durably deliver a worker question without exposing interactive UI.""" worker = self.get(id) - if worker.parent_id is None: - raise SubagentError("the root worker has no parent") note = { "id": uuid.uuid4().hex, "worker_id": worker.id, @@ -748,20 +1062,29 @@ def ask_parent(self, id: str, text: str) -> dict[str, Any]: "deliver": True, } self.store.append_event(self.session, "worker_notification", note) + self._signal() return note - def drain_notifications(self, parent_id: str | None = None) -> list[dict[str, Any]]: - """Acknowledge parent-bound result notes; human results require submit.""" - acknowledged = { - item["id"] for item in self.store.load_events(self.session, "worker_notification_ack") + def questions(self, id: str) -> list[dict[str, Any]]: + """Questions stay unresolved until a durable inbox reply is queued.""" + answered = { + answer + for item in self.store.load_events(self.get(id).session, "worker_inbox") + for answer in item.get("answers", []) } - notes = [ - item - for item in self.store.load_events(self.session, "worker_notification") - if item["deliver"] and item["parent_id"] == parent_id and item["id"] not in acknowledged + return [ + note + for note in self.store.load_events(self.session, "worker_notification") + if note["worker_id"] == id + and note.get("kind") == "question" + and note["id"] not in answered ] - for note in notes: - self.store.append_event(self.session, "worker_notification_ack", {"id": note["id"]}) + + def drain_notifications(self, parent_id: str | None = None) -> list[dict[str, Any]]: + """Durably deliver notes at an idle boundary; use consume for live history.""" + notes = self.pending_notifications(parent_id) + session = self.get(parent_id).session if parent_id is not None else self.session + self.consume(parent_id, self.store.load_for_model(session)) return notes async def shutdown(self) -> None: diff --git a/src/lecode/extras/worktree.py b/src/lecode/extras/worktree.py index b66bbad..eeccdbe 100644 --- a/src/lecode/extras/worktree.py +++ b/src/lecode/extras/worktree.py @@ -12,23 +12,29 @@ merge is left in progress (standard git flow) and the conflicted paths are reported — no auto-resolution, ``git merge --abort`` stays available. -Worker worktrees: ``create_worker`` pins a base commit and a destination -(path + branch) in a sidecar under ``.lecode/worktrees/.json`` for a -future explicit integration workflow. ``discover`` returns the *main* -repository root even when called from inside a linked worktree, so worktrees -and sidecars agree across restarts. +Worker worktrees pin their immediate parent in a durable sidecar. Reconcile +merges committed parent progress only into clean workers. Integration requires +exact reviewed worker and parent commits and caller-supplied, policy-checked validation. """ from __future__ import annotations +import asyncio +import fcntl +import hashlib import json import os import re +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager from dataclasses import dataclass +from datetime import UTC, datetime from pathlib import Path from typing import Any -from lecode.extras.proc import run_proc +from lecode.extras.proc import ProcResult, run_proc + +type ValidationRunner = Callable[[str, Path], Awaitable[ProcResult]] #: Per-git-command timeout. GIT_TIMEOUT_S = 30.0 @@ -93,33 +99,37 @@ def __init__(self, repo_root: Path | str) -> None: @classmethod async def discover(cls, cwd: Path | str) -> WorktreeManager: - """The manager for the *main* repository containing ``cwd``. + """Find the checkout, or the owner of a managed worktree's sidecars. - Inside a linked worktree ``--show-toplevel`` names the worktree, not - the repository that owns it; ``--git-common-dir`` names the main - ``.git``, whose parent is the root we want. + An unrelated linked checkout remains its own root, never the main + checkout merely because it shares a common Git directory. """ start = Path(cwd).expanduser().resolve() - common = await run_proc( - ["git", "rev-parse", "--git-common-dir"], cwd=cwd, timeout=GIT_TIMEOUT_S - ) - if common.exit_code == 0: - git_dir = Path(common.stdout.strip()) - if not git_dir.is_absolute(): - git_dir = start / git_dir - git_dir = git_dir.resolve() - if git_dir.name == ".git": - return cls(git_dir.parent) result = await run_proc( - ["git", "rev-parse", "--show-toplevel"], cwd=cwd, timeout=GIT_TIMEOUT_S + ["git", "rev-parse", "--show-toplevel"], cwd=start, timeout=GIT_TIMEOUT_S ) if result.exit_code != 0: raise WorktreeError(f"not a git repository: {cwd}") - return cls(Path(result.stdout.strip())) + root = Path(result.stdout.strip()).resolve() + manager = cls(root) + if root.parent.name == "worktrees" and root.parent.parent.name == ".lecode": + owner = root.parents[2] + branch = await manager._git("rev-parse", "--abbrev-ref", "HEAD") + managed = (owner / WORKTREE_ROOT / f"{root.name}.json").is_file() + if (managed or branch == f"{BRANCH_PREFIX}{root.name}") and ( + await manager._common_dir(owner) == await manager._common_dir(root) + ): + manager = cls(owner) + return manager async def _git(self, *args: str, cwd: Path | None = None) -> str: """Run git, returning stdout; non-zero exits become WorktreeError.""" - result = await run_proc(["git", *args], cwd=cwd or self.repo_root, timeout=GIT_TIMEOUT_S) + try: + result = await run_proc( + ["git", *args], cwd=cwd or self.repo_root, timeout=GIT_TIMEOUT_S + ) + except OSError as error: + raise WorktreeError(f"cannot run git at {cwd or self.repo_root}: {error}") from error if result.exit_code != 0: detail = result.stderr.strip() or result.stdout.strip() raise WorktreeError(f"git {args[0]} failed: {detail or f'exit {result.exit_code}'}") @@ -135,11 +145,23 @@ async def _common_dir(self, cwd: Path) -> Path: common = Path(result.stdout.strip()) return (common if common.is_absolute() else cwd / common).resolve() + @staticmethod + def _identity(path: Path) -> list[int]: + try: + stat = path.stat() + except OSError as error: + raise WorktreeError(f"cannot identify checkout: {path}: {error}") from error + return [stat.st_dev, stat.st_ino] + async def _commit(self, revision: str, *, cwd: Path) -> str: """Resolve one revision to a commit, without accepting arbitrary refs later.""" - return await self._git("rev-parse", "--verify", f"{revision}^{{commit}}", cwd=cwd) + return await self._git( + "rev-parse", "--verify", "--end-of-options", f"{revision}^{{commit}}", cwd=cwd + ) def _info(self, name: str) -> WorktreeInfo: + if not _NAME_RE.fullmatch(name): + raise WorktreeError(f"invalid worktree name: {name!r} (letters, digits, . _ -)") return WorktreeInfo( name=name, path=self.repo_root / WORKTREE_ROOT / name, @@ -152,10 +174,10 @@ def _require(self, name: str) -> WorktreeInfo: raise WorktreeError(f"no such worktree: {name} (expected at {info.path})") return info - def _exclude_worktree_root(self) -> None: + async def _exclude_worktree_root(self) -> None: """Best-effort: keep the worktree dir out of the main checkout's status.""" try: - exclude = self.repo_root / ".git" / "info" / "exclude" + exclude = await self._common_dir(self.repo_root) / "info" / "exclude" if not exclude.parent.is_dir(): return # .git is a file (linked worktree) — nothing to do existing = exclude.read_text(encoding="utf-8") if exclude.is_file() else "" @@ -176,6 +198,8 @@ async def _check_new(self, name: str) -> WorktreeInfo: info = self._info(name) if info.path.exists(): raise WorktreeError(f"worktree already exists: {info.path}") + if self._sidecar_path(name).exists(): + raise WorktreeError(f"worker sidecar already exists: {name}; use reconcile") branch_ref = await run_proc( ["git", "show-ref", "--verify", f"refs/heads/{info.branch}"], cwd=self.repo_root, @@ -188,7 +212,7 @@ async def _check_new(self, name: str) -> WorktreeInfo: async def create(self, name: str) -> WorktreeInfo: """Create ``.lecode/worktrees/`` on a fresh ``lecode/`` branch.""" info = await self._check_new(name) - self._exclude_worktree_root() + await self._exclude_worktree_root() await self._git("worktree", "add", "-b", info.branch, str(info.path)) return info @@ -240,6 +264,7 @@ def _validated_sidecar(self, name: str, info: WorktreeInfo) -> dict[str, Any] | and path == path.resolve() == expected_path and data["branch"] == info.branch and isinstance(data["base_commit"], str) + and bool(re.fullmatch(r"[0-9a-f]{40}|[0-9a-f]{64}", data["base_commit"])) and isinstance(data["dest_path"], str) and dest_path.is_absolute() and dest_path == dest_path.resolve() @@ -248,6 +273,15 @@ def _validated_sidecar(self, name: str, info: WorktreeInfo) -> dict[str, Any] | and isinstance(data["dest_common_dir"], str) and common_dir.is_absolute() and common_dir == common_dir.resolve() + and all( + key not in data + or ( + isinstance(data[key], list) + and len(data[key]) == 2 + and all(type(n) is int for n in data[key]) + ) + for key in ("repo_identity", "dest_identity", "checkout_identity") + ) and ( data.get("integrated_at") is None or isinstance(data.get("integrated_at"), str) ) @@ -286,6 +320,9 @@ async def _destination_path( destination = Path(root.stdout.strip()).resolve() if await self._common_dir(destination) != expected_common: raise WorktreeError(f"destination is not in this repository: {destination}") + actual_branch = await self._git("symbolic-ref", "-q", "HEAD", cwd=destination) + if actual_branch != f"refs/heads/{dest_branch}": + raise WorktreeError(f"destination branch changed: expected {dest_branch}") return destination async def create_worker( @@ -307,7 +344,7 @@ async def create_worker( dest_path, dest_branch, expected_common=common_dir ) base = await self._commit(base_commit, cwd=self.repo_root) - self._exclude_worktree_root() + await self._exclude_worktree_root() await self._git("worktree", "add", "-b", info.branch, str(info.path), base) self.write_sidecar( name, @@ -319,12 +356,332 @@ async def create_worker( "dest_path": str(destination), "dest_branch": dest_branch, "dest_common_dir": str(common_dir), + "repo_identity": self._identity(common_dir), + "dest_identity": self._identity(destination), + "checkout_identity": self._identity(info.path), "integrated_at": None, "integrated_head": None, }, ) return info + def _worker_data(self, name: str) -> tuple[WorktreeInfo, dict[str, Any]]: + info = self._info(name) + data = self._validated_sidecar(name, info) + if data is None: + raise WorktreeError(f"worker '{name}' has no pinned sidecar") + if data["dest_branch"] == info.branch or Path(data["dest_path"]) == info.path: + raise WorktreeError("worker cannot be its own destination") + return info, data + + @asynccontextmanager + async def _worker_lock(self, name: str) -> AsyncIterator[tuple[WorktreeInfo, dict[str, Any]]]: + """One stable inode per canonical repository + destination branch. + + Nonblocking flock also serializes separate opens within this process. + Polling is cancellable, unlike a blocking flock in a background thread. + """ + info, data = self._worker_data(name) + common = await self._common_dir(self.repo_root) + if common != Path(data["dest_common_dir"]): + raise WorktreeError("worker repository identity changed") + if "repo_identity" in data and self._identity(common) != data["repo_identity"]: + raise WorktreeError("worker repository identity changed") + key = hashlib.sha256(data["dest_branch"].encode()).hexdigest() + locks = common / "lecode-integration-locks" + try: + locks.mkdir(exist_ok=True) + lock = (locks / f"{key}.lock").open("a") + except OSError as error: + raise WorktreeError(f"cannot lock worker destination: {error}") from error + with lock: + while True: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + await asyncio.sleep(0.05) + try: + _, current = self._worker_data(name) + if any( + current.get(k) != data.get(k) + for k in ( + "path", + "branch", + "base_commit", + "dest_path", + "dest_branch", + "dest_common_dir", + "repo_identity", + "dest_identity", + "checkout_identity", + ) + ): + raise WorktreeError("worker sidecar changed while waiting for destination") + yield info, current + finally: + fcntl.flock(lock, fcntl.LOCK_UN) + # Never unlink: waiters must continue locking this same inode. + + async def _registration(self, path: Path) -> dict[str, str] | None: + listing = await self._git("worktree", "list", "--porcelain", "-z") + for entry in listing.split("\0\0"): + fields = dict(field.partition(" ")[::2] for field in entry.split("\0") if field) + if fields.get("worktree") == str(path): + return fields + return None + + async def _checkout( + self, path: Path, branch: str, common: Path, identity: list[int] | None = None + ) -> None: + """Reject missing, replaced, detached, switched, or unregistered checkouts.""" + if not path.is_dir() or path.is_symlink(): + raise WorktreeError(f"checkout missing or replaced: {path}") + if identity is not None and self._identity(path) != identity: + raise WorktreeError(f"checkout identity changed: {path}") + root = Path(await self._git("rev-parse", "--show-toplevel", cwd=path)).resolve() + if root != path or await self._common_dir(path) != common: + raise WorktreeError(f"checkout identity changed: {path}") + actual = await self._git("symbolic-ref", "-q", "HEAD", cwd=path) + if actual != f"refs/heads/{branch}": + raise WorktreeError(f"checkout branch changed: expected {branch} at {path}") + registration = await self._registration(path) + if registration is None or registration.get("branch") != actual: + raise WorktreeError(f"checkout is not registered: {path}") + git_dir = Path(await self._git("rev-parse", "--absolute-git-dir", cwd=path)) + if git_dir != common: + try: + backlink = Path((git_dir / "gitdir").read_text().strip()).resolve() + except OSError as error: + raise WorktreeError(f"cannot identify linked checkout: {path}") from error + if backlink != path / ".git": + raise WorktreeError(f"linked checkout identity changed: {path}") + if identity is not None and self._identity(path) != identity: + raise WorktreeError(f"checkout identity changed: {path}") + + async def _destination(self, data: dict[str, Any]) -> Path: + path = Path(data["dest_path"]) + await self._checkout( + path, data["dest_branch"], Path(data["dest_common_dir"]), data.get("dest_identity") + ) + if "dest_identity" in data and self._identity(path) != data["dest_identity"]: + raise WorktreeError(f"destination checkout identity changed: {path}") + if ( + "repo_identity" in data + and self._identity(Path(data["dest_common_dir"])) != data["repo_identity"] + ): + raise WorktreeError("destination repository identity changed") + return path + + async def _state(self, path: Path) -> tuple[bool, bool]: + dirty = bool(await self._git("status", "--porcelain", "--untracked-files=all", cwd=path)) + git_dir = Path(await self._git("rev-parse", "--absolute-git-dir", cwd=path)) + busy = any( + (git_dir / marker).exists() + for marker in ( + "MERGE_HEAD", + "CHERRY_PICK_HEAD", + "REVERT_HEAD", + "rebase-merge", + "rebase-apply", + "sequencer", + ) + ) + return dirty, busy + + async def _clean(self, path: Path) -> None: + dirty, busy = await self._state(path) + if dirty or busy: + raise WorktreeError( + f"checkout has uncommitted changes or an operation in progress: {path}" + ) + + async def _ancestor(self, ancestor: str, descendant: str) -> bool: + result = await run_proc( + ["git", "merge-base", "--is-ancestor", ancestor, descendant], + cwd=self.repo_root, + timeout=GIT_TIMEOUT_S, + ) + if result.exit_code not in (0, 1): + raise WorktreeError(f"cannot verify commit ancestry: {result.stderr.strip()}") + return result.exit_code == 0 + + async def _merge_destination(self, info: WorktreeInfo, target: str) -> None: + head = await self._commit("HEAD", cwd=info.path) + if not await self._ancestor(target, head): + await self._git("merge", "--ff", "--no-edit", "--no-autostash", target, cwd=info.path) + + async def reconcile(self, name: str, *, recreate: bool = False) -> WorktreeInspection: + """Check an idle worker and merge its parent's committed progress if clean. + + The caller must stop/join the worker first. Dirty/in-progress work is + retained and returned unchanged. Conflicting merges raise WorktreeError + and remain in the worker for resolution. Missing checkouts require + explicit recreation; missing uncommitted content cannot be recovered. + """ + async with self._worker_lock(name) as (info, data): + destination = await self._destination(data) + if not info.path.exists(): + if recreate is not True: + raise WorktreeError( + f"worker '{name}' checkout missing; use recreate=True explicitly. " + "Missing uncommitted content cannot be recovered." + ) + # Remove only this absent checkout's stale registration, not + # other worktrees' metadata or retained sidecar/session data. + registration = await self._registration(info.path) + if registration is not None: + if registration.get("branch") != f"refs/heads/{info.branch}": + raise WorktreeError("missing worker registration has changed branch") + await self._git("worktree", "remove", str(info.path)) + ref = await run_proc( + ["git", "show-ref", "--verify", f"refs/heads/{info.branch}"], + cwd=self.repo_root, + timeout=GIT_TIMEOUT_S, + ) + if ref.exit_code == 0: + await self._git("worktree", "add", str(info.path), info.branch) + else: + base = await self._commit(data["base_commit"], cwd=self.repo_root) + await self._git("worktree", "add", "-b", info.branch, str(info.path), base) + data = {**data, "checkout_identity": self._identity(info.path)} + self.write_sidecar(name, data) + await self._checkout( + info.path, + info.branch, + Path(data["dest_common_dir"]), + data.get("checkout_identity", []), + ) + dirty, busy = await self._state(info.path) + if not dirty and not busy: + target = await self._commit("HEAD", cwd=destination) + await self._merge_destination(info, target) + return await self.inspect(name) + + async def integrate( + self, + name: str, + *, + reviewed_head: str, + reviewed_parent_head: str, + validation: list[str], + validation_runner: ValidationRunner, + allow_unvalidated: bool = False, + ) -> MergeResult: + """Validate the exact reviewed candidate, then fast-forward its pinned parent. + + The parent model must inspect the actual diff and approve both full HEAD + hashes. Any parent mismatch requires re-review, even if the worker already + contains that parent. This mechanical gate cannot establish review honesty. The caller + owns worker idleness and must route validation_runner(cmd, worker_path) + through ToolRegistry; there is deliberately no default shell runner. + Only an explicit human decision may set allow_unvalidated=True. + """ + if not isinstance(validation, list) or any( + not isinstance(cmd, str) or not cmd.strip() for cmd in validation + ): + raise WorktreeError("validation must be a list of nonempty commands") + checks = tuple(validation) + if not checks and allow_unvalidated is not True: + raise WorktreeError("no validation checks: explicit human approval required") + for field, value in ( + ("reviewed_head", reviewed_head), + ("reviewed_parent_head", reviewed_parent_head), + ): + if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{40}|[0-9a-f]{64}", value): + raise WorktreeError(f"{field} must be the exact full reviewed commit hash") + async with self._worker_lock(name) as (info, data): + destination = await self._destination(data) + common = Path(data["dest_common_dir"]) + await self._checkout(info.path, info.branch, common, data.get("checkout_identity", [])) + await self._clean(destination) + await self._clean(info.path) + if await self._commit("HEAD", cwd=info.path) != reviewed_head: + raise WorktreeError("worker HEAD changed: re-review the actual diff and new HEAD") + target = await self._commit("HEAD", cwd=destination) + if target != reviewed_parent_head: + if await self._ancestor(reviewed_parent_head, target): + await self._merge_destination(info, target) + raise WorktreeError("destination HEAD changed since review: re-review required") + await self._merge_destination(info, target) + if await self._commit("HEAD", cwd=info.path) != reviewed_head: + raise WorktreeError( + "destination merged into worker: re-review the actual diff and new HEAD" + ) + + async def check_review() -> None: + if self._worker_data(name)[1] != data: + raise WorktreeError("worker sidecar changed during validation") + await self._destination(data) + await self._checkout( + info.path, info.branch, common, data.get("checkout_identity", []) + ) + await self._clean(info.path) + await self._clean(destination) + if await self._commit("HEAD", cwd=info.path) != reviewed_head: + raise WorktreeError("worker HEAD changed during validation: re-review required") + if await self._commit("HEAD", cwd=destination) != reviewed_parent_head: + raise WorktreeError( + "destination HEAD changed during validation: retry and re-review" + ) + if self._identity(info.path) != data.get("checkout_identity"): + raise WorktreeError(f"checkout identity changed: {info.path}") + if "dest_identity" in data and self._identity(destination) != data["dest_identity"]: + raise WorktreeError(f"destination checkout identity changed: {destination}") + + await check_review() + for cmd in checks: + try: + result = await validation_runner(cmd, info.path) + except Exception as error: + raise WorktreeError(f"validation failed ({cmd}): {error}") from error + if not isinstance(result, ProcResult) or result.exit_code != 0 or result.timed_out: + detail = ( + result.stderr or result.stdout + if isinstance(result, ProcResult) + else "invalid runner result" + ) + raise WorktreeError(f"validation failed ({cmd}): {detail}") + await check_review() + await self._git("merge", "--ff-only", "--no-autostash", reviewed_head, cwd=destination) + self.write_sidecar( + name, + { + **data, + "integrated_head": reviewed_head, + "integrated_at": datetime.now(UTC).isoformat(), + }, + ) + return MergeResult(True, [], f"integrated {reviewed_head} into {data['dest_branch']}") + + async def cleanup_worker(self, name: str, *, discard: bool = False) -> WorktreeInfo: + """Remove a clean, integrated worker; retain its sidecar and session data. + + discard=True is an explicit human data-loss decision, never inferred + from tool auto-approval. Identity checks cannot be bypassed by discard. + """ + async with self._worker_lock(name) as (info, data): + await self._checkout( + info.path, + info.branch, + Path(data["dest_common_dir"]), + data.get("checkout_identity", []), + ) + destination = await self._destination(data) + if discard is not True: + await self._clean(info.path) + head = await self._commit("HEAD", cwd=info.path) + target = await self._commit("HEAD", cwd=destination) + if not await self._ancestor(head, target): + raise WorktreeError("worker HEAD is not integrated into its pinned destination") + await self._git( + "worktree", "remove", *(["--force"] if discard is True else []), str(info.path) + ) + await self._git( + "branch", "-D" if discard is True else "-d", info.branch, cwd=destination + ) + return info + async def attach(self, name: str) -> WorktreeInfo: """Resume an existing worker: directory and branch must both exist.""" info = self._info(name) @@ -337,6 +694,11 @@ async def attach(self, name: str) -> WorktreeInfo: ) if branch_ref.exit_code != 0: raise WorktreeError(f"worktree '{name}' has no branch {info.branch}") + data = self._validated_sidecar(name, info) + common = Path(data["dest_common_dir"]) if data else await self._common_dir(self.repo_root) + await self._checkout( + info.path, info.branch, common, data.get("checkout_identity", []) if data else None + ) return info async def inspect(self, name: str) -> WorktreeInspection: @@ -347,6 +709,13 @@ async def inspect(self, name: str) -> WorktreeInspection: return WorktreeInspection( present=False, info=info, dirty=False, merge_in_progress=False, sidecar=sidecar ) + if sidecar is not None: + await self._checkout( + info.path, + info.branch, + Path(sidecar["dest_common_dir"]), + sidecar.get("checkout_identity", []), + ) porcelain = await self._git("status", "--porcelain", cwd=info.path) merge_head = await run_proc( ["git", "rev-parse", "-q", "--verify", "MERGE_HEAD"], diff --git a/src/lecode/hooks/decorator.py b/src/lecode/hooks/decorator.py index bcfd14c..7201640 100644 --- a/src/lecode/hooks/decorator.py +++ b/src/lecode/hooks/decorator.py @@ -2,9 +2,9 @@ The permission checker runs first — ``ToolRegistry._execute`` returns early on a checker Deny without ever calling ``tool.run``, so decorated hooks are never -consulted for checker-denied calls and can therefore only narrow. PreToolUse -fires before execution (Deny blocks, Ask requires approval unless the context -auto-approves, ``rewritten_input`` replaces the tool args); after execution, +consulted for checker-denied calls. PreToolUse fires before execution; rewritten +arguments are checked again so a rewrite cannot widen permissions. Deny blocks, +Ask uses normal human approval unless the context auto-approves. After execution, exactly one of PostToolUse (success) or PostToolUseFailure (``is_error`` result or an exception out of ``tool.run``) fires — both informational, their verdicts are recorded in the result metadata and cannot undo anything. On an @@ -15,9 +15,10 @@ from typing import Any -from lecode.agent.tools.base import Tool, ToolContext, ToolRegistry, ToolResult +from lecode.agent.tools.base import Tool, ToolContext, ToolRegistry, ToolResult, authorize_tool from lecode.hooks.events import POST_TOOL_USE_FAILURE from lecode.hooks.runner import HookDispatcher +from lecode.permission.checker import CheckResult, Decision def _wrap(tool: Tool, dispatcher: HookDispatcher) -> None: @@ -33,13 +34,19 @@ async def hooked_run(args: dict[str, Any], ctx: ToolContext) -> ToolResult: is_error=True, metadata={"hook_verdict": "deny"}, ) - if pre.verdict == "ask" and not ctx.auto_approve: - return ToolResult( - f"denied: hook requires approval ({pre.reason or 'PreToolUse hook'})", - is_error=True, - metadata={"needs_approval": True, "hook_verdict": "ask"}, - ) effective_args = pre.rewritten_input if pre.rewritten_input is not None else args + if effective_args != args or pre.verdict == "ask": + check = ( + ctx.permission_checker.check(tool.name, effective_args) + if effective_args != args + else CheckResult(Decision.ALLOW, "original arguments already authorized") + ) + if pre.verdict == "ask" and check.decision == Decision.ALLOW: + check = CheckResult(Decision.ASK, pre.reason or "PreToolUse hook") + denied = await authorize_tool(tool.name, effective_args, ctx, check) + if denied is not None: + denied.metadata["hook_verdict"] = pre.verdict + return denied try: result = await original_run(effective_args, ctx) except Exception as e: diff --git a/src/lecode/permission/checker.py b/src/lecode/permission/checker.py index 8650cd1..0b055e8 100644 --- a/src/lecode/permission/checker.py +++ b/src/lecode/permission/checker.py @@ -181,11 +181,11 @@ def read_only(self) -> bool: A readonly fallback with writable rule/grant exceptions is not a safe shared-checkout policy. Conservatively treat Ask as writable too. + Action-dependent tools must be treated as writable without arguments. """ if self._read_only or (self._parent is not None and self._parent.read_only): return True - mode = self._overlay.mode if self._overlay and self._overlay.mode else self._mode - if mode != "readonly": + if self._mode != "readonly" and (self._overlay is None or self._overlay.mode != "readonly"): return False rules = [self._rules] if self._overlay is not None: @@ -352,11 +352,18 @@ def _apply_doom_loop( # -- mode fallback --------------------------------------------------------- def _is_read_class(self, tool_name: str, args: dict[str, Any] | None = None) -> bool: - return ( - tool_name in READ_TOOLS - or is_read_equiv_mcp(tool_name) - or (tool_name == "workers" and args is not None and args.get("action") == "question") - ) + if tool_name == "workers": + # Control-plane replies cannot widen descendant permissions. Without + # an action, capability inference must still account for mutations. + action = args.get("action") if args is not None else None + return isinstance(action, str) and action in { + "list", + "question", + "send", + "inspect", + "review", + } + return tool_name in READ_TOOLS or is_read_equiv_mcp(tool_name) def _mode_fallback( self, mode: PermissionMode, tool_name: str, args: dict[str, Any] diff --git a/src/lecode/session/stats.py b/src/lecode/session/stats.py index 4494da2..812cc5a 100644 --- a/src/lecode/session/stats.py +++ b/src/lecode/session/stats.py @@ -9,7 +9,7 @@ from dataclasses import dataclass -from lecode.providers.catalog import Catalog, ModelNotFoundError +from lecode.providers.catalog import AmbiguousModelError, Catalog, ModelNotFoundError from lecode.session.model import EventRecord, MessageRecord, TombstoneRecord from lecode.session.storage import Session, SessionStore @@ -26,7 +26,7 @@ class Stats: created_at: str last_active: str | None tombstone_count: int - #: True when a counted worker usage event is marked incomplete. + #: True when any counted model call has missing usage or unknown cost. usage_incomplete: bool = False @@ -51,17 +51,24 @@ def session_stats(store: SessionStore, session: Session, catalog: Catalog | None for record in messages: role_counts[record.role] = role_counts.get(record.role, 0) + 1 usage = record.usage or {} + usage_incomplete |= bool(usage.get("incomplete")) or ( + record.role == "assistant" and not usage + ) in_tok, out_tok = _usage_tokens(usage) input_tokens += in_tok output_tokens += out_tok if usage.get("cost_usd") is not None: cost_usd += float(usage["cost_usd"]) - elif (in_tok or out_tok) and session.meta.model: + elif usage and (record.role == "assistant" or in_tok or out_tok): + if not session.meta.model: + usage_incomplete = True + continue if catalog is None: catalog = Catalog.default() try: pricing = catalog.get(session.meta.model).pricing - except ModelNotFoundError: + except (ModelNotFoundError, AmbiguousModelError): + usage_incomplete = True continue cost_usd += (in_tok * pricing.prompt + out_tok * pricing.completion) / 1_000_000 @@ -73,14 +80,15 @@ def session_stats(store: SessionStore, session: Session, catalog: Catalog | None usage = record.data.get("usage") or {} elif record.kind == "worker_usage": usage = record.data.get("usage") or record.data - if record.data.get("incomplete"): - usage_incomplete = True else: continue + usage_incomplete |= bool(record.data.get("incomplete") or usage.get("incomplete")) in_tok, out_tok = _usage_tokens(usage) input_tokens += in_tok output_tokens += out_tok cost_usd += float(usage.get("cost_usd") or 0.0) + if usage.get("cost_usd") is None: + usage_incomplete = True timestamps = [ r.ts for r in records if isinstance(r, MessageRecord | EventRecord | TombstoneRecord) diff --git a/src/lecode/slash/catalog.py b/src/lecode/slash/catalog.py index 05251c4..0457fa3 100644 --- a/src/lecode/slash/catalog.py +++ b/src/lecode/slash/catalog.py @@ -45,7 +45,7 @@ ("queue", "Show queued/steered messages"), ("tasks", "List background tasks"), ("runs", "Inspect agent runs"), - ("agent", "Control a worker"), + ("agent", "Control a worker; integrate WORKER_HASH PARENT_HASH after review"), ("copy", "Copy the last answer to the clipboard"), ("export", "Export the session as HTML"), ("import", "Import a session file"), diff --git a/src/lecode/slash/handlers.py b/src/lecode/slash/handlers.py index 12be7fb..41e7500 100644 --- a/src/lecode/slash/handlers.py +++ b/src/lecode/slash/handlers.py @@ -9,10 +9,13 @@ from __future__ import annotations +import json +from dataclasses import replace from pathlib import Path from typing import TYPE_CHECKING, cast, get_args from lecode.agent.tools.background import format_row +from lecode.agent.tools.workers import HUMAN_CONTROL_EXTRA from lecode.config.models import PermissionMode, ThinkingLevel from lecode.context.resources import load_text from lecode.extras.background import BACKGROUND_EXTRA @@ -762,8 +765,17 @@ async def cmd_runs(app: TuiApp, args: list[str]) -> None: app.open_agent_run(run.run_id) +_AGENT_USAGE = ( + "[id|number] [send TEXT|stop [tree]|resume [TEXT]|submit|focus|" + "inspect|integrate WORKER_HASH PARENT_HASH|cleanup|recover] | root" +) + + async def cmd_agent(app: TuiApp, args: list[str]) -> None: - """``/agent [send|stop|resume|submit|focus]``: human worker control.""" + """Human worker controls use the same tool permission gate as the model.""" + if args == ["root"]: + app.focus_worker(None) + return manager = app.worker_manager if manager is None: app.feed.info("(no workers this session)") @@ -778,7 +790,7 @@ async def cmd_agent(app: TuiApp, args: list[str]) -> None: f"{run.index}. {run.agent} · {run.status} · {run.description} · {run.run_id}" for run in workers ) - + "\n\ncontrol: /agent [send TEXT|stop [tree]|resume [TEXT]|submit|focus]" + + f"\n\ncontrol: /agent {_AGENT_USAGE}" ) return worker = app.resolve_worker(args[0]) @@ -791,27 +803,40 @@ async def cmd_agent(app: TuiApp, args: list[str]) -> None: action = args[1] rest = args[2:] try: - if action == "send": - if not rest: - raise ValueError("usage: /agent send ") - await manager.send(worker.id, " ".join(rest), from_human=True) - app.feed.info(f"sent to @{worker.agent} ({worker.id[:8]})") - elif action == "stop": - await manager.stop(worker.id, tree=bool(rest and rest[0] == "tree")) - app.feed.info(f"worker {worker.id[:8]} stopped") - elif action == "resume": - await manager.resume(worker.id, " ".join(rest) if rest else None) - app.feed.info(f"worker {worker.id[:8]} resumed") - elif action == "submit": - await app.submit_worker(worker.id) - app.feed.info(f"worker {worker.id[:8]} submitted") - elif action == "focus": + if action == "focus" and not rest: app.focus_worker(worker.id) - app.feed.info(f"composer focused on @{worker.agent} ({worker.id[:8]}); Esc returns") - else: - raise ValueError( - "usage: /agent [send TEXT|stop [tree]|resume [TEXT]|submit|focus]" + app.feed.info( + f"composer focused on @{worker.agent} ({worker.id[:8]}); " + "Esc returns to parent; /agent root returns to main" ) + return + params = {"action": action, "id": worker.id} + if action == "send" and rest: + params["text"] = " ".join(rest) + elif action == "resume": + if rest: + params["text"] = " ".join(rest) + elif action == "stop" and rest in ([], ["tree"]): + params["tree"] = bool(rest) + elif action == "integrate" and len(rest) == 2: + params["reviewed_head"] = rest[0] + params["reviewed_parent_head"] = rest[1] + elif action in {"submit", "inspect", "review", "cleanup", "recover"} and not rest: + pass + else: + raise ValueError(f"usage: /agent {_AGENT_USAGE}") + ctx = app.runtime.ctx + ctx = replace(ctx, extras={**ctx.extras, HUMAN_CONTROL_EXTRA: True}) + _, result = await ctx.extras["registry"].dispatch_result( + "human-worker-control", "workers", json.dumps(params), ctx + ) + if result.is_error: + app.feed.error(result.content) + else: + app.feed.info(result.content) + note = result.metadata.get("notification") + if note and note["new"] and note["deliver"]: + await app._on_worker_notification(note) except (KeyError, RuntimeError, SubagentError, ValueError) as e: app.feed.error(str(e)) @@ -1477,7 +1502,7 @@ def _complete_runs(app: TuiApp, args: list[str]) -> list[CompletionRow]: def _complete_agent(app: TuiApp, args: list[str]) -> list[CompletionRow]: if not args: - return [ + return [("root", "root", "return to main conversation")] + [ (run.run_id, f"{run.index} {run.agent} · {run.description}", run.status) for run in app.roster.runs() if run.worker @@ -1485,7 +1510,17 @@ def _complete_agent(app: TuiApp, args: list[str]) -> list[CompletionRow]: if len(args) == 1: return [ (action, action, "worker control") - for action in ("send", "stop", "resume", "submit", "focus") + for action in ( + "send", + "stop", + "resume", + "submit", + "focus", + "inspect", + "integrate", + "cleanup", + "recover", + ) ] if len(args) == 2 and args[1] == "stop": return [("tree", "tree", "stop descendants too")] @@ -1714,7 +1749,7 @@ def _complete_wt_exit(app: TuiApp, args: list[str]) -> list[CompletionRow]: "review": "[file…]", "notifications": "[on|off]", "runs": "[number|id]", - "agent": "[id|number] [send TEXT|stop [tree]|resume [TEXT]|submit|focus]", + "agent": _AGENT_USAGE, } diff --git a/src/lecode/tui/agents.py b/src/lecode/tui/agents.py index 2dc6d28..c3e2732 100644 --- a/src/lecode/tui/agents.py +++ b/src/lecode/tui/agents.py @@ -233,7 +233,9 @@ def finish( def cancel_running(self) -> list[AgentRun]: """Mark every still-running run cancelled (parent turn was cancelled).""" - cancelled = [run for run in self._runs.values() if run.status == "running"] + cancelled = [ + run for run in self._runs.values() if run.status == "running" and not run.worker + ] for run in cancelled: run.status = "cancelled" return cancelled diff --git a/src/lecode/tui/app.py b/src/lecode/tui/app.py index d352094..220f8bf 100644 --- a/src/lecode/tui/app.py +++ b/src/lecode/tui/app.py @@ -16,6 +16,8 @@ import contextlib import os import sys +from contextvars import ContextVar +from dataclasses import replace from pathlib import Path from typing import TYPE_CHECKING, Any @@ -102,7 +104,9 @@ ) from lecode.providers.openai_compat import ProviderError from lecode.providers.types import ContentPart +from lecode.session.model import MessageRecord from lecode.session.stats import session_stats +from lecode.slash.catalog import BUILTIN_COMMANDS from lecode.slash.handlers import build_registry from lecode.slash.registry import ( AmbiguousCommandError, @@ -169,6 +173,10 @@ EXIT_OK = 0 +class _PromptBlocked(Exception): + """Stop a generated sequence after its submission hook denies a prompt.""" + + def _register_shift_enter() -> None: """Map the Shift+Enter escape sequences to Ctrl-J (newline). @@ -243,6 +251,12 @@ def __init__( self._steer_queue: asyncio.Queue[MessageContent] = asyncio.Queue(maxsize=QUEUE_LIMIT) #: ``/btw`` side notes prepended to the next user submission. self._pending_notes: list[str] = [] + self._approved_prompt: ContextVar[tuple[str, str] | None] = ContextVar( + "approved_prompt", default=None + ) + self._command_input: ContextVar[list[str] | None] = ContextVar( + "command_input", default=None + ) #: Pending attachments for the next user submission (``/add``, ``@path``). self._attachments = AttachmentStore() #: Model catalog for ``/model``/``/models``; the live-fetched one when @@ -297,12 +311,13 @@ def __init__( self._status.input_tokens = stats.input_tokens self._status.output_tokens = stats.output_tokens self._status.cost_usd = stats.cost_usd + self._status.usage_incomplete = stats.usage_incomplete self._status.context_used = stats.context_tokens self._worker_manager = runtime.ctx.extras.get(WORKER_EXTRA) if self._worker_manager is not None: self._worker_manager.notify = self._on_worker_notification self._worker_manager.progress = self._on_worker_progress - self._hydrate_workers() + self._hydrate_workers() # Logbook suffix: the feed reads live context/cost from the statusline. self._feed.metrics = lambda: ( self._status.context_used, @@ -405,9 +420,36 @@ def catalog(self) -> Any: def submit_prompt(self, text: str, *, echo: str | None = None) -> None: """Render and queue/start a user prompt (skill commands, ``/retry``).""" - self._root_stopped = False - self._feed.user_message(text if echo is None else echo) - self._enqueue_or_start(text) + command = self._take_command_input() + + def dispatch(prompt: str) -> None: + if command is not None: + self._input_history.append_string(command) + self._root_stopped = False + self._feed.user_message(text if echo is None else echo) + self._enqueue_or_start(prompt) + + async def submit() -> None: + prompt = await self._guard_prompt(text) + if prompt is not None: + dispatch(prompt) + + hooks = self._runtime.hooks + approved = self._approved_prompt.get() + if approved is not None and approved[0] == text: + self._approved_prompt.set(None) + dispatch(approved[1]) + elif hooks is None or not hooks.handlers.get(USER_PROMPT_SUBMIT): + prompt = self._with_notes(text) + self._pending_notes.clear() + dispatch(prompt) + else: + self._spawn(submit()) + + def _take_command_input(self) -> str | None: + """Defer a submitting command's history entry until its prompt is approved.""" + pending = self._command_input.get() + return pending.pop() if pending else None def request_quit(self) -> None: """``/quit``/``/exit``: cancel any turn and leave the loop.""" @@ -497,6 +539,7 @@ def _reload_history(self) -> None: def reload_history(self) -> None: """Public wrapper used by undo/redo/rewind/clear/compact handlers.""" self._reload_history() + self._sync_usage() def turn_busy(self) -> bool: """Whether a turn or plan loop is in flight (switching commands refuse).""" @@ -504,8 +547,7 @@ def turn_busy(self) -> bool: def workers_active(self) -> bool: return self._worker_manager is not None and any( - worker.state in {"queued", "running", "waiting"} - for worker in self._worker_manager.list() + worker.is_active for worker in self._worker_manager.list() ) def resolve_worker(self, ref: str) -> Any | None: @@ -520,12 +562,18 @@ def resolve_worker(self, ref: str) -> Any | None: def focus_worker(self, worker_id: str | None) -> bool: """Switch the shared composer to a worker, retaining both drafts.""" - if worker_id is not None and self.resolve_worker(worker_id) is None: - return False + if worker_id is not None: + worker = self.resolve_worker(worker_id) + if worker is None: + return False + worker_id = worker.id current = self._focused_worker_id or "main" if self._input_area is not None: self._composer_drafts[current] = self._input_area.text self._focused_worker_id = worker_id + self._detail_run_id = worker_id + if self._roster_window is not None: + self._roster_window.vertical_scroll = 0 if self._input_area is not None: draft = self._composer_drafts.get(worker_id or "main", "") self._input_area.buffer.set_document(Document(draft, len(draft)), bypass_readonly=True) @@ -534,11 +582,17 @@ def focus_worker(self, worker_id: str | None) -> bool: def _hydrate_workers(self) -> None: """Load persisted workers once per attached root session.""" - if self._worker_manager is None: - return - for worker in self._worker_manager.load(): - self._roster.sync_worker(worker, context_window=self._status.context_window) - self._worker_usage = self._worker_totals() + if self._worker_manager is not None: + for worker in self._worker_manager.load(): + self._roster.sync_worker(worker, context_window=self._worker_context_window(worker)) + self._worker_usage = self._worker_totals() + self._sync_usage() + + def _worker_context_window(self, worker: Any) -> int: + try: + return self.catalog.get(worker.session.meta.model).context_window + except KeyError: + return self._config.agent.context_window def _worker_totals(self) -> tuple[int, int, float]: if self._worker_manager is None: @@ -553,9 +607,15 @@ def _worker_totals(self) -> tuple[int, int, float]: def current_session_usage(self) -> tuple[int, int, float, bool]: """Session totals with live worker usage substituted for persisted deltas.""" stats = session_stats(self._store, self._session) + worker_ids = ( + {worker.id for worker in self._worker_manager.list()} + if self._worker_manager is not None + else set() + ) persisted = [ event.get("usage", event) for event in self._store.load_events(self._session, "worker_usage") + if event.get("worker_id") in worker_ids ] live = self._worker_totals() recorded = ( @@ -569,6 +629,12 @@ def current_session_usage(self) -> tuple[int, int, float, bool]: if self._worker_manager is not None else stats.usage_incomplete ) + incomplete |= any( + isinstance(record, MessageRecord) + and record.role == "assistant" + and (record.usage is None or record.usage.get("incomplete", False)) + for record in self._store.read_records(self._session) + ) return ( stats.input_tokens + live[0] - recorded[0], stats.output_tokens + live[1] - recorded[1], @@ -576,14 +642,23 @@ def current_session_usage(self) -> tuple[int, int, float, bool]: incomplete, ) + def _sync_usage(self) -> None: + ( + self._status.input_tokens, + self._status.output_tokens, + self._status.cost_usd, + self._status.usage_incomplete, + ) = self.current_session_usage() + def _on_worker_progress(self, worker: Any) -> None: """WorkerManager callback: update exactly that roster entry and live totals.""" before = self._worker_usage - self._roster.sync_worker(worker, context_window=self._status.context_window) + self._roster.sync_worker(worker, context_window=self._worker_context_window(worker)) after = self._worker_totals() self._status.input_tokens += after[0] - before[0] self._status.output_tokens += after[1] - before[1] self._status.cost_usd += after[2] - before[2] + self._status.usage_incomplete |= worker.usage_incomplete self._worker_usage = after self._invalidate() @@ -592,7 +667,7 @@ async def _on_worker_notification(self, note: dict[str, Any]) -> None: worker = self.resolve_worker(note["worker_id"]) if worker is None: return - run = self._roster.sync_worker(worker, context_window=self._status.context_window) + run = self._roster.sync_worker(worker, context_window=self._worker_context_window(worker)) if note.get("kind") == "human_message": self._feed.info(f"[human → @{worker.agent} {worker.id[:8]}] {note['content']}") else: @@ -688,7 +763,18 @@ def queued_prompts(self) -> tuple[list[MessageContent], list[MessageContent]]: def _toolbar(self) -> ANSI: """The statusline as prompt_toolkit formatted text (ANSI via Rich).""" - text = render_statusline(self._status, self._theme, width=self._console.width) + state = self._status + if self._focused_worker_id is not None: + worker = self.resolve_worker(self._focused_worker_id) + if worker is not None: + state = replace( + state, + agent=worker.agent, + model=worker.session.meta.model or state.model, + context_used=worker.usage_totals.context_tokens, + context_window=self._worker_context_window(worker), + ) + text = render_statusline(state, self._theme, width=self._console.width) with self._console.capture() as capture: self._console.print(text, end="") return ANSI(capture.get()) @@ -852,8 +938,6 @@ def _enter(event: Any) -> None: self._accept_completion(buffer) return text = buffer.text - if text.strip(): - buffer.append_to_history() buffer.reset() self._spawn(self._submit(text)) @@ -862,8 +946,6 @@ def _alt_enter(event: Any) -> None: if self._approval.is_pending or self._question.is_pending: return text = event.current_buffer.text - if text.strip(): - event.current_buffer.append_to_history() event.current_buffer.reset() self._spawn(self._submit(text, steer=True)) @@ -882,6 +964,11 @@ def _ctrl_c(event: Any) -> None: if self._question.is_pending: self._question.dismiss() return + if self._focused_worker_id is not None: + worker = self.resolve_worker(self._focused_worker_id) + if worker is not None and worker.is_active: + self._spawn(self._stop_viewed_worker(worker.id)) + return if self.cancel_turn(): self._spawn(self._fire_hook(INTERRUPT)) return @@ -1314,6 +1401,9 @@ def set_catalog(self, catalog: Any, *, origin: str, count: int) -> None: self._arg_rows_cache = None # an open "/model " picker may now have rows with contextlib.suppress(Exception): # unknown model — keep the default self._status.context_window = catalog.get(self._config.llm.model).context_window + if self._worker_manager is not None: + for worker in self._worker_manager.list(): + self._roster.sync_worker(worker, context_window=self._worker_context_window(worker)) if origin == "live": self._feed.info(f"models: {count} fetched live from the provider") else: @@ -1357,45 +1447,67 @@ async def _switch_hooks(self, old: Session, new: Session) -> None: await dispatch_event(SESSION_END, envelope, end_handlers) await hooks.fire(SESSION_START) + async def _allow_prompt(self, text: str) -> bool: + verdict = await self._fire_hook(USER_PROMPT_SUBMIT, prompt=text) + if verdict is not None and verdict.verdict == "deny": + self._feed.info(f"prompt blocked by hook: {verdict.reason or 'UserPromptSubmit hook'}") + return False + return True + + def _with_notes(self, text: str) -> str: + return "\n".join(self._pending_notes) + "\n\n" + text if self._pending_notes else text + + async def _guard_prompt(self, text: str) -> str | None: + """Approve the final text, consuming only the notes included in it.""" + note_count = len(self._pending_notes) + prompt = self._with_notes(text) + if not await self._allow_prompt(prompt): + return None + del self._pending_notes[:note_count] + return prompt + async def _submit(self, text: str, *, steer: bool = False) -> None: """Route one submitted line: shell-outs, slash commands, or LLM input.""" text = text.rstrip("\n") if not text.strip(): return + focused_worker_id = self._focused_worker_id if text.startswith("!!"): await self._run_shell(text[2:].strip(), share_with_llm=True, steer=steer) elif text.startswith("!"): await self._run_shell(text[1:].strip(), share_with_llm=False, steer=steer) elif text.startswith("/"): await self.handle_command(text) - elif text.startswith(".") and self._submit_persona(text, steer=steer): + elif text.startswith(".") and await self._submit_persona(text, steer=steer): pass # .persona : handled (persona system-prompt overlay) else: - if self._focused_worker_id is not None: - worker = self.resolve_worker(self._focused_worker_id) + if focused_worker_id is not None: + worker = self.resolve_worker(focused_worker_id) if worker is None: self._feed.error("focused worker no longer exists") self.focus_worker(None) return + prompt = await self._guard_prompt(text) + if prompt is None: + return + self._input_history.append_string(text) self._feed.user_message(f"[@{worker.agent}] {text}") - await self._worker_manager.send(worker.id, text, steer, from_human=True) + await self._worker_manager.send(worker.id, prompt, steer, from_human=True) self._feed.info(f"sent to @{worker.agent} ({worker.id[:8]})") return if self.loop_running(): self._feed.info("a plan loop is running — /loop stop first") return - verdict = await self._fire_hook(USER_PROMPT_SUBMIT, prompt=text) - if verdict is not None and verdict.verdict == "deny": - self._feed.info( - f"prompt blocked by hook: {verdict.reason or 'UserPromptSubmit hook'}" - ) - return mentions, cleaned = parse_mentions(text, self._runtime.agents) invocable = {a.name for a in self._runtime.agents.subagents()} targets = [name for name in mentions if name in invocable] if targets and not self._turn_running(): + prompt = await self._guard_prompt(cleaned or text) + if prompt is None: + return + self._input_history.append_string(text) self._feed.user_message(text) - worker = await self._start_human_worker(targets[0], cleaned or text) + worker = await self._start_human_worker(targets[0], prompt) if worker is not None: # Compatibility: direct @agent work remains awaitable through # the established turn-task seam, but execution stays in the @@ -1405,9 +1517,17 @@ async def _submit(self, text: str, *, steer: bool = False) -> None: self._root_stopped = False prepared = self._prepare_message(text) echo = self._attachment_echo() + error = check_modalities(self._attachments.list(), self._runner.model, self.catalog) + if error is not None: + self._feed.error(error) + return + prepared = await self._guard_prompt(prepared) + if prepared is None: + return content = self._with_attachments(prepared) if content is None: return # modality error already rendered; attachments kept + self._input_history.append_string(text) if not self._turn_running(): # Queued messages are listed in the chatbox instead of the # transcript; they echo when the model actually sees them. @@ -1431,7 +1551,7 @@ async def _start_human_worker(self, agent: str, prompt: str) -> Any | None: except (RuntimeError, SubagentError, WorktreeError) as e: self._feed.error(str(e)) return None - self._roster.sync_worker(worker, context_window=self._status.context_window) + self._roster.sync_worker(worker, context_window=self._worker_context_window(worker)) self._feed.info( f"worker {worker.id[:8]} started for @{agent} · /agent {worker.id[:8]} focus" ) @@ -1469,7 +1589,7 @@ def _attachment_echo(self) -> str: def _prepare_message(self, text: str) -> str: """Pull ``@path`` media refs into attachments, then apply ``@agent`` - routing notes and ``/btw`` pending notes.""" + routing context.""" before = len(self._attachments) text = extract_attachment_refs(text, self._cwd, self._attachments) for attachment in self._attachments.list()[before:]: @@ -1484,13 +1604,9 @@ def _prepare_message(self, text: str) -> str: f"(The user mentioned agent(s) {names}; no subagent was " f"dispatched — treat the mention as context.)\n{cleaned}" ) - if self._pending_notes: - notes = "\n".join(self._pending_notes) - self._pending_notes = [] - text = f"{notes}\n\n{text}" return text - def _submit_persona(self, text: str, *, steer: bool) -> bool: + async def _submit_persona(self, text: str, *, steer: bool) -> bool: """``.persona ``: submit with a persona system-prompt overlay. Returns ``True`` when the input was consumed (known persona). @@ -1506,9 +1622,18 @@ def _submit_persona(self, text: str, *, steer: bool) -> bool: if not rest.strip(): self._feed.error(f"usage: .{name} ") return True + prompt = rest.strip() + overlay = body.strip() + if self._turn_running(): + prompt = f"{overlay}\n\n{prompt}" + overlay = None + prompt = await self._guard_prompt(prompt) + if prompt is None: + return True + self._input_history.append_string(text) if not self._turn_running(): self._feed.user_message(text) - self._enqueue_or_start(rest.strip(), steer=steer, overlay=body.strip()) + self._enqueue_or_start(prompt, steer=steer, overlay=overlay) return True def _enqueue_or_start( @@ -1537,6 +1662,8 @@ async def _run_shell(self, cmd: str, *, share_with_llm: bool, steer: bool) -> No """``!cmd`` shows output locally; ``!!cmd`` also feeds it to the LLM.""" if not cmd: return + if not share_with_llm: + self._input_history.append_string("!" + cmd) self._feed.user_message(("!!" if share_with_llm else "!") + cmd) self._activity("running shell") self._shell_task = asyncio.current_task() @@ -1555,7 +1682,10 @@ async def _run_shell(self, cmd: str, *, share_with_llm: bool, steer: bool) -> No is_error = result.exit_code != 0 or result.timed_out self._feed.tool_result("shell", output or "(no output)", is_error=is_error) if share_with_llm: - self._enqueue_or_start(f"!{cmd}\n\n{output}", steer=steer) + prompt = await self._guard_prompt(f"!{cmd}\n\n{output}") + if prompt is not None: + self._input_history.append_string("!!" + cmd) + self._enqueue_or_start(prompt, steer=steer) async def handle_command(self, text: str) -> None: """Dispatch a ``/command`` line through the slash registry.""" @@ -1566,13 +1696,52 @@ async def handle_command(self, text: str) -> None: try: command = self._commands.match(query) except UnknownCommandError: + self._input_history.append_string(text) self._feed.info(f"unknown command: /{query}") return except AmbiguousCommandError as e: + self._input_history.append_string(text) matches = ", ".join(f"/{name}" for name in e.matches) self._feed.info(f"ambiguous command: /{query} ({matches})") return - await command.handler(self, args) + if ( + command.name == "agent" + and len(args) > 2 + and args[1] in {"send", "resume"} + and not await self._allow_prompt(" ".join(args[2:])) + ): + return + approved_prompt = None + if command.name == "retry" and not self.turn_busy(): + last_user = next( + (m for m in reversed(self._store.load_messages(self._session)) if m.role == "user"), + None, + ) + content = last_user.message.get("content") if last_user is not None else None + if isinstance(content, str): + prompt = await self._guard_prompt(content) + if prompt is None: + return + approved_prompt = (content, prompt) + # /retry mutates history before submitting; approve before that mutation, + # then let its exact prompt consume this task-local approval once. + token = self._approved_prompt.set(approved_prompt) + defer_history = command.name in {"retry", "chain", "loop", "review", "redo"} or ( + command.name not in {name for name, _ in BUILTIN_COMMANDS} + ) + # Non-submitting commands retain their original ordering: undo/rewind + # must leave a tombstone last, and session switches log in the old session. + if not defer_history: + self._input_history.append_string(text) + command_token = self._command_input.set([text] if defer_history else []) + try: + await command.handler(self, args) + remaining = self._take_command_input() + if remaining is not None: + self._input_history.append_string(remaining) + finally: + self._approved_prompt.reset(token) + self._command_input.reset(command_token) # -- inline permission prompt ---------------------------------------------- @@ -1649,33 +1818,15 @@ async def _request_approval( ) self._invalidate() - async def _confirm_worker_worktree(self, *args: Any, **kwargs: Any) -> bool: + async def _confirm_worker_worktree(self, question: str) -> bool: """Ask the TUI user before a dirty write worker gets a separate worktree.""" - question = str( - kwargs.get("question") or next((arg for arg in args if isinstance(arg, str)), "") - ) - worker = kwargs.get("worker") - worktree = kwargs.get("worktree") or kwargs.get("cwd") - for arg in args: - if not isinstance(arg, str) and worker is None: - worker = arg - elif not isinstance(arg, str) and worktree is None: - worktree = arg - worker_id = getattr(worker, "id", None) or getattr(worker, "worker_id", None) or "new" - agent = getattr(worker, "agent", None) or kwargs.get("agent") or "write" - path = getattr(worktree, "path", worktree) or self._cwd - target = f"@{agent} worker {str(worker_id)[:8]} in {path}" future = self._approval.request( "dirty worktree", - target, question, - worker=str(worker_id), - conversation=str(path), + "", allow_always=False, ) self._show_approval_head() - if question: - self._feed.info(question) self._status.state = StatusLineState.AWAITING_APPROVAL self._invalidate() try: @@ -1688,6 +1839,8 @@ async def _confirm_worker_worktree(self, *args: Any, **kwargs: Any) -> bool: StatusLineState.AWAITING_APPROVAL if self._approval.is_pending else StatusLineState.RUNNING + if self.turn_busy() + else StatusLineState.IDLE ) self._invalidate() @@ -1734,6 +1887,11 @@ def cancel_turn(self) -> bool: return True return False + async def _stop_viewed_worker(self, worker_id: str) -> None: + await self._worker_manager.stop(worker_id) + await self._fire_hook(INTERRUPT) + self._invalidate() + def cancel_action(self) -> bool: """Cancel a non-turn action (plan loop, ``!cmd`` shell-out); ``True`` if one was cancelled. Checked by Ctrl-C after :meth:`cancel_turn`.""" @@ -1814,13 +1972,12 @@ async def _run_turn(self, text: MessageContent, *, overlay: str | None = None) - self._signals.emit(STOP) self._status.state = StatusLineState.IDLE self._status.activity = None + self._sync_usage() + self._status.usage_incomplete |= cancelled if result is not None: if result.final_text: self._last_response = result.final_text totals = result.usage_totals - self._status.input_tokens += totals.input_tokens - self._status.output_tokens += totals.output_tokens - self._status.cost_usd += totals.cost_usd # Last API call's prompt size — the real context fill. self._status.context_used = totals.context_tokens or totals.input_tokens self._feed.turn_stats( @@ -1886,8 +2043,16 @@ async def _run_subagent_turn(self, name: str, prompt: str) -> None: def start_loop(self, plan_path: Path, max_iterations: int) -> None: """``/loop``: run a plan loop over this session in the background.""" + command = self._take_command_input() async def run_iteration(prompt: str) -> str: + nonlocal command + prompt = await self._guard_prompt(prompt) + if prompt is None: + raise _PromptBlocked + if command is not None: + self._input_history.append_string(command) + command = None message: dict[str, Any] = {"role": "user", "content": prompt} self._history.append(message) self._store.append_message(self._session, message) @@ -1909,6 +2074,8 @@ async def _loop() -> None: max_iterations=max_iterations, on_progress=self._feed.info, ) + except _PromptBlocked: + return except asyncio.CancelledError: self._feed.info("loop stopped") return @@ -1933,15 +2100,28 @@ async def _loop() -> None: def start_chain(self, topic: str) -> None: """``/chain``: brainstorm → plan → code → review as the turn task.""" - self._turn_task = asyncio.ensure_future(self._run_chain(topic)) + self._turn_task = asyncio.ensure_future(self._run_chain(topic, self._take_command_input())) - async def _run_chain(self, topic: str) -> None: + async def _run_chain(self, topic: str, command: str | None = None) -> None: """Run the chain; each phase renders as it completes (not persisted).""" + app = self + + class GuardedRunner(AgentRunner): + async def run(self, messages: list[dict[str, Any]], **kwargs: Any) -> Any: + nonlocal command + prompt = await app._guard_prompt(messages[-1]["content"]) + if prompt is None: + raise _PromptBlocked + if command is not None: + app._input_history.append_string(command) + command = None + messages[-1] = {"role": "user", "content": prompt} + return await super().run(messages, **kwargs) def factory() -> AgentRunner: # Fresh runner per phase; no session binding → chain is a side # computation rendered to the feed, not written to the session. - return AgentRunner( + return GuardedRunner( self._runner.provider, self._runtime.registry, self._runtime.ctx, @@ -1961,6 +2141,8 @@ def on_phase(phase: str, output: str) -> None: cwd=self._cwd, on_phase=on_phase, ) + except _PromptBlocked: + return except asyncio.CancelledError: self._feed.info("turn cancelled") return @@ -2012,6 +2194,10 @@ def _on_event(self, event: Any) -> None: self._feed.llm_call(event.model, event.turn) self._activity("thinking") elif isinstance(event, LlmResponse): + self._status.input_tokens += event.input_tokens + self._status.output_tokens += event.output_tokens + self._status.cost_usd += event.cost_usd + self._status.usage_incomplete |= getattr(event, "usage_incomplete", False) if event.input_tokens > 0 and event.prompt_chars > 0: # Calibrate the live estimate: EMA of chars/token, clamped. ratio = min(max(event.prompt_chars / event.input_tokens, 2.0), 8.0) @@ -2023,6 +2209,7 @@ def _on_event(self, event: Any) -> None: event.output_tokens, event.cost_usd, ) + self._invalidate() elif isinstance(event, Done): self._feed.stream_end() self._spawn(self._notifier.task_finish()) @@ -2047,7 +2234,7 @@ def _on_child_event(self, progress: SubagentProgress) -> None: parent's own stream. """ run = self._roster.observe(progress) - if isinstance(progress.event, (Error, Done)): + if not run.worker and isinstance(progress.event, (Error, Done)): self._feed.agent_summary(run) self._invalidate() @@ -2076,8 +2263,12 @@ def detail_run_id(self) -> str | None: def open_agent_run(self, run_id: str) -> bool: """Show one run's live detail panel; ``False`` for an unknown run.""" - if self._roster.get(run_id) is None: + run = self._roster.get(run_id) + if run is None: return False + if run.worker: + return self.focus_worker(run_id) + self.focus_worker(None) self._detail_run_id = run_id if self._roster_window is not None: self._roster_window.vertical_scroll = 0 @@ -2085,8 +2276,7 @@ def open_agent_run(self, run_id: str) -> bool: return True def close_agent_run(self) -> None: - self._detail_run_id = None - self._invalidate() + self.focus_worker(None) def _roster_visible(self) -> bool: return self._detail_run_id is not None or self._roster.has_running() @@ -2134,8 +2324,8 @@ def cycle_agent(self) -> str: def print_totals(self) -> None: """Cost-reporting requirement: full-session totals on exit.""" - stats = session_stats(self._store, self._session) + input_tokens, output_tokens, cost, incomplete = self.current_session_usage() self._feed.info( - f"Session {self._session.name}: tokens {stats.input_tokens} in / " - f"{stats.output_tokens} out · cost ${stats.cost_usd:.4f}" + f"Session {self._session.name}: tokens {input_tokens} in / " + f"{output_tokens} out · cost ${cost:.4f}" + (" incomplete" if incomplete else "") ) diff --git a/src/lecode/tui/permission.py b/src/lecode/tui/permission.py index 840f0b0..1eeb64d 100644 --- a/src/lecode/tui/permission.py +++ b/src/lecode/tui/permission.py @@ -36,9 +36,9 @@ class PendingApproval: def approval_prompt_text(tool_name: str, target: str, *, allow_always: bool = True) -> str: """The one-line ask: ``allow bash 'ls'? (y)once (a)lways (n)deny — ESC denies``.""" - shown = target if len(target) <= _TARGET_MAX_LEN else target[: _TARGET_MAX_LEN - 1] + "…" if not allow_always: - return f"confirm {tool_name} '{shown}'? (y)es (n)o — ESC denies" + return f"{target}\n(y)es (n)o — ESC denies" + shown = target if len(target) <= _TARGET_MAX_LEN else target[: _TARGET_MAX_LEN - 1] + "…" return f"allow {tool_name} '{shown}'? (y)once (a)lways (n)deny — ESC denies" diff --git a/src/lecode/tui/statusline.py b/src/lecode/tui/statusline.py index bbfc66c..0b7a132 100644 --- a/src/lecode/tui/statusline.py +++ b/src/lecode/tui/statusline.py @@ -3,7 +3,7 @@ Three-line layout, every element labelled:: dir: · commit: · branch: · diff: - model: · cost: <$0.00> · ctx: ▓▓▓░░ 84.0k/200k 42% + total cost: <$0.00> · ctx: ▓▓▓░░ 84.0k/200k 42% · model: session: · agent: · in: 1.2k · out: 0.4k · Not user-configurable. Lines are truncated to the terminal width. @@ -68,6 +68,7 @@ class StatusState: input_tokens: int = 0 output_tokens: int = 0 cost_usd: float = 0.0 + usage_incomplete: bool = False state: StatusLineState = StatusLineState.IDLE queued: int = 0 steered: int = 0 @@ -148,20 +149,26 @@ def labelled(line: Text, label: str, value: str, style: str, *, first: bool = Fa if value: labelled(line1, label, value, theme.muted) - # Line 2: model · cost · ctx meter x/y pct% + # Keep total cost and missing usage visible even with a long model name. bar, pct = context_meter(state.context_used, state.context_window) line2 = Text() - labelled(line2, "model", state.model, theme.text, first=True) - if state.reasoning is not None: - line2.append_text(sep.copy()) - line2.append(state.reasoning, style=theme.muted) - labelled(line2, "cost", format_cost(state.cost_usd), theme.muted) + labelled( + line2, + "total cost", + format_cost(state.cost_usd) + (" incomplete" if state.usage_incomplete else ""), + theme.warning if state.usage_incomplete else theme.muted, + first=True, + ) labelled( line2, "ctx", f"{bar} {human_tokens(state.context_used)}/{human_tokens(state.context_window)} {pct}%", theme.text, ) + labelled(line2, "model", state.model, theme.text) + if state.reasoning is not None: + line2.append_text(sep.copy()) + line2.append(state.reasoning, style=theme.muted) # Line 3: session · agent · in/out tokens · state state_seg, state_color = _state_segment(state) diff --git a/tests/test_agent_builder.py b/tests/test_agent_builder.py index 4ec947d..db1e16f 100644 --- a/tests/test_agent_builder.py +++ b/tests/test_agent_builder.py @@ -14,6 +14,7 @@ def cwd(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) return tmp_path @@ -31,6 +32,21 @@ def test_default_runtime(cwd): assert runtime.registry.names() == sorted(expected) assert runtime.ctx.auto_approve is False assert runtime.system_prompt.startswith("You are lecode") + assert "- general: General-purpose coding subagent" in runtime.system_prompt + assert "- explore: Fast read-only" in runtime.system_prompt + + +def test_subagent_discovery_respects_overrides(cwd): + from tests.test_agents import write_agent + + agents = cwd / ".lecode" / "agents" + write_agent(agents, "general", "description: Custom primary\nmode: primary") + write_agent(agents, "secret", "description: Hidden\nmode: subagent\nhidden: true") + write_agent(agents, "helper", "description: Custom helper\nmode: subagent") + runtime = build_runtime(Config(), cwd) + assert "- general:" not in runtime.system_prompt + assert "- secret:" not in runtime.system_prompt + assert "- helper: Custom helper" in runtime.system_prompt def test_read_only_mode_denies_writes(cwd): @@ -78,15 +94,14 @@ def test_session_runtime_installs_workers(cwd, tmp_path): assert runtime.ctx.extras["workers"].session is session -def test_workers_schema_omits_unimplemented_integration_actions(cwd, tmp_path): +def test_workers_schema_exposes_reviewed_integration_controls(cwd, tmp_path): from lecode.session.storage import SessionStore store = SessionStore(config_dir=tmp_path / "cfg") session = store.create("workers", cwd) runtime = build_runtime(Config(), cwd, session=session, store=store) actions = runtime.registry.get("workers").parameters["properties"]["action"]["enum"] - assert "integrate" not in actions - assert "cleanup" not in actions + assert {"review", "integrate", "cleanup", "recover"} <= set(actions) def test_agent_name_applies_overlay(cwd): diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 1351f91..18825fd 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -3,10 +3,11 @@ from __future__ import annotations import asyncio -from contextlib import asynccontextmanager +import json import pytest from tests.fakes import FakeProvider, sample_catalog +from tests.test_workers import setup as setup from lecode.agent.runner import ( CONTINUE_PROMPT, @@ -26,6 +27,7 @@ ) from lecode.agent.tools.base import Tool, ToolRegistry from lecode.agent.tools.base import ToolResult as ToolExecResult +from lecode.extras.background import BACKGROUND_EXTRA from lecode.providers.openai_compat import ProviderError from lecode.providers.types import TokenDelta from lecode.session.model import EventRecord, MessageRecord @@ -57,25 +59,6 @@ async def run(self, args, ctx) -> ToolExecResult: raise AssertionError("unreachable") -class TaskLikeTool(EchoTool): - def __init__(self) -> None: - super().__init__() - self.name = "task" - - -class LeaseManager: - def __init__(self) -> None: - self.suspensions = 0 - - def consume(self, id, history): - return [] - - @asynccontextmanager - async def suspend(self, id): - self.suspensions += 1 - yield - - def make_runner(tool_ctx, script, **kwargs) -> tuple[AgentRunner, FakeProvider]: provider = FakeProvider(script) kwargs.setdefault("catalog", sample_catalog()) @@ -156,20 +139,6 @@ async def test_tool_round_trip(tool_ctx): assert result_events[0].is_error is False -@pytest.mark.parametrize("names,expected", [(["task"], 1), (["task", "echo"], 0)]) -async def test_worker_suspends_only_all_task_batches(tool_ctx, names, expected): - manager = LeaseManager() - tool_ctx.extras.update({"workers": manager, "worker_id": "worker"}) - registry = ToolRegistry([TaskLikeTool(), EchoTool()]) - calls = [ - {"id": f"c{i}", "name": name, "arguments": '{"text":"ok"}'} for i, name in enumerate(names) - ] - provider = FakeProvider([{"tool_calls": calls}, {"text": "done"}]) - runner = AgentRunner(provider, registry, tool_ctx) - assert (await runner.run([{"role": "user", "content": "go"}])).final_text == "done" - assert manager.suspensions == expected - - async def test_parallel_tool_calls_paired_by_id(tool_ctx): script = [ { @@ -417,6 +386,300 @@ async def test_steer_queue_drained_before_input_queue(tool_ctx): assert drained == ["steer me", "regular input"] +async def test_root_resume_repairs_interrupted_tool_before_worker_notification(setup): + manager, ctx, _, store, session = setup + + class Background: + def drain_notifications(self): + return ["background finished"] + + ctx.extras[BACKGROUND_EXTRA] = Background() + store.append_message(session, {"role": "user", "content": "resume"}) + store.append_message( + session, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "interrupted", + "type": "function", + "function": {"name": "task", "arguments": "{}"}, + } + ], + }, + ) + store.append_event( + session, + "worker_notification", + { + "id": "finished", + "worker_id": "child", + "parent_id": None, + "state": "completed", + "kind": "completed", + "content": "finished work", + "deliver": True, + }, + ) + provider = FakeProvider([{"text": "recovered"}]) + runner = AgentRunner(provider, ctx.extras["registry"], ctx, session=session, store=store) + + result = await runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)] + ) + + assert result.final_text == "recovered" + messages = provider.requests[0]["messages"] + repaired = next(message for message in messages if message["role"] == "tool") + assert repaired["tool_call_id"] == "interrupted" + assert "outcome unknown" in repaired["content"] + assert messages.index(repaired) < next( + i for i, message in enumerate(messages) if message.get("content") == "background finished" + ) + assert messages[-1]["content"] == "[worker child completed] finished work" + assert not manager._outstanding(store.load_for_model(session)) + + +@pytest.mark.parametrize("queue_name", ["input_queue", "steer_queue"]) +async def test_child_question_root_human_reply_wakes_completion(setup, queue_name): + manager, ctx, _, store, session = setup + asked_human = asyncio.Event() + queue = asyncio.Queue() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + assert not manager._outstanding(messages) + outstanding = set() + for message in messages: + if message["role"] == "tool": + outstanding.remove(message["tool_call_id"]) + else: + assert not outstanding, "human reply preceded tool results" + outstanding.update(c["id"] for c in message.get("tool_calls", [])) + if messages[1]["content"] == "child": + if len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "ask", + "name": "workers", + "arguments": '{"action":"question","text":"which color?"}', + } + ] + } + else: + assert messages[-1]["content"] == "use blue" + entry = {"text": "blue result"} + elif len(messages) == 2: + entry = { + "tool_calls": [ + {"id": "spawn", "name": "task", "arguments": '{"prompt":"child"}'} + ] + } + elif any("blue result" in str(m.get("content")) for m in messages): + entry = {"text": "reviewed blue result"} + elif any(m.get("tool_call_id") == "reply" for m in messages): + entry = {"text": "waiting for child"} + elif any(m.get("content") == "use blue" for m in messages): + entry = { + "tool_calls": [ + { + "id": "reply", + "name": "workers", + "arguments": json.dumps( + {"action": "send", "id": manager.list()[0].id, "text": "use blue"} + ), + } + ] + } + else: + assert any("which color?" in str(m.get("content")) for m in messages) + entry = {"text": "Human, which color?"} + async for event in self._stream(entry): + yield event + + def on_event(event): + if isinstance(event, Token) and event.text == "Human, which color?": + asked_human.set() + + store.append_message(session, {"role": "user", "content": "root"}) + runner = AgentRunner( + Provider([]), + ctx.extras["registry"], + ctx, + session=session, + store=store, + **{queue_name: queue}, + ) + task = asyncio.create_task( + runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)], on_event + ) + ) + try: + async with asyncio.timeout(2): + await asked_human.wait() + assert not task.done() + queue.put_nowait("use blue") + result = await task + assert result.final_text == "reviewed blue result" + assert manager.list()[0].result.final_text == "blue result" + assert manager.list()[0].state == "completed" + assert [m["content"] for m in store.load_for_model(session)].count("use blue") == 1 + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await manager.shutdown() + + +@pytest.mark.parametrize("supervisor", ["root", "worker"]) +@pytest.mark.parametrize("max_turns", [1, 2]) +async def test_last_turn_background_spawn_reports_limit_with_unresolved_child( + setup, supervisor, max_turns +): + from lecode.extras.subagents import SubagentError + + manager, ctx, _, store, session = setup + ctx.config.agent.max_turns = max_turns + release = asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + if messages[1]["content"] == "child": + await release.wait() + entry = {"text": "child done"} + elif len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "spawn", + "name": "task", + "arguments": '{"prompt":"child","run_in_background":true}', + } + ], + "usage": {"input_tokens": 5, "cost_usd": 0.25}, + } + else: + entry = {"text": "premature final"} + async for event in self._stream(entry): + yield event + + provider = Provider([]) + ctx.extras["provider"] = provider + events = [] + try: + async with asyncio.timeout(2): + if supervisor == "worker": + parent = await manager.start(ctx, agent="explore", prompt="parent") + with pytest.raises(SubagentError, match="max_turns"): + await manager.wait(parent.id) + result = parent.result + assert parent.state == "failed" + child = manager.children(parent.id)[0] + assert child.id in parent.error + stopped_session = parent.session + else: + store.append_message(session, {"role": "user", "content": "root"}) + runner = AgentRunner( + provider, ctx.extras["registry"], ctx, session=session, store=store + ) + result = await runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)], + events.append, + ) + child = manager.children(None)[0] + assert any(isinstance(e, Error) and child.id in e.message for e in events) + assert events[-1] == Done("max_turns", max_turns) + stopped_session = session + assert result.stop_reason == "max_turns" + assert result.turns == max_turns + assert result.usage_totals.cost_usd == 0.25 + assert result.final_text == ("premature final" if max_turns == 2 else "") + assert child.is_active + stopped = store.load_events(stopped_session, "run_stopped")[-1] + assert stopped["reason"] == "max_turns" and child.id in stopped["message"] + release.set() + await manager.wait(child.id) + finally: + release.set() + await manager.shutdown() + + +@pytest.mark.parametrize("cancel", [False, True]) +async def test_completion_queue_race_preserves_priority_inputs_and_cleans_waiters(setup, cancel): + from tests.test_workers import GatedProvider + + manager, ctx, _, store, session = setup + child_provider = GatedProvider() + ctx.extras["provider"] = child_provider + + class Queue(asyncio.Queue): + def __init__(self): + super().__init__() + self.waiting = asyncio.Event() + self.received = asyncio.Event() + self.waiters = [] + + async def get(self): + self.waiters.append(asyncio.current_task()) + self.waiting.set() + item = await super().get() + self.received.set() + return item + + steer, inputs = Queue(), Queue() + child = await manager.start(ctx, agent="explore", prompt="child") + await child_provider.started.get() + provider = FakeProvider([{"text": "waiting"}, {"text": "reviewed"}]) + runner = AgentRunner( + provider, + ctx.extras["registry"], + ctx, + session=session, + store=store, + steer_queue=steer, + input_queue=inputs, + ) + task = asyncio.create_task(runner.run([{"role": "user", "content": "root"}])) + try: + async with asyncio.timeout(2): + await steer.waiting.wait() + await inputs.waiting.wait() + for queue, prefix in ((inputs, "input"), (steer, "steer")): + queue.put_nowait(f"{prefix} one") + queue.put_nowait(f"{prefix} two") + await inputs.received.wait() + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + child_provider.release.set() + await task + texts = [ + m["content"] + for m in store.load_for_model(session) + if m["role"] == "user" and not m["content"].startswith("[worker") + ] + assert texts == ["steer one", "steer two", "input one", "input two"] + assert all(waiter.done() for queue in (steer, inputs) for waiter in queue.waiters) + assert steer.empty() and inputs.empty() + if not cancel: + assert [ + m["content"] + for m in provider.requests[1]["messages"] + if m["role"] == "user" and not m["content"].startswith("[worker") + ] == ["root", *texts] + child_provider.release.set() + await manager.wait(child.id) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + child_provider.release.set() + await manager.shutdown() + + # -- automatic compaction --------------------------------------------------------- @@ -623,3 +886,119 @@ async def test_mid_turn_threshold_triggers_earlier(tool_ctx, tmp_path): assert _compact_events(store, session) started = [e for e in events if isinstance(e, CompactionStarted)] assert [(e.context_tokens, e.threshold) for e in started] == [(600, 500)] + + +@pytest.mark.parametrize( + ("model", "usage", "incomplete"), + [ + ("missing/model", {"input_tokens": 7, "output_tokens": 2}, True), + ("missing/model", {"input_tokens": 0, "output_tokens": 0}, True), + ("missing/model", {"input_tokens": 7, "output_tokens": 2, "cost_usd": 0}, False), + ("missing/model", {"cost_usd": 0}, False), + ("openai/gpt-5-", {"input_tokens": 7}, True), + ("openai/gpt-5-", {"cost_usd": 0}, False), + ("free/model", {"input_tokens": 7, "output_tokens": 2}, False), + ("free/model", {"input_tokens": 0, "output_tokens": 0}, False), + ("free/model", None, True), + ("free/model", {"cost_usd": 0, "incomplete": True}, True), + ], +) +async def test_usage_completeness_reaches_events_totals_and_storage( + tool_ctx, tmp_path, monkeypatch, model, usage, incomplete +): + from lecode.providers.catalog import Pricing + from lecode.session.stats import session_stats + + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + catalog = sample_catalog() + free = catalog.all()[0].model_copy( + update={"id": "free/model", "pricing": Pricing(prompt=0, completion=0)} + ) + tool_ctx.catalog = catalog.merge([free]) + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("usage-completeness", tmp_path, model=model) + runner, _ = make_runner( + tool_ctx, [{"text": "done", "usage": usage}], catalog=None, store=store, session=session + ) + runner.model = model + events = [] + result = await runner.run([{"role": "user", "content": "hi"}], events.append) + response = next(event for event in events if isinstance(event, LlmResponse)) + recorded = store.load_messages(session)[0].usage + stats = session_stats(store, session, catalog=tool_ctx.catalog) + + assert response.usage_incomplete is incomplete + assert result.usage_totals.usage_incomplete is incomplete + assert bool(recorded.get("incomplete")) is incomplete + assert stats.usage_incomplete is incomplete + assert response.cost_usd == result.usage_totals.cost_usd == stats.cost_usd == 0 + assert response.input_tokens == stats.input_tokens == (usage or {}).get("input_tokens", 0) + assert response.output_tokens == stats.output_tokens == (usage or {}).get("output_tokens", 0) + + +@pytest.mark.parametrize("usage", [None, {"input_tokens": 9, "output_tokens": 2, "cost_usd": 0.25}]) +@pytest.mark.parametrize("text", ["partial", ""]) +async def test_cancel_preserves_known_usage_or_unknown_marker( + tool_ctx, tmp_path, monkeypatch, usage, text +): + from lecode.providers.types import Usage + from lecode.session.stats import session_stats + + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + ready = asyncio.Event() + + class PartialProvider: + async def stream_chat(self, *args, **kwargs): + if text: + yield TokenDelta(text=text) + if usage is not None: + yield Usage(usage=usage) + ready.set() + await asyncio.Event().wait() + + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("partial-usage", tmp_path, model=tool_ctx.config.llm.model) + runner = AgentRunner(PartialProvider(), ToolRegistry(), tool_ctx, store=store, session=session) + task = asyncio.create_task(runner.run([{"role": "user", "content": "hi"}])) + await ready.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + recorded = store.load_messages(session) + assert len(recorded) == 1 + assert recorded[0].usage == ( + usage or {"input_tokens": 0, "output_tokens": 0, "cost_usd": 0, "incomplete": True} + ) + stats = session_stats(store, session) + assert stats.usage_incomplete is (usage is None) + assert stats.cost_usd == (usage or {}).get("cost_usd", 0) + + +@pytest.mark.parametrize("review_usage", [None, {"input_tokens": 3}, {"cost_usd": 0}]) +async def test_pierre_usage_completeness_is_persisted( + tool_ctx, tmp_path, monkeypatch, review_usage +): + from lecode.session.stats import session_stats + + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + tool_ctx.config.pierre.enabled = True + tool_ctx.config.pierre.model = "missing/reviewer" + store = SessionStore(config_dir=tmp_path / "cfg") + session = store.create("review-usage", tmp_path, model=tool_ctx.config.llm.model) + runner, _ = make_runner( + tool_ctx, + [{"text": "done", "usage": {"cost_usd": 0.25}}, {"text": "review", "usage": review_usage}], + store=store, + session=session, + ) + result = await runner.run([{"role": "user", "content": "hi"}]) + incomplete = review_usage != {"cost_usd": 0} + assert result.review == "review" + assert result.usage_totals.usage_incomplete is incomplete + assert store.load_events(session, "pierre")[0]["usage"]["incomplete"] is incomplete + stats = session_stats(store, session) + assert stats.usage_incomplete is incomplete + assert stats.cost_usd == result.usage_totals.cost_usd == 0.25 diff --git a/tests/test_agents.py b/tests/test_agents.py index 9fa4e06..759a62b 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -25,10 +25,13 @@ def write_agent(root, name: str, frontmatter: str, body: str = "Agent prompt.") def test_builtins_present_without_files(dirs): _, _, project = dirs registry = load_agents(cwd=project) - assert {"build", "plan", "explore"} <= set(registry.names()) + assert {"build", "plan", "explore", "general"} == set(registry.names()) assert registry.get("build").mode == "primary" assert registry.get("plan").mode == "primary" assert registry.get("explore").mode == "subagent" + assert registry.get("general").mode == "subagent" + assert registry.get("general").builtin + assert registry.overlay_for("general") is None def test_builtin_overlays(dirs): @@ -58,16 +61,17 @@ def test_project_wins_on_collision(dirs): assert registry.get("same").description == "project version" -def test_user_file_overrides_builtin(dirs): +@pytest.mark.parametrize("name", ["plan", "general"]) +def test_user_file_overrides_builtin(dirs, name): _, _, project = dirs write_agent( project / ".lecode" / "agents", - "plan", + name, "description: custom plan\nmode: primary", body="My own planner.", ) registry = load_agents(cwd=project) - plan = registry.get("plan") + plan = registry.get(name) assert plan.description == "custom plan" assert plan.body == "My own planner." assert plan.overlay is None # user file replaces the built-in wholesale @@ -197,7 +201,7 @@ def test_subagents_listing(dirs): write_agent(project / ".lecode" / "agents", "helper", "description: h\nmode: all") write_agent(project / ".lecode" / "agents", "sub", "description: s\nmode: subagent") registry = load_agents(cwd=project) - assert [a.name for a in registry.subagents()] == ["explore", "helper", "sub"] + assert [a.name for a in registry.subagents()] == ["explore", "general", "helper", "sub"] def test_cycle_wraps_around(dirs): diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 5bd70d0..8741593 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -42,6 +42,7 @@ def test_defaults_validate_from_empty(): assert config.compaction.buffer_tokens == 20000 assert config.agent.max_turns == 500 assert config.tools.enabled == {} + assert config.worktree.validation == [] assert config.permissions.mode == "yolo" assert config.notifications.volume == 0.5 assert config.mcp.enable_exa is True @@ -205,3 +206,24 @@ def test_deep_merge_semantics(): "c": [9], "d": True, } + + +def test_worktree_validation_config(global_dir, project): + root, cwd = project + (global_dir / "config.toml").write_text( + 'schema_version = 1\n[worktree]\nvalidation = ["check-one", "check-two"]\n' + ) + loaded = load_config(cwd=cwd) + assert loaded.config.worktree.validation == ["check-one", "check-two"] + assert loaded.warnings == [] + (root / ".lecode").mkdir() + (root / ".lecode" / "config.toml").write_text("[worktree]\nvalidation = []\n") + assert load_config(cwd=cwd).config.worktree.validation == [] + assert "validation = []" in (root / ".lecode" / "config.toml").read_text() + + +@pytest.mark.parametrize("value", ['"check"', "[1]", "true"]) +def test_worktree_validation_rejects_wrong_types(global_dir, tmp_path, value): + (global_dir / "config.toml").write_text(f"[worktree]\nvalidation = {value}\n") + with pytest.raises(ValueError, match=r"worktree\.validation"): + load_config(cwd=tmp_path) diff --git a/tests/test_hooks_decorator.py b/tests/test_hooks_decorator.py index cc946f4..77e00aa 100644 --- a/tests/test_hooks_decorator.py +++ b/tests/test_hooks_decorator.py @@ -147,6 +147,63 @@ async def test_rewritten_input_reaches_tool(tmp_path): assert "rewritten" in message["content"] +async def test_hook_cannot_rewrite_readonly_inspection_into_cleanup(tmp_path): + tool = RecordingTool() + tool.name = "workers" + rewritten = {"action": "cleanup", "id": "child"} + hook = f"echo '{json.dumps({'verdict': 'allow', 'rewritten_input': rewritten})}'" + registry = apply_hooks(ToolRegistry([tool]), dispatcher(tmp_path, pre=[hook])) + ctx = make_ctx(tmp_path, mode="readonly", auto_approve=True) + _, result = await registry.dispatch_result( + "c1", "workers", json.dumps({"action": "inspect", "id": "child"}), ctx + ) + assert result.is_error + assert "denied" in result.content + assert tool.calls == [] + + +async def test_rewritten_target_requires_its_own_approval(tmp_path): + from lecode.config.models import PermissionRule + from lecode.permission import Deny + + tool = RecordingTool() + tool.name = "bash" + config = Config() + config.permissions.rules.ask["bash"] = [PermissionRule(pattern="sensitive")] + seen = [] + + async def refuse(name, args, reason): + seen.append((name, args)) + return Deny() + + hook = f"echo '{json.dumps({'verdict': 'allow', 'rewritten_input': {'command': 'sensitive'}})}'" + registry = apply_hooks(ToolRegistry([tool]), dispatcher(tmp_path, pre=[hook])) + ctx = make_ctx(tmp_path, config=config, callback=refuse) + _, result = await registry.dispatch_result("c1", "bash", '{"command": "safe"}', ctx) + assert result.is_error + assert seen == [("bash", {"command": "sensitive"})] + assert tool.calls == [] + + +async def test_hook_ask_uses_interactive_approver(tmp_path): + from lecode.permission import AllowOnce + + tool = RecordingTool() + seen = [] + + async def approve(name, args, reason): + seen.append((name, args, reason)) + return AllowOnce() + + hook = f"echo '{json.dumps({'verdict': 'ask', 'reason': 'check before execution'})}'" + registry = apply_hooks(ToolRegistry([tool]), dispatcher(tmp_path, pre=[hook])) + ctx = make_ctx(tmp_path, callback=approve) + _, result = await registry.dispatch_result("c1", "rec", "{}", ctx) + assert not result.is_error + assert seen == [("rec", {}, "check before execution")] + assert tool.calls == [{}] + + async def test_defer_runs_with_original_args(tmp_path): tool = RecordingTool() defer = f"echo '{json.dumps({'verdict': 'defer'})}'" diff --git a/tests/test_permission_checker.py b/tests/test_permission_checker.py index 41183e3..df5846c 100644 --- a/tests/test_permission_checker.py +++ b/tests/test_permission_checker.py @@ -289,24 +289,77 @@ def test_read_only_denies_writes_even_in_yolo(): assert "read-only" in checker.check("bash", {"command": "ls"}).reason +@pytest.mark.parametrize("strict", [False, True]) @pytest.mark.parametrize( ("args", "expected"), [ ({"action": "question", "text": "Need a choice"}, Decision.ALLOW), - ({"action": "list"}, Decision.DENY), - ({"action": "send", "id": "w", "text": "continue"}, Decision.DENY), + ({"action": "list"}, Decision.ALLOW), + ({"action": "inspect", "id": "w"}, Decision.ALLOW), + ({"action": "send", "id": "w", "text": "continue"}, Decision.ALLOW), ({"action": "stop", "id": "w"}, Decision.DENY), ({"action": "resume", "id": "w"}, Decision.DENY), ({"action": "submit", "id": "w"}, Decision.DENY), ({"action": "integrate"}, Decision.DENY), ({"action": "cleanup"}, Decision.DENY), + ({"action": "recover"}, Decision.DENY), + ({"action": "unknown"}, Decision.DENY), + ({}, Decision.DENY), + ({"action": None}, Decision.DENY), + ({"action": 1}, Decision.DENY), + ({"action": True}, Decision.DENY), + ({"action": ["send"]}, Decision.DENY), + ({"action": {"send": True}}, Decision.DENY), ], ) -def test_strict_readonly_allows_only_workers_question(args, expected): - checker = _checker({"mode": "yolo"}, read_only=True) +def test_readonly_workers_control_plane(args, expected, strict): + checker = _checker({"mode": "yolo" if strict else "readonly"}, read_only=strict) assert checker.check("workers", args).decision == expected +@pytest.mark.parametrize("action", ["list", "question", "send", "inspect"]) +@pytest.mark.parametrize("decision", [Decision.ASK, Decision.DENY]) +@pytest.mark.parametrize("source", ["global", "overlay"]) +def test_readonly_worker_controls_honor_rules_through_descendants(action, decision, source): + rules = {decision: {"workers": [{"pattern": "*"}]}} + parent = _checker( + {"mode": "yolo", "rules": rules if source == "global" else {}}, read_only=True + ).for_agent(AgentOverlay(extra_rules=_ruleset(**rules) if source == "overlay" else _ruleset())) + child = parent.for_child( + AgentOverlay(extra_rules=_ruleset(allow={"workers": [{"pattern": "*"}]})), + session_perms=SessionPermissions([("workers", "*")]), + ).for_child(AgentOverlay(mode="yolo")) + result = child.check("workers", {"action": action, "id": "w", "text": "reply"}) + assert result.decision == decision + assert result.matched_rule == PermissionRule(pattern="*") + + +@pytest.mark.parametrize("action", ["list", "question", "send", "inspect"]) +def test_readonly_worker_controls_honor_denied_tools(action): + checker = _checker({"mode": "readonly"}).for_agent(AgentOverlay(denied_tools=("workers",))) + child = checker.for_child(AgentOverlay(mode="yolo")) + assert child.check("workers", {"action": action}).decision == Decision.DENY + + +def test_readonly_parent_reply_does_not_widen_child_permissions(): + parent = _checker({"mode": "yolo"}, read_only=True) + child = parent.for_child( + AgentOverlay( + mode="yolo", + extra_rules=_ruleset( + allow={"write": [{"pattern": "*"}], "workers": [{"pattern": "*"}]} + ), + ), + session_perms=SessionPermissions([("write", "*"), ("workers", "*")]), + ) + assert parent.check("workers", {"action": "send", "id": "w", "text": "write"}).decision == ( + Decision.ALLOW + ) + assert child.read_only is True + assert child.check("write", {"path": "a.py"}).decision == Decision.DENY + assert child.check("workers", {"action": "integrate"}).decision == Decision.DENY + + def test_read_only_not_widened_by_overlay_allow_rule(): overlay = AgentOverlay(extra_rules=_ruleset(allow={"bash": [{"pattern": "*"}]})) checker = _checker({"mode": "yolo"}, read_only=True).for_agent(overlay) @@ -499,6 +552,41 @@ def test_readonly_with_writable_exceptions_is_not_safe_for_shared_checkout(sourc assert strict.check("write", {"path": "allowed/a"}).decision == Decision.DENY +@pytest.mark.parametrize("source", ["global", "overlay", "grant"]) +@pytest.mark.parametrize("decision", [Decision.ALLOW, Decision.ASK]) +def test_readonly_workers_exceptions_are_not_safe_for_shared_checkout(source, decision): + rules = {decision: {"workers": [{"pattern": "*"}]}} + checker = _checker( + { + "mode": "yolo" if source == "overlay" else "readonly", + "rules": rules if source == "global" else {}, + }, + session_perms=SessionPermissions([("workers", "*")] if source == "grant" else []), + ).for_agent( + AgentOverlay( + mode="readonly", extra_rules=_ruleset(**rules) if source == "overlay" else _ruleset() + ) + ) + assert checker.read_only is False + assert checker.check("workers", {"action": "integrate"}).decision == ( + Decision.ALLOW if source == "grant" else decision + ) + strict = checker.for_child(read_only=True) + assert strict.read_only is True + assert strict.check("workers", {"action": "integrate"}).decision == Decision.DENY + + +def test_child_readonly_capability_uses_base_policy_despite_yolo_overlay(): + parent = _checker({"mode": "readonly"}, session_perms=SessionPermissions([("write", "*")])) + child = parent.for_child(AgentOverlay(mode="yolo")) + assert parent.read_only is False + assert child.read_only is True + assert child.check("write", {"path": "a.py"}).decision == Decision.DENY + assert child.check("workers", {"action": "send", "id": "w", "text": "reply"}).decision == ( + Decision.ALLOW + ) + + @pytest.mark.parametrize("parent_decision", list(Decision)) @pytest.mark.parametrize("overlay_decision", list(Decision)) def test_per_call_overlay_composes_with_full_policy(parent_decision, overlay_decision): diff --git a/tests/test_session_stats.py b/tests/test_session_stats.py index 16e4d2e..17de081 100644 --- a/tests/test_session_stats.py +++ b/tests/test_session_stats.py @@ -14,6 +14,8 @@ @pytest.fixture def store(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + monkeypatch.chdir(tmp_path) return SessionStore() @@ -57,6 +59,7 @@ def test_cost_unknown_model_skipped(store): stats = session_stats(store, s, catalog=sample_catalog()) assert stats.cost_usd == 0.0 assert stats.input_tokens == 5 + assert stats.usage_incomplete def test_openai_style_usage_keys(store): @@ -148,8 +151,10 @@ def test_worker_usage_includes_failed_and_cancelled_dispatches(store, session): def test_worker_usage_incomplete_flag(store, session): - store.append_event(session, "worker_usage", {"input_tokens": 1, "output_tokens": 1}) - assert session_stats(store, session).usage_incomplete is False + store.append_event( + session, "worker_usage", {"input_tokens": 1, "output_tokens": 1, "cost_usd": 0} + ) + assert session_stats(store, session, catalog=sample_catalog()).usage_incomplete is False store.append_event( session, "worker_usage", @@ -158,3 +163,111 @@ def test_worker_usage_incomplete_flag(store, session): stats = session_stats(store, session) assert stats.usage_incomplete is True assert stats.input_tokens == 1_001_000 + 2 + + +@pytest.mark.parametrize("kind", ["message", "pierre", "worker_usage"]) +@pytest.mark.parametrize("nested", [False, True]) +def test_incomplete_zero_usage_propagates(store, kind, nested): + session = store.create("incomplete", cwd="/tmp", model=GPT5_MINI) + usage = {"input_tokens": 0, "output_tokens": 0, "cost_usd": 0, "incomplete": True} + if kind == "message": + store.append_message(session, {"role": "assistant", "content": "done"}, usage=usage) + else: + data = {"usage": usage} if nested else usage + store.append_event(session, kind, data) + stats = session_stats(store, session, catalog=sample_catalog()) + assert stats.cost_usd == 0 + assert stats.usage_incomplete + + +@pytest.mark.parametrize("model", ["missing/model", "openai/gpt-5-", None]) +@pytest.mark.parametrize("usage", [None, {"input_tokens": 0}, {"cost_usd": 0}]) +def test_missing_usage_and_pricing_are_not_known_zero(store, model, usage): + session = store.create("zero", cwd="/tmp", model=model) + store.append_message(session, {"role": "assistant", "content": "done"}, usage=usage) + stats = session_stats(store, session, catalog=sample_catalog()) + assert stats.cost_usd == 0 + assert stats.usage_incomplete is (usage != {"cost_usd": 0}) + + +@pytest.mark.parametrize("incomplete", [False, True]) +async def test_worker_live_usage_checkpoint_and_root_restore_count_once( + tool_ctx, store, tmp_path, incomplete +): + from lecode.agent.runner import LlmResponse + from lecode.extras.workers import Worker, WorkerManager + + root = store.create("root", cwd=tmp_path) + store.append_message( + root, + {"role": "assistant", "content": "root"}, + usage={"input_tokens": 10, "output_tokens": 1, "cost_usd": 0.5}, + ) + child = store.create("child", cwd=tmp_path) + worker = Worker("child", None, 1, "explore", "delegated", "completed", child, tmp_path) + worker.dispatch_id = "dispatch" + manager = WorkerManager( + tool_ctx.config, cwd=tmp_path, root_ctx=tool_ctx, store=store, session=root + ) + manager._record(worker) + stale_snapshot = store.load_events(root, "worker")[-1] + await manager._event(worker, LlmResponse("model", 1, 5, 2, 0, usage_incomplete=incomplete)) + assert worker.usage_incomplete is incomplete + assert worker.usage_totals.cost_usd == 0 + manager._record_usage(worker) + await manager._event(worker, LlmResponse("model", 2, 7, 3, 0.25)) + assert worker.usage_incomplete is incomplete + checkpoint = store.load_events(child, "worker_usage_checkpoint")[-1] + assert checkpoint["usage_incomplete"] is incomplete + # The child checkpoint is durable even if the last root snapshot was lost. + store.append_event(root, "worker", stale_snapshot) + await manager.shutdown() + + for _ in range(2): + restored = WorkerManager( + tool_ctx.config, cwd=tmp_path, root_ctx=tool_ctx, store=store, session=root + ) + try: + loaded = restored.load()[0] + assert loaded.usage_totals == worker.usage_totals + assert loaded.usage_incomplete is incomplete + restored.load() + stats = session_stats(store, root) + assert (stats.input_tokens, stats.output_tokens, stats.context_tokens) == (22, 6, 10) + assert stats.cost_usd == 0.75 + assert stats.usage_incomplete is incomplete + assert len(store.load_events(root, "worker_usage")) == 2 + finally: + await restored.shutdown() + + +@pytest.mark.parametrize("flag", [{}, {"incomplete": True}, {"usage_incomplete": True}]) +async def test_worker_checkpoint_completeness_loads_old_and_new_json( + tool_ctx, store, tmp_path, flag +): + from lecode.extras.workers import Worker, WorkerManager + + root = store.create("root", cwd=tmp_path) + child = store.create("child", cwd=tmp_path) + worker = Worker("child", None, 1, "explore", "delegated", "completed", child, tmp_path) + manager = WorkerManager( + tool_ctx.config, cwd=tmp_path, root_ctx=tool_ctx, store=store, session=root + ) + manager._record(worker) + snapshot = store.load_events(root, "worker")[-1] + snapshot["usage_totals"].pop("usage_incomplete") + store.append_event(root, "worker", snapshot) + store.append_event( + child, + "worker_usage_checkpoint", + {"dispatch_id": None, "input_tokens": 2, "cost_usd": 0, **flag}, + ) + try: + loaded = manager.load()[0] + assert loaded.usage_incomplete is bool(flag) + stats = session_stats(store, root) + assert stats.usage_incomplete is bool(flag) + assert stats.input_tokens == 2 + assert stats.cost_usd == 0 + finally: + await manager.shutdown() diff --git a/tests/test_slash_features.py b/tests/test_slash_features.py index c4ca124..e7b9936 100644 --- a/tests/test_slash_features.py +++ b/tests/test_slash_features.py @@ -3,9 +3,18 @@ from __future__ import annotations +import shlex +import shutil + +import pytest from tests.test_tui_app import make_app +from tests.test_worker_controls import commit, review_heads from tests.test_worktree import make_repo +from lecode.config.models import PermissionRule, PermissionRuleSet +from lecode.context.agents import AgentDefinition, AgentRegistry +from lecode.permission.checker import AgentOverlay, Deny + async def test_agent_focus_command_targets_persistent_worker(tmp_path, monkeypatch): app, _, out = make_app(tmp_path, monkeypatch, [{"text": "done"}]) @@ -53,6 +62,172 @@ async def test_agent_submit_wakes_root_once(tmp_path, monkeypatch): await manager.shutdown() +@pytest.mark.parametrize( + "action", ["send hello", "stop", "resume hello", "submit", "cleanup", "recover"] +) +async def test_agent_mutations_use_workers_permission_gate(tmp_path, monkeypatch, action): + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "done"}]) + manager = app.worker_manager + worker = await manager.start(app.runtime.ctx, agent="explore", prompt="Inspect", origin="human") + await manager.wait(worker.id) + app.runtime.ctx.permission_checker = app.runtime.ctx.permission_checker.for_agent( + AgentOverlay(denied_tools=("workers",)) + ) + try: + await app.handle_command(f"/agent 1 {action}") + assert "denied" in out.getvalue() + assert worker.state == "completed" + assert manager.pending(worker.id) == [] + finally: + await manager.shutdown() + + +async def test_agent_review_integrate_cleanup_and_retained_transcript(tmp_path, monkeypatch): + await make_repo(tmp_path) + (tmp_path / ".git/info/exclude").write_text("/cfg/\n/global-skills/\n") + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "Implemented"}]) + app.runtime.ctx.extras["agents"] = AgentRegistry( + {"writer": AgentDefinition(name="writer", description="Write", body="", mode="subagent")} + ) + manager = app.worker_manager + worker = await manager.start( + app.runtime.ctx, agent="writer", prompt="Implement", origin="human" + ) + await manager.wait(worker.id) + head = commit(worker.cwd) + reviewed = await review_heads(app.runtime, worker) + manager.confirm = lambda _: True + try: + await app.handle_command("/agent 1 inspect") + assert head in out.getvalue() and "+worker change" in out.getvalue() + await app.handle_command(f"/agent 1 integrate {head}") + assert "integrate WORKER_HASH PARENT_HASH" in out.getvalue() + assert not (tmp_path / "change.txt").exists() + await app.handle_command(f"/agent 1 integrate {head} {reviewed['reviewed_parent_head']}") + assert (tmp_path / "change.txt").exists(), out.getvalue() + assert (tmp_path / "change.txt").read_text() == "worker change\n" + await app.handle_command("/agent 1 cleanup") + assert not worker.cwd.exists() + await app.handle_command("/agent 1") + assert app.detail_run_id == worker.id + assert "Implemented" in str(app._roster_text()) + finally: + await manager.shutdown() + + +async def test_agent_recover_warns_and_requires_human_confirmation(tmp_path, monkeypatch): + await make_repo(tmp_path) + (tmp_path / ".git/info/exclude").write_text("/cfg/\n/global-skills/\n") + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "Implemented"}]) + app.runtime.ctx.extras["agents"] = AgentRegistry( + {"writer": AgentDefinition(name="writer", description="Write", body="", mode="subagent")} + ) + manager = app.worker_manager + worker = await manager.start( + app.runtime.ctx, agent="writer", prompt="Implement", origin="human" + ) + await manager.wait(worker.id) + shutil.rmtree(worker.cwd) + questions = [] + + def confirm(question): + questions.append(question) + return len(questions) > 1 + + manager.confirm = confirm + try: + await app.handle_command("/agent 1 recover") + assert not worker.cwd.exists() and "declined" in out.getvalue() + await app.handle_command("/agent 1 recover") + assert worker.cwd.is_dir() + assert worker.id in questions[0] and "unrecoverable" in questions[0] + finally: + await manager.shutdown() + + +@pytest.mark.parametrize( + "restriction", + [ + "ancestor-deny", + "ancestor-ask", + "root-deny-workers", + "root-ask-workers", + "root-deny-bash", + "root-ask-bash", + "root-mode", + "both", + "allowed", + ], +) +async def test_nested_agent_integration_keeps_live_ancestor_and_root_permissions( + tmp_path, monkeypatch, restriction +): + await make_repo(tmp_path) + (tmp_path / ".git/info/exclude").write_text("/cfg/\n/global-skills/\n") + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "done"}] * 2) + app.runtime.ctx.extras["agents"] = AgentRegistry( + {"writer": AgentDefinition(name="writer", description="Write", body="", mode="subagent")} + ) + base = app.runtime.ctx.permission_checker + approvals = [] + + async def reject(name, args, reason, **kwargs): + approvals.append((name, reason)) + return Deny() + + app.runtime.ctx.approval_callback = reject + if restriction in {"ancestor-deny", "both"}: + app.runtime.ctx.permission_checker = base.for_agent(AgentOverlay(denied_tools=("bash",))) + elif restriction == "ancestor-ask": + app.runtime.ctx.permission_checker = base.for_agent( + AgentOverlay(extra_rules=PermissionRuleSet(ask={"bash": [PermissionRule(pattern="*")]})) + ) + manager = app.worker_manager + parent = await manager.start(app.runtime.ctx, agent="writer", prompt="Parent", origin="human") + await manager.wait(parent.id) + nested = manager._runtime(parent) + child = await manager.start(nested.ctx, agent="writer", prompt="Child") + await manager.wait(child.id) + head = commit(child.cwd) + reviewed = await review_heads(nested, child) + assert parent.depth == 1 and child.depth == 2 + # Persisted child grants cannot override an ancestor's Deny or Ask. + for worker in (parent, child): + manager._runtime(worker).ctx.session_perms.grant("bash", "*") + manager.store.grant_permission(worker.session, "bash", "*") + marker = tmp_path.parent / f"{tmp_path.name}-validation" + app.config.worktree.validation = [f"printf validated > {shlex.quote(str(marker))}"] + # Simulate a root agent switch while retaining the real cached supervisor. + app.runtime.ctx.permission_checker = base.for_agent(AgentOverlay(denied_tools=("write",))) + if restriction.startswith("root-deny-") or restriction == "both": + tool = "workers" if restriction == "both" else restriction.removeprefix("root-deny-") + app.runtime.ctx.permission_checker = base.for_agent(AgentOverlay(denied_tools=(tool,))) + elif restriction.startswith("root-ask-"): + tool = restriction.removeprefix("root-ask-") + app.runtime.ctx.permission_checker = base.for_agent( + AgentOverlay(extra_rules=PermissionRuleSet(ask={tool: [PermissionRule(pattern="*")]})) + ) + elif restriction == "root-mode": + app.set_permission_mode("readonly") + try: + await app.handle_command( + f"/agent {child.id} integrate {head} {reviewed['reviewed_parent_head']}" + ) + if restriction == "allowed": + assert marker.read_text() == "validated", out.getvalue() + assert (parent.cwd / "change.txt").read_text() == "worker change\n" + else: + assert "denied" in out.getvalue(), out.getvalue() + assert not marker.exists() + assert not (parent.cwd / "change.txt").exists() + assert bool(approvals) is ("ask" in restriction) + if approvals and approvals[0][0] == "bash": + assert child.id in approvals[0][1] and str(child.cwd) in approvals[0][1] + assert not (tmp_path / "change.txt").exists() + finally: + await manager.shutdown() + + # -- /init ------------------------------------------------------------------------- diff --git a/tests/test_subagents.py b/tests/test_subagents.py index cfb797f..5c0d723 100644 --- a/tests/test_subagents.py +++ b/tests/test_subagents.py @@ -174,6 +174,7 @@ async def test_completed_run_persists_activity_trail(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) store = SessionStore(config_dir=tmp_path / "cfg") + monkeypatch.chdir(tmp_path) session = store.create("agent-run", tmp_path) runtime = build_runtime(Config(), tmp_path, session=session, store=store) runtime.ctx.extras["provider"] = provider @@ -208,6 +209,7 @@ async def test_cancelled_run_persists_cancelled_status(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) store = SessionStore(config_dir=tmp_path / "cfg") + monkeypatch.chdir(tmp_path) session = store.create("agent-run", tmp_path) runtime = build_runtime(Config(), tmp_path, session=session, store=store) runtime.ctx.extras["provider"] = provider @@ -238,6 +240,7 @@ async def test_failed_run_persists_error_status(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) store = SessionStore(config_dir=tmp_path / "cfg") + monkeypatch.chdir(tmp_path) session = store.create("agent-run", tmp_path) runtime = build_runtime(Config(), tmp_path, session=session, store=store) runtime.ctx.extras["provider"] = NeverProvider() @@ -313,6 +316,7 @@ async def test_primary_agent_not_invocable(tmp_path, monkeypatch): async def test_missing_provider(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) runtime = build_runtime(Config(), tmp_path) # no provider seam installed with pytest.raises(SubagentError, match="no provider available"): await run_subagent( @@ -567,6 +571,7 @@ async def test_task_result_event_carries_run_id_metadata(tmp_path, monkeypatch): monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) store = SessionStore(config_dir=tmp_path / "cfg") + monkeypatch.chdir(tmp_path) session = store.create("agent-run", tmp_path) runtime = build_runtime(Config(), tmp_path, session=session, store=store) runtime.ctx.extras["provider"] = provider @@ -604,22 +609,36 @@ async def test_roster_panel_visible_while_child_runs(tmp_path, monkeypatch): app._runner.provider = provider app._runtime.ctx.extras["provider"] = provider await app._submit("go") - await wait_for(lambda: provider.child_started.is_set()) - - assert app._roster_visible() - rows = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) - assert "Scan repo" in rows - - run_id = app.roster.runs()[0].run_id - assert app.open_agent_run(run_id) - detail = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) - assert "Scan repo" in detail - - app._turn_task.cancel() - await app._turn_task - assert app.roster.runs()[0].status == "cancelled" - app.close_agent_run() - assert not app._roster_visible() + try: + await wait_for(lambda: provider.child_started.is_set()) + + assert app._roster_visible() + rows = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) + assert "Scan repo" in rows + + run = app.roster.runs()[0] + assert app.open_agent_run(run.run_id) + detail = "".join(fragment[1] for fragment in to_formatted_text(app._roster_text())) + assert "Scan repo" in detail + + app._turn_task.cancel() + await app._turn_task + worker = app._worker_manager.get(run.run_id) + assert worker.state == "running" + assert worker.is_active + assert run.status == "running" + app.close_agent_run() + assert app._roster_visible() + + await app._worker_manager.stop(worker.id) + assert worker.state == "stopped" + assert not worker.is_active + assert run.status == "stopped" + assert not app._roster_visible() + finally: + app._turn_task.cancel() + await asyncio.gather(app._turn_task, return_exceptions=True) + await app._worker_manager.shutdown() async def test_runs_command_lists_and_opens_detail(tmp_path, monkeypatch): diff --git a/tests/test_tool_bash.py b/tests/test_tool_bash.py index fe7c3af..8cf20a0 100644 --- a/tests/test_tool_bash.py +++ b/tests/test_tool_bash.py @@ -17,12 +17,17 @@ async def test_echo(tool_ctx): result = await bash.make_tool().run({"command": "echo hello"}, tool_ctx) assert not result.is_error assert "hello" in result.content + proc = result.metadata["proc_result"] + assert proc.exit_code == 0 + assert proc.stdout == "hello\n" + assert not proc.timed_out async def test_nonzero_exit_reported(tool_ctx): result = await bash.make_tool().run({"command": "exit 7"}, tool_ctx) assert result.is_error assert "exit code 7" in result.content + assert result.metadata["proc_result"].exit_code == 7 async def test_stderr_merged(tool_ctx): @@ -34,6 +39,7 @@ async def test_timeout(tool_ctx): result = await bash.make_tool().run({"command": "sleep 30", "timeout": 0.3}, tool_ctx) assert result.is_error assert "timed out" in result.content + assert result.metadata["proc_result"].timed_out async def test_idle_timeout(tool_ctx): @@ -43,6 +49,7 @@ async def test_idle_timeout(tool_ctx): assert result.is_error assert "no output" in result.content assert "start" in result.content + assert result.metadata["proc_result"].timed_out async def test_truncation_and_overflow_file(tool_ctx, tmp_path): @@ -54,6 +61,7 @@ async def test_truncation_and_overflow_file(tool_ctx, tmp_path): assert len(files) == 1 assert files[0].read_text().count("repeated-output-line") > 1000 assert len(result.content) < MAX_OUTPUT_BYTES + 500 + assert result.metadata["proc_result"].truncated async def test_rtk_rewrite_applied(tool_ctx, tmp_path, monkeypatch): diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index cd1cc52..d8c309a 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -110,6 +110,7 @@ def test_layout_is_chatbox_above_statusline(tmp_path, monkeypatch): assert isinstance(children[3], ConditionalContainer) # picker panel sizes to its rows assert isinstance(children[4], Window) and children[4].height == 3 assert app._live_buffer is not None + assert pt_app.full_screen is False def test_roster_panel_renders_as_prompt_toolkit_text(tmp_path, monkeypatch): @@ -180,6 +181,243 @@ async def test_worker_focus_preserves_drafts_and_routes_composer(tmp_path, monke await manager.shutdown() +@pytest.mark.parametrize( + "command", + ["blocked follow-up", "/agent 1 send blocked follow-up", "/agent 1 resume blocked follow-up"], +) +async def test_focused_prompt_hook_denial_prevents_delivery(tmp_path, monkeypatch, command): + config = Config() + config.hooks = {"UserPromptSubmit": ['echo \'{"verdict":"deny","reason":"closed"}\'']} + app, _, out = make_app(tmp_path, monkeypatch, [{"text": "initial"}], config=config) + manager = app.worker_manager + worker = await manager.start( + app.runtime.ctx, agent="explore", prompt="initial", origin="human", background=True + ) + await manager.wait(worker.id) + before = app.store.load_for_model(worker.session) + assert app.focus_worker(worker.id) + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + inp.send_text(command + "\r") + await wait_for(lambda: "prompt blocked by hook: closed" in out.getvalue()) + assert app.store.load_for_model(worker.session) == before + assert manager.pending(worker.id) == [] + assert not any( + event.get("text") in {command, "blocked follow-up"} + for event in app.store.load_events(app.session, "input") + ) + assert not any( + note.get("content") == "blocked follow-up" + for note in app.store.load_events(app.session, "worker_notification") + ) + inp.send_text("/quit\r") + assert await task == 0 + + +@pytest.mark.parametrize("route", ["main", "focused", "mention", "expanded", "send", "resume"]) +async def test_human_dispatch_hooks_once_and_preserves_prompt(tmp_path, monkeypatch, route): + import json + + config = Config() + config.hooks = {"UserPromptSubmit": ["cat >> prompts.jsonl"]} + app, provider, _ = make_app(tmp_path, monkeypatch, [{"text": "answer"}] * 2, config=config) + if route in {"focused", "send", "resume"}: + worker = await app.worker_manager.start( + app.runtime.ctx, agent="explore", prompt="initial", origin="human", background=True + ) + await app.worker_manager.wait(worker.id) + app.focus_worker(worker.id) + prompt = "@explore inspect" if route == "mention" else "expanded instructions" + before = len(provider.requests) + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + if route == "expanded": + app.submit_prompt(prompt, echo="/skill") + elif route in {"send", "resume"}: + inp.send_text(f"/agent 1 {route} {prompt}\r") + else: + inp.send_text(prompt + "\r") + await wait_for(lambda: len(provider.requests) > before) + sent = [m["content"] for m in provider.requests[-1]["messages"] if m["role"] == "user"] + assert ("inspect" if route == "mention" else prompt) in sent + hooks = [json.loads(line) for line in (tmp_path / "prompts.jsonl").read_text().splitlines()] + assert [hook["prompt"] for hook in hooks] == ["inspect" if route == "mention" else prompt] + inp.send_text("/quit\r") + assert await task == 0 + + +@pytest.mark.parametrize("deny", [False, True]) +async def test_retry_hook_runs_before_undo_and_only_once(tmp_path, monkeypatch, deny): + import json + + config = Config() + config.hooks = { + "UserPromptSubmit": [ + "cat >> prompts.jsonl; echo >> prompts.jsonl; " + 'if test -e deny; then echo \'{"verdict":"deny"}\'; fi' + ] + } + app, _, out = make_app( + tmp_path, monkeypatch, [{"text": "first"}, {"text": "retry"}], config=config + ) + await app._submit("original prompt") + await app._turn_task + before = app.store.load_for_model(app.session) + if deny: + (tmp_path / "deny").touch() + await app.handle_command("/retry") + if deny: + await wait_for(lambda: "prompt blocked by hook" in out.getvalue()) + assert app.store.load_for_model(app.session) == before + else: + await wait_for( + lambda: "retry" in [m.get("content") for m in app.store.load_for_model(app.session)] + ) + prompts = [ + json.loads(line)["prompt"] for line in (tmp_path / "prompts.jsonl").read_text().splitlines() + ] + assert prompts == ["original prompt", "original prompt"] + + +@pytest.mark.parametrize("command", ["go", ".reviewer go", "/review note.py"]) +async def test_notes_hook_checks_composed_prompt_before_persistence(tmp_path, monkeypatch, command): + import json + + config = Config() + config.hooks = { + "UserPromptSubmit": [ + 'payload=$(cat); echo "$payload" >> prompts.jsonl; ' + 'case "$payload" in *BLOCKED*) echo \'{"verdict":"deny"}\' ;; esac' + ] + } + app, provider, out = make_app(tmp_path, monkeypatch, [{"text": "unused"}], config=config) + (tmp_path / "note.py").write_text("print('safe')\n") + await app._submit("/btw BLOCKED") + before = app.store.read_records(app.session) + await app._submit(command) + await wait_for(lambda: "prompt blocked by hook" in out.getvalue()) + assert not provider.requests + assert app.store.read_records(app.session) == before + assert app._pending_notes == ["BLOCKED"] + prompts = [ + json.loads(line)["prompt"] for line in (tmp_path / "prompts.jsonl").read_text().splitlines() + ] + assert len(prompts) == 1 + if command.startswith("/review"): + assert prompts[0].startswith("BLOCKED\n\n") + assert "meticulous code reviewer" in prompts[0] and "print('safe')" in prompts[0] + else: + assert prompts == ["BLOCKED\n\ngo"] + + +@pytest.mark.parametrize("route", ["chain", "loop"]) +@pytest.mark.parametrize("deny_after", [0, 1]) +async def test_generated_prompts_are_guarded_at_each_submission( + tmp_path, monkeypatch, route, deny_after +): + import json + + config = Config() + config.hooks = { + "UserPromptSubmit": [ + 'payload=$(cat); echo "$payload" >> prompts.jsonl; ' + 'case "$payload" in *BLOCKED*) echo \'{"verdict":"deny"}\' ;; esac' + ] + } + app, provider, out = make_app(tmp_path, monkeypatch, [{"text": "BLOCKED"}] * 4, config=config) + topic = "BLOCKED" if deny_after == 0 else "safe" + plan = tmp_path / "plan.md" + plan.write_text(f"- [ ] {topic}\n") + if route == "loop" and deny_after: + original = provider.stream_chat + + def stream_chat(*args, **kwargs): + plan.write_text("- [ ] BLOCKED\n") + return original(*args, **kwargs) + + monkeypatch.setattr(provider, "stream_chat", stream_chat) + before = app.store.read_records(app.session) + await app._submit(f"/chain {topic}" if route == "chain" else "/loop plan.md 3") + await (app._turn_task if route == "chain" else app._loop_task) + assert "prompt blocked by hook" in out.getvalue() + assert len(provider.requests) == deny_after + prompts = [ + json.loads(line)["prompt"] for line in (tmp_path / "prompts.jsonl").read_text().splitlines() + ] + assert len(prompts) == deny_after + 1 + for prompt, request in zip(prompts[:-1], provider.requests, strict=True): + assert prompt == request["messages"][-1]["content"] + assert "BLOCKED" in prompts[-1] + if deny_after == 0: + assert app.store.read_records(app.session) == before + assert not any( + m.get("role") == "user" and "BLOCKED" in m.get("content", "") + for m in app.store.load_for_model(app.session) + ) + + +@pytest.mark.parametrize("state", ["running", "stopped", "completed"]) +async def test_worker_focus_and_detail_follow_same_target_and_escape_parent( + tmp_path, monkeypatch, state +): + from dataclasses import replace + + app, _, _ = make_app(tmp_path, monkeypatch, [{"text": "answer A"}, {"text": "answer B"}]) + manager = app.worker_manager + parent = await manager.start(app.runtime.ctx, agent="explore", prompt="A", origin="human") + await manager.wait(parent.id) + child_ctx = replace(app.runtime.ctx, extras={**app.runtime.ctx.extras, "worker_id": parent.id}) + if state != "completed": + child_ctx.extras["provider"] = BlockingProvider() + child = await manager.start(child_ctx, agent="explore", prompt="B", origin="human") + if state == "completed": + await manager.wait(child.id) + else: + await wait_for(lambda: bool(child_ctx.extras["provider"].requests)) + if state == "stopped": + await manager.stop(child.id) + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + try: + await wait_for(lambda: app._input_area is not None) + app._input_area.buffer.text = "main draft" + await app.handle_command("/agent 1") + inp.send_text("/agent 2 focus\r") + await wait_for(lambda: app._focused_worker_id == child.id) + assert app.detail_run_id == app._focused_worker_id == child.id + if state == "completed": + assert "answer B" in str(app._roster_text()) + assert "answer A" not in str(app._roster_text()) + inp.send_text("child draft") + await wait_for(lambda: app._input_area.text == "child draft") + inp.send_text("\x1b") + await wait_for(lambda: app._focused_worker_id == parent.id) + assert app.detail_run_id == parent.id + assert "answer A" in str(app._roster_text()) + inp.send_text("parent draft") + await wait_for(lambda: app._input_area.text == "parent draft") + inp.send_text("\x1b") + await wait_for(lambda: app._focused_worker_id is None) + assert app.detail_run_id is None + assert app._input_area.text == "main draft" + await app.handle_command("/agent 2 focus") + assert app._input_area.text == "child draft" + await app.handle_command("/agent 1") + assert app.detail_run_id == app._focused_worker_id == parent.id + assert app._input_area.text == "parent draft" + inp.send_text("\x15") # clear the draft before entering a command + await wait_for(lambda: app._input_area.text == "") + inp.send_text("/agent root\r") + await wait_for(lambda: app._focused_worker_id is None, timeout=1) + assert app.detail_run_id is None + assert app._input_area.text == "main draft" + finally: + app.request_quit() + await task + + async def test_worker_detail_renders_persisted_child_transcript(tmp_path, monkeypatch): """Worker detail reads the child session, not just the bounded live trail.""" (tmp_path / "note.txt").write_text("persisted tool result", encoding="utf-8") @@ -672,6 +910,71 @@ async def test_ctrl_c_cancels_running_turn(tmp_path, monkeypatch): assert app.cancel_turn() is False # nothing running now +async def test_ctrl_c_denies_only_pending_approval(tmp_path, monkeypatch): + from lecode.permission import Deny + + app, provider, _ = make_blocking_app(tmp_path, monkeypatch) + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + inp.send_text("root work\r") + await wait_for(lambda: len(provider.requests) == 1) + callback = app.runtime.ctx.approval_callback + first = asyncio.create_task(callback("bash", {"command": "ls"}, "")) + second = asyncio.create_task(callback("bash", {"command": "pwd"}, "")) + await wait_for(lambda: app._approval.is_pending) + inp.send_text("\x03") + assert await first == Deny() + assert not second.done() + assert not app._turn_task.done() + inp.send_text("n") + assert await second == Deny() + provider.blocked = False + provider.release.set() + await wait_for(lambda: not app._turn_running()) + inp.send_text("/quit\r") + assert await task == 0 + + +@pytest.mark.parametrize("direct", [False, True]) +async def test_ctrl_c_stops_focused_worker_not_root_or_waiter(tmp_path, monkeypatch, direct): + app, provider, out = make_blocking_app(tmp_path, monkeypatch) + app.runtime.ctx.extras["provider"] = provider + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + inp.send_text("@explore inspect\r" if direct else "root work\r") + await wait_for(lambda: len(provider.requests) == 1) + if direct: + worker = app.worker_manager.list()[0] + else: + inp.send_text("root queued\r") + await wait_for(lambda: app.queued_prompts() == ([], ["root queued"])) + worker = await app.worker_manager.start( + app.runtime.ctx, + agent="explore", + prompt="child work", + origin="human", + background=True, + ) + await wait_for(lambda: len(provider.requests) == 2) + assert app.focus_worker(worker.id) + inp.send_text("child draft") + await wait_for(lambda: app._input_area.text == "child draft") + inp.send_text("\x03") + await wait_for(lambda: worker.state == "stopped") + assert app._input_area.text == "child draft" + if not direct: + assert not app._turn_task.done() + assert app.queued_prompts() == ([], ["root queued"]) + assert "turn cancelled" not in out.getvalue() + provider.blocked = False + provider.release.set() + await wait_for(lambda: not app._turn_running()) + inp.send_text("\x15/quit\r") + assert await task == 0 + + async def test_provider_error_renders_and_recovers(tmp_path, monkeypatch): from lecode.providers.openai_compat import ProviderError @@ -772,6 +1075,113 @@ async def test_statusline_totals_update_after_turn(tmp_path, monkeypatch): assert app._status.context_used == 10 +async def test_live_total_counts_root_and_nested_workers_once(tmp_path, monkeypatch): + from dataclasses import replace + + release = asyncio.Event() + + class PausedProvider(FakeProvider): + async def _stream(self, entry): + if entry.get("pause"): + await release.wait() + async for event in super()._stream(entry): + yield event + + app, _, _ = make_app(tmp_path, monkeypatch, []) + provider = PausedProvider( + [ + { + "tool_calls": [{"name": "read", "arguments": '{"path":"note.txt"}'}], + "usage": {"input_tokens": 100, "output_tokens": 10, "cost_usd": 0.1}, + }, + { + "pause": True, + "text": "root done", + "usage": {"input_tokens": 200, "output_tokens": 20, "cost_usd": 0.2}, + }, + {"text": "child", "usage": {"input_tokens": 300, "output_tokens": 30, "cost_usd": 0.3}}, + { + "text": "grandchild", + "usage": {"input_tokens": 400, "output_tokens": 40, "cost_usd": 0.4}, + }, + ] + ) + app.runner.provider = provider + app.runtime.ctx.extras["provider"] = provider + (tmp_path / "note.txt").write_text("note") + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + inp.send_text("root work\r") + await wait_for(lambda: len(provider.requests) == 2) + assert app.status.cost_usd == pytest.approx(0.1) + worker = await app.worker_manager.start( + app.runtime.ctx, agent="explore", prompt="child", origin="human", background=True + ) + await app.worker_manager.wait(worker.id) + child_ctx = replace( + app.runtime.ctx, extras={**app.runtime.ctx.extras, "worker_id": worker.id} + ) + child = await app.worker_manager.start( + child_ctx, agent="explore", prompt="grandchild", origin="human", background=True + ) + await app.worker_manager.wait(child.id) + assert app.status.cost_usd == pytest.approx(0.8) + release.set() + await wait_for(lambda: not app._turn_running()) + assert app.status.cost_usd == pytest.approx(1.0) + assert app.current_session_usage()[:3] == (1000, 100, pytest.approx(1.0)) + app.reload_history() + assert app.status.cost_usd == pytest.approx(1.0) + inp.send_text("/quit\r") + assert await task == 0 + + +async def test_selected_worker_context_and_incomplete_cost_survive_idle(tmp_path, monkeypatch): + config = Config() + config.agent.subagent_model = "anthropic/claude-sonnet-4" + app, _, out = make_app( + tmp_path, + monkeypatch, + [ + {"text": "main", "usage": {"input_tokens": 100, "output_tokens": 10, "cost_usd": 0.1}}, + { + "text": "child", + "usage": {"input_tokens": 2000, "output_tokens": 20, "cost_usd": 0.2}, + }, + {"text": "missing usage"}, + ], + config=config, + ) + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app._input_area is not None) + inp.send_text("main work\r") + await wait_for(lambda: "answer:" in out.getvalue()) + main_window = app.status.context_window + worker = await app.worker_manager.start( + app.runtime.ctx, agent="explore", prompt="child", origin="human", background=True + ) + await app.worker_manager.wait(worker.id) + assert app.focus_worker(worker.id) + rendered = "".join(part[1] for part in to_formatted_text(app._toolbar())) + assert "2.0k/200.0k" in rendered + assert "anthropic/claude-sonnet-4" in rendered + assert "total cost: $0.3000" in rendered + inp.send_text("follow up\r") + await wait_for(lambda: "sent to @explore" in out.getvalue()) + await app.worker_manager.wait(worker.id) + assert not app.roster.has_running() + assert app.focus_worker(None) + assert app.status.context_window == main_window + assert app.status.context_used == 100 + rendered = "".join(part[1] for part in to_formatted_text(app._toolbar())) + assert "total cost: $0.3000 incomplete" in rendered + assert out.getvalue().count(f"worker {worker.id[:8]}") == 2 + inp.send_text("/quit\r") + assert await task == 0 + + # -- end-to-end pipe-input smokes --------------------------------------------- @@ -895,6 +1305,7 @@ def attach_worktree(self, manager, info, original_cwd) -> None: def cli_env(tmp_path, monkeypatch): """Isolated cwd + config dir; deps check, provider, prompt and TuiApp faked.""" monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) monkeypatch.chdir(tmp_path) monkeypatch.setattr("lecode.cli.check_dependencies", lambda: None) monkeypatch.setattr("lecode.cli.build_provider", lambda config, api_key=None: object()) diff --git a/tests/test_tui_permission.py b/tests/test_tui_permission.py index 0753f86..bd8cd56 100644 --- a/tests/test_tui_permission.py +++ b/tests/test_tui_permission.py @@ -365,24 +365,32 @@ async def test_worker_approval_is_attributed_at_fifo_head(tmp_path, monkeypatch) async def test_dirty_worker_confirmation_is_one_shot(tmp_path, monkeypatch): - from types import SimpleNamespace - app, out = make_app(tmp_path, monkeypatch, []) - task = asyncio.ensure_future( - app._confirm_worker_worktree( - "Uncommitted changes will not enter the worker's committed-HEAD worktree. Continue?", - worker=SimpleNamespace(id="worker-1234", agent="build"), - worktree=tmp_path / "worker-tree", - ) + question = ( + f"@build worker worker-1234 in {tmp_path / 'worker-tree'}: " + "Uncommitted changes will not enter the committed-HEAD worktree. Continue?" ) - await wait_for(lambda: app._approval.pending is not None) - pending = app._approval.pending - assert pending is not None and pending.allow_always is False - assert "@build worker worker-1" in out.getvalue() - assert "worker-tree" in out.getvalue() - assert "(y)es (n)o" in out.getvalue() - app._resolve_approval(AllowAlways(pattern="*")) - assert await task is False + with create_pipe_input() as inp: + task = asyncio.create_task(app.run(input=inp, output=DummyOutput())) + await wait_for(lambda: app.worker_manager.confirm is not None) + first = asyncio.create_task(app.worker_manager.confirm(question)) + second = asyncio.create_task(app.worker_manager.confirm("second worker question?")) + await wait_for(lambda: "(y)es (n)o" in out.getvalue()) + assert question in " ".join(out.getvalue().split()) + assert "second worker question?" not in out.getvalue() + inp.send_text("a") + await asyncio.sleep(0.05) + assert not first.done() and not second.done() + inp.send_text("y") + assert await first is True + await wait_for(lambda: "second worker question?" in out.getvalue()) + assert not second.done() + inp.send_text("n") + assert await second is False + assert app.status.state is StatusLineState.IDLE + assert app.store.load_grants(app.session) == [] + inp.send_text("/quit\r") + assert await task == 0 async def test_pipe_approval_y_runs_asked_tool(tmp_path, monkeypatch): diff --git a/tests/test_tui_statusline.py b/tests/test_tui_statusline.py index 7907b67..da58f23 100644 --- a/tests/test_tui_statusline.py +++ b/tests/test_tui_statusline.py @@ -58,19 +58,30 @@ def test_line1_omits_missing_git_fields(state, theme): def test_line2_model_cost_context(state, theme): line2 = render_statusline(state, theme, width=200).plain.splitlines()[1] - assert line2 == "model: openai/gpt-5-mini · cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42%" + assert line2 == "total cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42% · model: openai/gpt-5-mini" def test_line2_shows_reasoning_override_after_model(state, theme): state.reasoning = "High" line2 = render_statusline(state, theme, width=200).plain.splitlines()[1] - assert line2 == "model: openai/gpt-5-mini · High · cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42%" + assert ( + line2 + == "total cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42% · model: openai/gpt-5-mini · High" + ) def test_line2_omits_reasoning_at_baseline(state, theme): line2 = render_statusline(state, theme, width=200).plain.splitlines()[1] assert "High" not in line2 - assert line2 == "model: openai/gpt-5-mini · cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42%" + assert line2 == "total cost: $0.0123 · ctx: ▓▓▓░░ 84.0k/200.0k 42% · model: openai/gpt-5-mini" + + +def test_incomplete_total_stays_visible_with_long_model(state, theme): + state.model = "very-long-provider/" + "model" * 20 + state.usage_incomplete = True + line2 = render_statusline(state, theme, width=80).plain.splitlines()[1] + assert "total cost: $0.0123 incomplete" in line2 + assert "84.0k/200.0k" in line2 def test_line3_session_agent_tokens_state(state, theme): diff --git a/tests/test_tui_streaming_pty.py b/tests/test_tui_streaming_pty.py index 0643448..2116ca4 100644 --- a/tests/test_tui_streaming_pty.py +++ b/tests/test_tui_streaming_pty.py @@ -139,6 +139,8 @@ async def _run_pty_app_once( master, slave = pty.openpty() fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", ROWS, COLS, 0, 0)) monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "global-skills")) + monkeypatch.chdir(tmp_path) monkeypatch.setenv("TERM", "xterm-256color") saved_stdout = os.dup(1) @@ -215,6 +217,10 @@ def reader_thread() -> None: # mid-flow fails fast instead of every driver wait timing out on # a dead tty (the two CI flakes looked exactly like that). done, _ = await asyncio.wait({task, driven}, return_when=asyncio.FIRST_COMPLETED) + if driven in done and driven.exception() is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + return await driven if task in done and driven not in done: driven.cancel() with contextlib.suppress(asyncio.CancelledError): @@ -313,6 +319,161 @@ async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: assert any("to @explore" in line for line in lines) +async def test_ctrl_c_focused_child_keeps_root_alive(tmp_path, monkeypatch): + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + os.write(master, b"delegate work\r") + await _wait_for(lambda: screen.display, "workers") + os.write(master, b"/agent 1 focus\r") + await _wait_for(lambda: screen.display, "to @explore") + os.write(master, b"\x03") + await _wait_for(lines, "ROOT SURVIVED") + assert not any("turn cancelled" in line for line in lines()) + os.write(master, b"/quit\r") + return lines() + + await _run_pty_app( + tmp_path, + monkeypatch, + [ + {"tool_calls": [{"name": "task", "arguments": '{"prompt":"child work"}'}]}, + {"text": ["still working "] * 200}, + {"text": "ROOT SURVIVED"}, + ], + drive, + delay=0.03, + ) + + +async def test_denied_focused_prompt_never_executes(tmp_path, monkeypatch): + def configure(config: Config) -> None: + config.hooks = { + "UserPromptSubmit": [ + 'case $(cat) in *blocked-prompt*) echo \'{"verdict":"deny"}\' ;; esac' + ] + } + + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + os.write(master, b"@explore initial\r") + await _wait_for(lines, "started for @explore") + os.write(master, b"/agent 1 focus\r") + await _wait_for(lambda: screen.display, "to @explore") + os.write(master, b"blocked-prompt\r") + await _wait_for(lines, "prompt blocked by hook") + os.write(master, b"/quit\r") + return lines() + + await _run_pty_app( + tmp_path, + monkeypatch, + [ + {"text": "initial"}, + {"text": "NEVER_EXECUTED"}, + ], + drive, + configure=configure, + delay=0.03, + ) + store = SessionStore() + for meta in store.list_sessions(): + messages = store.load_for_model(store.open(meta.id)) + assert not any(m.get("content") in {"blocked-prompt", "NEVER_EXECUTED"} for m in messages) + + +async def test_focus_switch_replaces_visible_worker_transcript(tmp_path, monkeypatch): + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + visible = lambda: screen.display # noqa: E731 + os.write(master, b"@explore FIRST\r") + await _wait_for(lines, "explore · FIRST · worker") + os.write(master, b"@explore SECOND\r") + await _wait_for(lines, "explore · SECOND · worker") + os.write(master, b"/agent 1\r") + await _wait_for(visible, "ANSWER_A") + os.write(master, b"/agent 2 focus\r") + await _wait_for(visible, "ANSWER_B") + assert not any("ANSWER_A" in line for line in visible()) + await _wait_for(visible, "to @explore") + os.write(master, b"/agent root\r") + await _wait_for(visible, "| message |") + assert not any("transcript:" in line for line in visible()) + os.write(master, b"/agent 2 focus\r") + await _wait_for(visible, "to @explore") + os.write(master, b"\x1b") + await _wait_for(visible, "| message |") + assert not any("transcript:" in line for line in visible()) + os.write(master, b"/quit\r") + return lines() + + await _run_pty_app( + tmp_path, monkeypatch, [{"text": "ANSWER_A"}, {"text": "ANSWER_B"}], drive, delay=0.03 + ) + + +async def test_denied_note_composition_never_reaches_model(tmp_path, monkeypatch): + def configure(config: Config) -> None: + config.hooks = { + "UserPromptSubmit": ['case $(cat) in *BLOCKED*) echo \'{"verdict":"deny"}\' ;; esac'] + } + + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + os.write(master, b"/btw BLOCKED\r") + await _wait_for(lines, "noted") + os.write(master, b"go\r") + await _wait_for(lines, "prompt blocked by hook") + os.write(master, b"/quit\r") + return lines() + + await _run_pty_app( + tmp_path, monkeypatch, [{"text": "NEVER_EXECUTED"}], drive, configure=configure, delay=0.03 + ) + store = SessionStore() + session = store.open(store.list_sessions()[0].id) + assert store.load_for_model(session) == [] + assert not any(event.get("text") == "go" for event in store.load_events(session, "input")) + + +async def test_idle_incomplete_total_and_selected_context_render(tmp_path, monkeypatch): + def configure(config: Config) -> None: + config.agent.subagent_model = "anthropic/claude-sonnet-4" + + async def drive(master: int, screen: pyte.HistoryScreen) -> list[str]: + lines = lambda: _screen_lines(screen) # noqa: E731 + visible = lambda: screen.display # noqa: E731 + os.write(master, b"main work\r") + await _wait_for(lines, "answer:") + os.write(master, b"@explore child work\r") + await _wait_for(visible, "total cost: $0.3000") + os.write(master, b"/agent 1 focus\r") + await _wait_for(visible, "2.0k/200.0k") + os.write(master, b"follow up\r") + await _wait_for(visible, "total cost: $0.3000 incomplete") + os.write(master, b"\x1b") + await _wait_for(visible, "| message |") + await _wait_for(visible, "0.1k/1.0M") + assert any("total cost: $0.3000 incomplete" in line for line in visible()) + os.write(master, b"/quit\r") + return lines() + + await _run_pty_app( + tmp_path, + monkeypatch, + [ + {"text": "main", "usage": {"input_tokens": 100, "output_tokens": 10, "cost_usd": 0.1}}, + { + "text": "child", + "usage": {"input_tokens": 2000, "output_tokens": 20, "cost_usd": 0.2}, + }, + {"text": "unknown usage"}, + ], + drive, + configure=configure, + delay=0.03, + ) + + async def test_slash_menu_renders_and_no_match_row(tmp_path, monkeypatch): """Typing '/mod' shows the dropdown on the real terminal (several rows at once, not clipped); an unknown prefix shows the inert 'No matching diff --git a/tests/test_worker_controls.py b/tests/test_worker_controls.py new file mode 100644 index 0000000..f370c7c --- /dev/null +++ b/tests/test_worker_controls.py @@ -0,0 +1,488 @@ +"""Production worker controls through the registered, permission-gated tools.""" + +import asyncio +import json +import shlex +import shutil +import subprocess +import sys + +import pytest +from tests.fakes import FakeProvider + +from lecode.agent.builder import build_runtime +from lecode.agent.tools.base import ToolResult +from lecode.agent.tools.bash import BashTool +from lecode.config.models import Config +from lecode.context.agents import AgentDefinition, AgentRegistry +from lecode.extras.proc import ProcResult +from lecode.extras.subagents import SubagentError +from lecode.permission.checker import AgentOverlay +from lecode.session.storage import SessionStore + + +def git(cwd, *args): + return subprocess.run( + ["git", *args], cwd=cwd, check=True, capture_output=True, text=True + ).stdout.strip() + + +def commit(cwd, text="worker change\n"): + (cwd / "change.txt").write_text(text) + git(cwd, "add", "change.txt") + git(cwd, "-c", "user.name=Test", "-c", "user.email=test@example.com", "commit", "-m", "change") + return git(cwd, "rev-parse", "HEAD") + + +async def dispatch(runtime, action, **args): + _, result = await runtime.registry.dispatch_result( + "control", "workers", json.dumps({"action": action, **args}), runtime.ctx + ) + return result + + +async def review_heads(runtime, worker): + result = await dispatch(runtime, "review", id=worker.id) + assert not result.is_error, result.content + return result.metadata + + +@pytest.fixture +async def workflow(tmp_path, monkeypatch): + monkeypatch.setenv("LECODE_CONFIG_DIR", str(tmp_path / "config")) + monkeypatch.setenv("LECODE_SKILLS_DIR", str(tmp_path / "skills")) + cwd = tmp_path / "repo" + cwd.mkdir() + monkeypatch.chdir(cwd) + git(cwd, "init", "-b", "main") + commit(cwd, "base\n") + config = Config() + config.memory.enabled = config.lsp.enabled = config.pierre.enabled = False + store = SessionStore(tmp_path / "config") + runtime = build_runtime( + config, + cwd, + session=store.create("main", cwd), + store=store, + auto_approve=True, + agent_registry=AgentRegistry( + { + "writer": AgentDefinition( + name="writer", description="Write code", body="", mode="subagent" + ) + } + ), + ) + runtime.ctx.extras["provider"] = FakeProvider([{"text": "done"}] * 12) + try: + yield runtime + finally: + await runtime.ctx.extras["workers"].shutdown() + + +async def start(runtime): + _, result = await runtime.registry.dispatch_result( + "assignment", "task", '{"agent":"writer","prompt":"Implement the assignment"}', runtime.ctx + ) + assert not result.is_error, result.content + return runtime.ctx.extras["workers"].get(result.metadata["worker_id"]) + + +async def test_registered_review_returns_exact_diff_and_commit(workflow): + worker = await start(workflow) + head = commit(worker.cwd) + result = await dispatch(workflow, "review", id=worker.id) + assert not result.is_error, result.content + assert result.metadata["reviewed_head"] == head + parent_head = git(workflow.ctx.cwd, "rev-parse", "HEAD") + assert result.metadata["reviewed_parent_head"] == parent_head + assert f"reviewed_head {head} and reviewed_parent_head {parent_head}" in result.content + assert "-base" in result.content and "+worker change" in result.content + assert "main" in result.content and "Implement the assignment" in result.content + assert ( + "review this exact diff against assignment then integrate reviewed_head" in result.content + ) + + +@pytest.mark.parametrize("approved", [True, False, None]) +async def test_empty_validation_requires_explicit_human_confirmation(workflow, approved): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + questions = [] + if approved is not None: + + def confirm(question): + questions.append(question) + return approved + + workflow.ctx.extras["workers"].confirm = confirm + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error is (approved is not True), result.content + assert (workflow.ctx.cwd / "change.txt").read_text() == ( + "worker change\n" if approved else "base\n" + ) + if approved is not None: + assert len(questions) == 1 + assert worker.id in questions[0] and "@writer" in questions[0] + assert str(worker.cwd) in questions[0] and "validation" in questions[0] + else: + assert "unavailable" in result.content + + +async def test_stale_review_never_merges(workflow): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + commit(worker.cwd, "unreviewed\n") + workflow.ctx.extras["workers"].confirm = lambda _: True + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "re-review" in result.content + assert (workflow.ctx.cwd / "change.txt").read_text() == "base\n" + + +async def test_parent_rewind_requires_review_without_refreshing_hash(workflow): + base = git(workflow.ctx.cwd, "rev-parse", "HEAD") + parent_head = commit(workflow.ctx.cwd, "parent progress\n") + worker = await start(workflow) + head = commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + assert reviewed == {"reviewed_head": head, "reviewed_parent_head": parent_head} + git(workflow.ctx.cwd, "reset", "--keep", base) + workflow.ctx.config.worktree.validation = ["touch must-not-run"] + for _ in range(2): + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "destination HEAD changed since review" in result.content + assert git(workflow.ctx.cwd, "rev-parse", "HEAD") == base + assert git(worker.cwd, "rev-parse", "HEAD") == head + assert not (worker.cwd / "must-not-run").exists() + + +@pytest.mark.parametrize("deny", [False, True]) +async def test_configured_validation_runs_in_child_through_permission_gate(workflow, deny): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + script = "from pathlib import Path; assert Path('change.txt').read_text() == 'worker change\\n'" + workflow.ctx.config.worktree.validation = [ + f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}" + ] + if deny: + workflow.ctx.permission_checker = workflow.ctx.permission_checker.for_agent( + AgentOverlay(denied_tools=("bash",)) + ) + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error is deny, result.content + assert (workflow.ctx.cwd / "change.txt").read_text() == ( + "base\n" if deny else "worker change\n" + ) + if deny: + assert "denied" in result.content + + +async def test_failed_validation_never_merges(workflow): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + workflow.ctx.config.worktree.validation = ["exit 7"] + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "7" in result.content + assert (workflow.ctx.cwd / "change.txt").read_text() == "base\n" + + +async def test_worker_overlay_constrains_validation(workflow): + workflow.ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition( + name="writer", + description="Write", + body="", + mode="subagent", + overlay=AgentOverlay(denied_tools=("bash",)), + ) + } + ) + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + workflow.ctx.config.worktree.validation = ["touch must-not-run"] + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "denied by agent overlay: bash" in result.content + assert not (worker.cwd / "must-not-run").exists() + assert (workflow.ctx.cwd / "change.txt").read_text() == "base\n" + + +async def test_cleanup_refuses_dirty_and_unmerged_then_retains_transcript(workflow): + worker = await start(workflow) + (worker.cwd / "change.txt").write_text("dirty\n") + dirty = await dispatch(workflow, "cleanup", id=worker.id) + assert dirty.is_error and "uncommitted" in dirty.content + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + unmerged = await dispatch(workflow, "cleanup", id=worker.id) + assert unmerged.is_error and "not integrated" in unmerged.content + workflow.ctx.extras["workers"].confirm = lambda _: True + integrated = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert not integrated.is_error, integrated.content + cleaned = await dispatch(workflow, "cleanup", id=worker.id) + assert not cleaned.is_error, cleaned.content + assert not worker.cwd.exists() + manager = workflow.ctx.extras["workers"] + assert manager.get(worker.id).session_id == worker.session_id + assert manager.store.load_for_model(worker.session)[-1]["content"] == "done" + + +@pytest.mark.parametrize("approve", [False, True, None]) +async def test_missing_checkout_recovery_needs_human_confirmation(workflow, approve): + worker = await start(workflow) + head = commit(worker.cwd) + shutil.rmtree(worker.cwd) + questions = [] + if approve is not None: + + async def confirm(question): + questions.append(question) + return approve + + workflow.ctx.extras["workers"].confirm = confirm + result = await dispatch(workflow, "recover", id=worker.id) + assert result.is_error is (approve is not True), result.content + assert worker.cwd.exists() is (approve is True) + if approve: + assert git(worker.cwd, "rev-parse", "HEAD") == head + if approve is not None: + assert worker.id in questions[0] and str(worker.cwd) in questions[0] + assert "uncommitted" in questions[0] and "unrecoverable" in questions[0] + + +@pytest.mark.parametrize("action", ["send", "resume"]) +async def test_write_followup_reconciles_parent_commits(workflow, action): + worker = await start(workflow) + (workflow.ctx.cwd / "parent.txt").write_text("parent progress\n") + git(workflow.ctx.cwd, "add", "parent.txt") + git( + workflow.ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.com", + "commit", + "-m", + "parent", + ) + result = await dispatch(workflow, action, id=worker.id, text="Continue") + assert not result.is_error, result.content + await workflow.ctx.extras["workers"].wait(worker.id) + assert (worker.cwd / "parent.txt").read_text() == "parent progress\n" + + +def supervisor_runtime(runtime, worker): + return runtime.ctx.extras["workers"]._runtime(worker) + + +async def test_nested_integration_requires_immediate_supervisor_and_pinned_target(workflow): + parent = await start(workflow) + nested = supervisor_runtime(workflow, parent) + child = await start(nested) + commit(child.cwd, "nested change\n") + review = await dispatch(nested, "inspect", id=child.id) + assert not review.is_error and parent.worktree.branch in review.content + manager = workflow.ctx.extras["workers"] + manager.confirm = lambda _: True + denied = await dispatch(workflow, "integrate", id=child.id, **review.metadata) + assert denied.is_error and "immediate supervisor" in denied.content + result = await dispatch(nested, "integrate", id=child.id, **review.metadata) + assert not result.is_error, result.content + assert (parent.cwd / "change.txt").read_text() == "nested change\n" + assert (workflow.ctx.cwd / "change.txt").read_text() == "base\n" + + +async def test_maintenance_reserves_workspace_against_send_resume_and_child_start(workflow): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + entered, release = asyncio.Event(), asyncio.Event() + + async def confirm(_): + entered.set() + await release.wait() + return True + + manager = workflow.ctx.extras["workers"] + manager.confirm = confirm + integration = asyncio.create_task(dispatch(workflow, "integrate", id=worker.id, **reviewed)) + try: + async with asyncio.timeout(3): + await entered.wait() + for action in ("send", "resume"): + result = await dispatch(workflow, action, id=worker.id, text="Race") + assert result.is_error and "maintenance" in result.content + nested = supervisor_runtime(workflow, worker) + _, result = await nested.registry.dispatch_result( + "nested", "task", '{"agent":"writer","prompt":"Race"}', nested.ctx + ) + assert result.is_error and "maintenance" in result.content + assert manager.pending(worker.id) == [] + finally: + release.set() + result = await integration + assert not result.is_error, result.content + + +async def test_missing_followup_fails_without_implicit_recreation(workflow): + worker = await start(workflow) + shutil.rmtree(worker.cwd) + result = await dispatch(workflow, "send", id=worker.id, text="Continue") + assert not result.is_error, result.content + with pytest.raises(SubagentError, match="checkout missing"): + await workflow.ctx.extras["workers"].wait(worker.id) + assert not worker.cwd.exists() + + +async def test_dirty_review_requires_checkpoint_before_reviewed_head(workflow): + worker = await start(workflow) + (worker.cwd / "change.txt").write_text("uncommitted work\n") + (worker.cwd / "new.txt").write_text("new work\n") + review = await dispatch(workflow, "review", id=worker.id) + assert not review.is_error, review.content + assert "+uncommitted work" in review.content and "new.txt" in review.content + assert "OWN branch" in review.content and "reviewed_head" not in review.metadata + + +@pytest.mark.parametrize( + "args", + [ + {"action": "cleanup", "discard": True}, + {"action": "recover", "recreate": True}, + {"action": "integrate", "reviewed_head": "main"}, + {"action": "integrate", "reviewed_head": "a" * 40}, + {"action": "integrate", "reviewed_parent_head": "b" * 40}, + {"action": "integrate", "reviewed_head": "a" * 40, "reviewed_parent_head": "main"}, + {"action": "integrate", "reviewed_head": "a" * 40, "reviewed_parent_head": "b" * 39}, + {"action": "integrate", "reviewed_head": "a" * 40, "reviewed_parent_head": 42}, + {"action": "integrate", "reviewed_head": "a" * 40, "allow_unvalidated": True}, + {"action": "integrate", "reviewed_head": "a" * 40, "validation": ["true"]}, + {"action": "resume", "text": 3}, + ], +) +async def test_control_arguments_cannot_override_human_or_validation_gates(workflow, args): + worker = await start(workflow) + result = await dispatch(workflow, id=worker.id, **args) + assert result.is_error + assert worker.cwd.is_dir() + + +async def test_validation_hooks_are_fresh_for_each_child_and_block_execution(workflow, tmp_path): + workers = [await start(workflow), await start(workflow)] + reviews = [] + for worker in workers: + commit(worker.cwd) + reviews.append(await review_heads(workflow, worker)) + log = tmp_path / "validation-hooks.jsonl" + script = ( + "import json,sys; from pathlib import Path; data=json.load(sys.stdin); " + f"p=Path({str(log)!r}); " + "p.open('a').write(json.dumps(data)+'\\n'); " + "print(json.dumps({'verdict':'deny','reason':'validation hook blocked'}))" + ) + workflow.ctx.config.hooks = { + "PreToolUse": [f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}"] + } + workflow.ctx.config.worktree.validation = ["touch should-not-exist"] + for worker, reviewed in zip(workers, reviews, strict=True): + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "validation hook blocked" in result.content + assert not (worker.cwd / "should-not-exist").exists() + entries = [json.loads(line) for line in log.read_text().splitlines()] + assert [entry["cwd"] for entry in entries] == [str(w.cwd) for w in workers] + assert [entry["session"]["id"] for entry in entries] == [w.session_id for w in workers] + assert all(entry["tool"]["args"] == {"command": "touch should-not-exist"} for entry in entries) + _, root_result = await workflow.registry.dispatch_result( + "root", "bash", '{"command":"pwd"}', workflow.ctx + ) + assert not root_result.is_error and str(workflow.ctx.cwd) in root_result.content + assert len(log.read_text().splitlines()) == 2 + + +async def test_readonly_inspection_respects_workers_deny(workflow): + worker = await start(workflow) + commit(worker.cwd) + workflow.ctx.permission_checker = workflow.ctx.permission_checker.for_child(read_only=True) + review = await dispatch(workflow, "inspect", id=worker.id) + assert not review.is_error, review.content + workflow.ctx.permission_checker = workflow.ctx.permission_checker.for_agent( + AgentOverlay(denied_tools=("workers",)) + ) + denied = await dispatch(workflow, "inspect", id=worker.id) + assert denied.is_error and "denied" in denied.content + + +async def test_child_creation_reserves_retained_parent_against_cleanup(workflow, tmp_path): + parent = await start(workflow) + nested = supervisor_runtime(workflow, parent) + entered, release = tmp_path / "checkout-entered", tmp_path / "checkout-release" + hook = workflow.ctx.cwd / ".git/hooks/post-checkout" + hook.write_text( + f"#!/bin/sh\ntouch {shlex.quote(str(entered))}\n" + f"while [ ! -e {shlex.quote(str(release))} ]; do sleep 0.01; done\n" + ) + hook.chmod(0o755) + creation = asyncio.create_task( + nested.registry.dispatch_result( + "child", "task", '{"agent":"writer","prompt":"Nested"}', nested.ctx + ) + ) + try: + async with asyncio.timeout(3): + while not entered.exists(): + await asyncio.sleep(0.01) + result = await dispatch(workflow, "cleanup", id=parent.id) + assert result.is_error and "idle" in result.content + assert parent.cwd.is_dir() + finally: + release.touch() + await creation + + +async def test_cleanup_preserves_retained_child_destination(workflow): + parent = await start(workflow) + child = await start(supervisor_runtime(workflow, parent)) + result = await dispatch(workflow, "cleanup", id=parent.id) + assert result.is_error and "child workspaces first" in result.content + assert parent.cwd.is_dir() and child.cwd.is_dir() + child_cleanup = await dispatch(workflow, "cleanup", id=child.id) + assert not child_cleanup.is_error, child_cleanup.content + parent_cleanup = await dispatch(workflow, "cleanup", id=parent.id) + assert not parent_cleanup.is_error, parent_cleanup.content + + +async def test_truncated_diff_cannot_be_presented_as_exact_review(workflow): + worker = await start(workflow) + commit(worker.cwd, "large change\n" * 100_000) + review = await dispatch(workflow, "inspect", id=worker.id) + assert review.is_error and "truncated" in review.content + assert "reviewed_head" not in review.metadata + + +@pytest.mark.parametrize( + "tool_result", + [ + ToolResult("looks successful"), + ToolResult("tool error", is_error=True, metadata={"proc_result": ProcResult(0, "", "")}), + ], +) +async def test_validation_never_infers_success_without_successful_exit( + workflow, monkeypatch, tool_result +): + worker = await start(workflow) + commit(worker.cwd) + reviewed = await review_heads(workflow, worker) + workflow.ctx.config.worktree.validation = ["true"] + + async def result_only(args, ctx): + return tool_result + + monkeypatch.setattr(BashTool, "run", result_only) + result = await dispatch(workflow, "integrate", id=worker.id, **reviewed) + assert result.is_error and "validation failed" in result.content + assert (workflow.ctx.cwd / "change.txt").read_text() == "base\n" diff --git a/tests/test_workers.py b/tests/test_workers.py index 35a3757..25ee7a3 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -1,21 +1,22 @@ """WorkerManager's public persistence and scheduling contract.""" import asyncio +import json +import shlex import subprocess +import sys from dataclasses import replace -from unittest.mock import Mock import pytest from tests.fakes import FakeProvider +from tests.test_worker_controls import workflow as workflow from lecode.agent.builder import build_runtime from lecode.agent.runner import ( AgentRunner, LlmCall, LlmResponse, - RunResult, Token, - UsageTotals, ) from lecode.agent.runner import Done as RunnerDone from lecode.config.models import Config @@ -23,7 +24,7 @@ from lecode.extras.subagents import SubagentError from lecode.extras.workers import WORKER_CURRENT_EXTRA, WorkerManager from lecode.extras.worktree import WorktreeError, WorktreeManager -from lecode.permission import PermissionChecker +from lecode.hooks import dispatcher_from_config from lecode.permission.checker import AgentOverlay from lecode.session.stats import session_stats from lecode.session.storage import SessionInUseError, SessionStore @@ -70,34 +71,93 @@ def setup(tmp_path, monkeypatch): return manager, runtime.ctx, provider, store, session -@pytest.fixture -def checker_contract(setup): - """Isolate the sibling-owned checker seam in runner/scheduler tests. +@pytest.mark.asyncio +@pytest.mark.parametrize("readonly", [False, True]) +async def test_default_general_task_inherits_permissions_and_isolates_writes(setup, readonly): + from tests.test_worker_controls import commit, git - The real checker integration is tested separately, once for_child lands. - This double does not purport to verify inherited permission restrictions. - """ - _, ctx, _, _, _ = setup - parent = ctx.permission_checker - checker = Mock(wraps=parent) - checker.mode = parent.mode - checker.read_only = parent.read_only + manager, ctx, _, _, _ = setup + git(ctx.cwd, "init", "-b", "main") + commit(ctx.cwd, "base\n") + ctx.auto_approve = True + if readonly: + ctx.permission_checker = ctx.permission_checker.for_agent(AgentOverlay(mode="readonly")) + ctx.extras["provider"] = FakeProvider( + [ + { + "tool_calls": [ + { + "id": "write", + "name": "write", + "arguments": json.dumps({"path": "new.txt", "content": "implemented"}), + } + ] + }, + {"text": "done"}, + ] + ) + try: + _, result = await ctx.extras["registry"].dispatch_result( + "create", + "task", + '{"agent":"general","prompt":"Implement","run_in_background":true}', + ctx, + ) + assert not result.is_error, result.content + worker = manager.get(result.metadata["worker_id"]) + await manager.wait(worker.id) + assert worker.agent == "general" + assert (worker.worktree is None) == readonly + assert (worker.cwd == ctx.cwd) == readonly + assert not (ctx.cwd / "new.txt").exists() + if not readonly: + assert (worker.cwd / "new.txt").read_text() == "implemented" + else: + messages = manager.store.load_for_model(worker.session) + assert any(m["role"] == "tool" and "denied" in m["content"].lower() for m in messages) + finally: + await manager.shutdown() - def derive(overlay=None, **kw): - child = PermissionChecker(ctx.config, mode=parent.mode, _overlay=overlay, **kw) - derived = Mock(wraps=child) - derived.mode = child.mode - derived.read_only = child.read_only - derived.for_child = Mock(side_effect=derive) - return derived - checker.for_child = Mock(side_effect=derive) - ctx.permission_checker = checker - return checker +@pytest.mark.asyncio +async def test_worker_creation_discovery_and_control_errors(setup): + manager, ctx, _, _, _ = setup + registry = ctx.extras["registry"] + try: + for name in ("missing", "build", "plan"): + _, result = await registry.dispatch_result( + "bad-agent", "task", json.dumps({"agent": name, "prompt": "work"}), ctx + ) + assert result.is_error + assert "available: explore, general" in result.content + _, result = await registry.dispatch_result( + "bad-id", "workers", '{"action":"send","id":"general","text":"work"}', ctx + ) + assert result.is_error + assert "unknown worker id" in result.content + assert "task(agent='general'" in result.content + assert "returned worker_id" in result.content + _, result = await registry.dispatch_result("empty", "workers", '{"action":"list"}', ctx) + assert "Create one with task" in result.content + _, result = await registry.dispatch_result( + "no-git", "task", '{"agent":"general","prompt":"work"}', ctx + ) + assert result.is_error + assert "Restart the session from a Git repository with a committed HEAD" in result.content + assert "shell cd does not change" in result.content + assert manager.list() == [] + assert "does not create workers" in registry.get("workers").description + assert "actual worker_id" in registry.get("task").description + assert ( + "workers" + in registry.get("task").parameters["properties"]["run_in_background"]["description"] + ) + finally: + await manager.shutdown() @pytest.mark.asyncio -async def test_followup_is_persisted_and_replayed(setup, checker_contract): +async def test_followup_is_persisted_and_replayed(setup): manager, ctx, provider, store, _ = setup worker = await manager.start(ctx, agent="explore", prompt="question") try: @@ -119,7 +179,7 @@ async def test_followup_is_persisted_and_replayed(setup, checker_contract): @pytest.mark.asyncio -async def test_task_tool_creates_persisted_workers_and_nests(setup, checker_contract): +async def test_task_tool_creates_persisted_workers_and_nests(setup): manager, ctx, _, _, _ = setup ctx.extras["provider"] = FakeProvider( [ @@ -139,7 +199,7 @@ async def test_task_tool_creates_persisted_workers_and_nests(setup, checker_cont @pytest.mark.asyncio -async def test_background_task_tool_delivers_worker_notification(setup, checker_contract): +async def test_background_task_tool_delivers_worker_notification(setup): manager, ctx, _, _, _ = setup try: _, result = await ctx.extras["registry"].dispatch_result( @@ -157,7 +217,7 @@ async def test_background_task_tool_delivers_worker_notification(setup, checker_ @pytest.mark.asyncio -async def test_task_uses_worker_manager_when_tui_events_are_installed(setup, checker_contract): +async def test_task_uses_worker_manager_when_tui_events_are_installed(setup): manager, ctx, _, _, _ = setup ctx.extras["subagent_events"] = lambda event: None try: @@ -171,7 +231,7 @@ async def test_task_uses_worker_manager_when_tui_events_are_installed(setup, che @pytest.mark.asyncio -async def test_worker_forwards_all_runner_events_to_tui_callback(setup, checker_contract): +async def test_worker_forwards_all_runner_events_to_tui_callback(setup): manager, ctx, _, _, _ = setup seen = [] ctx.extras["subagent_events"] = seen.append @@ -190,7 +250,7 @@ async def test_worker_forwards_all_runner_events_to_tui_callback(setup, checker_ @pytest.mark.asyncio -async def test_parent_cancellation_does_not_cancel_managed_worker(setup, checker_contract): +async def test_parent_cancellation_does_not_cancel_managed_worker(setup): manager, ctx, _, _, _ = setup provider = GatedProvider() ctx.extras["provider"] = provider @@ -214,7 +274,7 @@ async def test_parent_cancellation_does_not_cancel_managed_worker(setup, checker @pytest.mark.asyncio -async def test_worker_controls_enforce_descendant_hierarchy_and_questions(setup, checker_contract): +async def test_worker_controls_enforce_descendant_hierarchy_and_questions(setup): from lecode.agent.tools.workers import make_tool manager, ctx, _, _, _ = setup @@ -246,9 +306,7 @@ def git(cwd, *args): @pytest.mark.asyncio -async def test_write_worktree_pins_parent_head_and_readonly_shares_parent_cwd( - setup, checker_contract -): +async def test_write_worktree_pins_parent_head_and_readonly_shares_parent_cwd(setup): manager, ctx, _, _, _ = setup git(ctx.cwd, "init", "-b", "main") git( @@ -346,7 +404,7 @@ async def test_depth_and_agent_eligibility_fail_before_creating_sessions(setup): @pytest.mark.asyncio -async def test_strict_cap_and_shielded_wait(setup, checker_contract): +async def test_strict_cap_and_shielded_wait(setup): manager, ctx, _, _, _ = setup provider = GatedProvider() ctx.extras["provider"] = provider @@ -387,9 +445,7 @@ async def test_send_does_not_wake_stopped_worker_even_with_interrupt(setup): @pytest.mark.asyncio -async def test_failed_usage_survives_restart_and_child_locks_last_until_shutdown( - setup, checker_contract -): +async def test_failed_usage_survives_restart_and_child_locks_last_until_shutdown(setup): manager, ctx, _, store, session = setup ctx.extras["provider"] = FakeProvider( [ @@ -460,7 +516,7 @@ async def test_load_interrupted_uses_child_usage_if_root_snapshot_lagged(setup): @pytest.mark.asyncio -async def test_interrupt_repairs_unanswered_calls_without_replaying_inputs(setup, checker_contract): +async def test_interrupt_repairs_unanswered_calls_without_replaying_inputs(setup): manager, ctx, provider, store, _ = setup worker = await manager.start(ctx, agent="explore", prompt="original") await manager.stop(worker.id) @@ -498,9 +554,7 @@ async def test_interrupt_repairs_unanswered_calls_without_replaying_inputs(setup @pytest.mark.asyncio -async def test_running_inbox_is_durable_but_only_consumed_after_safe_boundary( - setup, checker_contract -): +async def test_running_inbox_is_durable_but_only_consumed_after_safe_boundary(setup): manager, ctx, _, store, _ = setup provider = GatedProvider() ctx.extras["provider"] = provider @@ -526,14 +580,14 @@ async def test_running_inbox_is_durable_but_only_consumed_after_safe_boundary( @pytest.mark.asyncio -async def test_completion_delivery_background_only_and_human_submit(setup, checker_contract): +async def test_completion_delivery_and_idempotent_human_submit(setup): manager, ctx, _, _, _ = setup notifications = [] manager.notify = notifications.append try: foreground = await manager.start(ctx, agent="explore", prompt="foreground") await manager.wait(foreground.id) - assert manager.drain_notifications() == [] + assert [n["worker_id"] for n in manager.drain_notifications()] == [foreground.id] assert notifications == [] background = await manager.start(ctx, agent="explore", prompt="background", background=True) await manager.wait(background.id) @@ -555,7 +609,7 @@ async def test_completion_delivery_background_only_and_human_submit(setup, check @pytest.mark.asyncio -async def test_submit_rejects_delegated_worker(setup, checker_contract): +async def test_submit_rejects_delegated_worker(setup): manager, ctx, _, _, _ = setup try: worker = await manager.start(ctx, agent="explore", prompt="delegated") @@ -567,7 +621,7 @@ async def test_submit_rejects_delegated_worker(setup, checker_contract): @pytest.mark.asyncio -async def test_completed_worker_followup_after_restart(setup, checker_contract): +async def test_completed_worker_followup_after_restart(setup): manager, ctx, provider, store, _ = setup worker = await manager.start(ctx, agent="explore", prompt="first question") await manager.wait(worker.id) @@ -584,49 +638,7 @@ async def test_completed_worker_followup_after_restart(setup, checker_contract): @pytest.mark.asyncio -async def test_suspended_supervisors_release_capacity_and_reacquire( - setup, checker_contract, monkeypatch -): - manager, ctx, _, _, _ = setup - children_started = asyncio.Queue() - release = asyncio.Event() - - async def run(runner, messages, on_event=None): - id = runner.ctx.extras[WORKER_CURRENT_EXTRA] - assert sum(w.state == "running" for w in manager.list()) <= 10 - if manager.get(id).depth == 1: - child = await manager.start(runner.ctx, agent="explore", prompt="child") - async with manager.suspend(id): - assert manager.get(id).state == "waiting" - await manager.wait(child.id) - assert manager.get(id).state == "running" - assert sum(w.state == "running" for w in manager.list()) <= 10 - else: - children_started.put_nowait(id) - await release.wait() - return RunResult("done", 1, "done", UsageTotals()) - - monkeypatch.setattr(AgentRunner, "run", run) - try: - parents = [await manager.start(ctx, agent="explore", prompt="parent") for _ in range(10)] - async with asyncio.timeout(2): - for _ in range(10): - await children_started.get() - assert sum(w.state == "waiting" for w in manager.list()) == 10 - assert sum(w.state == "running" for w in manager.list()) == 10 - with pytest.raises(RuntimeError, match="supervisor"): - async with manager.suspend(manager.children(parents[0].id)[0].id): - pass - release.set() - async with asyncio.timeout(2): - await asyncio.gather(*(manager.wait(w.id) for w in parents)) - assert all(w.state == "completed" for w in manager.list()) - finally: - await manager.shutdown() - - -@pytest.mark.asyncio -async def test_stop_is_individual_unless_tree_requested(setup, checker_contract): +async def test_stop_is_individual_unless_tree_requested(setup): manager, ctx, _, _, _ = setup provider = GatedProvider() ctx.extras["provider"] = provider @@ -651,9 +663,7 @@ async def test_stop_is_individual_unless_tree_requested(setup, checker_contract) @pytest.mark.asyncio -async def test_runtime_is_rebuilt_with_fresh_grants_and_cwd_bound_extras( - setup, checker_contract, monkeypatch -): +async def test_runtime_is_rebuilt_with_fresh_grants_and_cwd_bound_extras(setup, monkeypatch): manager, ctx, _, _, _ = setup ctx.session_perms.grant("bash", "*") ctx.extras["registry"].unregister("bash") @@ -676,26 +686,11 @@ async def run(runner, *args, **kwargs): assert child.extras["background"] is not ctx.extras["background"] assert child.extras["registry"] is not ctx.extras["registry"] assert "bash" not in child.extras["registry"].names() - assert checker_contract.for_child.call_args_list[0].args == ( - ctx.extras["agents"].get("explore").overlay, - ) - assert checker_contract.for_child.call_args_list[0].kwargs == {"cwd": ctx.cwd} - assert checker_contract.for_child.call_args_list[-1].args == ( - ctx.extras["agents"].get("explore").overlay, - ) - assert checker_contract.for_child.call_args_list[-1].kwargs == { - "cwd": child.cwd, - "session_perms": child.session_perms, - "read_only": True, - } + assert child.permission_checker.read_only finally: await manager.shutdown() -@pytest.mark.skipif( - not hasattr(PermissionChecker, "for_child"), - reason="sibling-owned PermissionChecker.for_child has not landed", -) @pytest.mark.asyncio async def test_real_parent_checker_restrictions_reach_child_dispatch(setup): manager, ctx, _, _, _ = setup @@ -717,7 +712,7 @@ async def test_real_parent_checker_restrictions_reach_child_dispatch(setup): @pytest.mark.asyncio -async def test_usage_records_do_not_double_count_child_transcript(setup, checker_contract): +async def test_usage_records_do_not_double_count_child_transcript(setup): manager, ctx, _, store, session = setup ctx.extras["provider"] = FakeProvider( [ @@ -741,7 +736,7 @@ async def test_usage_records_do_not_double_count_child_transcript(setup, checker @pytest.mark.asyncio -async def test_stopped_worker_does_not_return_a_stale_result(setup, checker_contract): +async def test_stopped_worker_does_not_return_a_stale_result(setup): manager, ctx, _, _, _ = setup try: worker = await manager.start(ctx, agent="explore", prompt="question") @@ -801,6 +796,8 @@ async def confirm(question): manager.confirm = confirm worker = await manager.start(ctx, agent="writer", prompt="write") assert len(confirmations) == 1 + assert isinstance(confirmations[0], str) + assert worker.id in confirmations[0] and str(ctx.cwd) in confirmations[0] assert not (worker.cwd / dirty.name).exists() assert dirty.read_text() == "not committed" git(ctx.cwd, "checkout", "--detach") @@ -819,9 +816,7 @@ async def confirm(question): ], ) @pytest.mark.asyncio -async def test_worker_model_precedence_is_recorded( - setup, checker_contract, agent_model, subagent_model, expected -): +async def test_worker_model_precedence_is_recorded(setup, agent_model, subagent_model, expected): manager, ctx, provider, _, _ = setup ctx.config.llm.model = "main-model" ctx.config.agent.subagent_model = subagent_model @@ -837,3 +832,1274 @@ async def test_worker_model_precedence_is_recorded( assert worker.session.meta.model == expected finally: await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restart", [False, True]) +async def test_followup_pins_session_model_in_provider_context_and_events(setup, restart): + manager, ctx, provider, store, session = setup + ctx.config.agent.subagent_model = "pinned-model" + events = [] + ctx.extras["subagent_events"] = events.append + worker = await manager.start(ctx, agent="explore", prompt="initial") + await manager.wait(worker.id) + ctx.config.agent.subagent_model = "new-default" + ctx.extras["agents"] = AgentRegistry( + {"explore": replace(ctx.extras["agents"].get("explore"), model="new-agent-model")} + ) + if restart: + await manager.shutdown() + manager = WorkerManager(ctx.config, cwd=ctx.cwd, root_ctx=ctx, store=store, session=session) + worker = manager.load()[0] + try: + await manager.send(worker.id, "follow up") + await manager.wait(worker.id) + assert [r["model"] for r in provider.requests] == ["pinned-model"] * 2 + assert manager._runtime(worker).ctx.config.llm.model == worker.session.meta.model + assert {p.event.model for p in events if isinstance(p.event, (LlmCall, LlmResponse))} == { + "pinned-model" + } + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reason", ["max_turns", "empty", "context_overflow"]) +async def test_non_success_worker_stop_preserves_result_and_allows_resume(setup, reason): + manager, ctx, _, store, session = setup + tool = { + "tool_calls": [{"id": "read", "name": "read", "arguments": '{"file_path":"missing"}'}], + "usage": {"input_tokens": 900, "cost_usd": 0.25}, + } + if reason == "max_turns": + ctx.config.agent.max_turns = 1 + script = [tool] + elif reason == "empty": + script = [{"usage": {"input_tokens": 7, "cost_usd": 0.25}}] * 4 + else: + ctx.config.agent.context_window = 1000 + ctx.config.compaction.buffer_tokens = 200 + ctx.config.compaction.on_overflow = "pause" + script = [tool, tool, {"text": "summary"}, tool] + provider = FakeProvider(script) + ctx.extras["provider"] = provider + try: + worker = await manager.start(ctx, agent="explore", prompt="work") + with pytest.raises(SubagentError, match=reason): + await manager.wait(worker.id) + assert worker.state == "failed" + assert worker.result.stop_reason == reason + assert worker.result.usage_totals.cost_usd > 0 + assert worker.usage_totals.cost_usd == worker.result.usage_totals.cost_usd + note = store.load_events(session, "worker_notification")[-1] + assert note["state"] == "failed" and reason in note["content"] + ctx.config.agent.max_turns = 10 + ctx.config.compaction.enabled = False + await manager.resume(worker.id, "continue explicitly") + assert (await manager.wait(worker.id)).stop_reason == "done" + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_mixed_tool_batches_release_slots_only_after_ordinary_work(setup, cancel): + manager, ctx, _, store, _ = setup + parents_ready = asyncio.Event() + children_ready = asyncio.Event() + release_children = asyncio.Event() + parent_calls = child_calls = 0 + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + nonlocal parent_calls, child_calls + assert sum(w.state == "running" for w in manager.list()) <= 10 + prompt = messages[1]["content"] + if prompt == "parent" and len(messages) == 2: + parent_calls += 1 + if parent_calls == 10: + parents_ready.set() + await parents_ready.wait() + entry = { + "tool_calls": [ + {"id": "child", "name": "task", "arguments": '{"prompt":"child"}'}, + { + "id": "ordinary", + "name": "read", + "arguments": '{"file_path":"sample.txt"}', + }, + ] + } + elif prompt == "child": + child_calls += 1 + if child_calls == 10: + children_ready.set() + await release_children.wait() + entry = {"text": "child result"} + else: + assert [m["tool_call_id"] for m in messages if m["role"] == "tool"] == [ + "child", + "ordinary", + ] + entry = {"text": "reviewed"} + async for event in self._stream(entry): + yield event + + ctx.cwd.joinpath("sample.txt").write_text("ordinary result") + ctx.extras["provider"] = Provider([]) + try: + parents = [await manager.start(ctx, agent="explore", prompt="parent") for _ in range(10)] + async with asyncio.timeout(3): + await children_ready.wait() + assert sum(w.state == "waiting" for w in manager.list()) == 10 + assert sum(w.state == "running" for w in manager.list()) == 10 + if cancel: + async with asyncio.timeout(2): + await asyncio.gather(*(manager.stop(w.id) for w in parents)) + assert all(w.state == "stopped" and not w.is_active for w in parents) + release_children.set() + if cancel: + async with asyncio.timeout(2): + await asyncio.gather( + *( + manager.wait(child.id) + for parent in parents + for child in manager.children(parent.id) + ) + ) + assert all(w.state == "stopped" for w in parents) + for worker in parents: + assert [ + m["tool_call_id"] + for m in store.load_for_model(worker.session) + if m["role"] == "tool" + ] == ["ordinary"] + return + async with asyncio.timeout(3): + assert all( + r.final_text == "reviewed" + for r in await asyncio.gather(*(manager.wait(w.id) for w in parents)) + ) + for worker in parents: + assert [ + m["tool_call_id"] + for m in store.load_for_model(worker.session) + if m["role"] == "tool" + ] == ["child", "ordinary"] + finally: + release_children.set() + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize("child_error", [False, True]) +async def test_root_reviews_unresolved_delegation_but_not_human_workers( + setup, background, child_error +): + manager, ctx, _, store, session = setup + child_started = asyncio.Event() + root_finished = asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + prompt = messages[1]["content"] + if prompt == "unrelated human": + await root_finished.wait() + entry = {"text": "human result"} + elif prompt == "child": + child_started.set() + await asyncio.sleep(0.02) + entry = {"text": "delegated result", "usage": {"input_tokens": 7}} + if child_error: + entry = {"error": RuntimeError("delegated result")} + elif len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "delegate", + "name": "task", + "arguments": json.dumps( + {"prompt": "child", "run_in_background": background} + ), + } + ] + } + elif any("delegated result" in str(m.get("content")) for m in messages): + entry = {"text": "reviewed delegated result"} + else: + await child_started.wait() + entry = {"text": "premature final"} + async for event in self._stream(entry): + yield event + + provider = Provider([]) + ctx.extras["provider"] = provider + try: + human = await manager.start(ctx, agent="explore", prompt="unrelated human", origin="human") + store.append_message(session, {"role": "user", "content": "root"}) + runner = AgentRunner(provider, ctx.extras["registry"], ctx, session=session, store=store) + async with asyncio.timeout(2): + result = await runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)] + ) + assert result.final_text == "reviewed delegated result" + assert human.state == "running" + history = store.load_for_model(session) + deliveries = [ + m + for m in history + if m["role"] in {"user", "tool"} and "delegated result" in str(m.get("content")) + ] + assert len(deliveries) == 1 + assert session_stats(store, session).input_tokens == (0 if child_error else 7) + finally: + root_finished.set() + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supervisor", ["root", "readonly", "writable"]) +@pytest.mark.parametrize("background", [False, True]) +async def test_child_question_is_answered_at_safe_parent_boundary(setup, supervisor, background): + manager, ctx, _, store, session = setup + nested = supervisor != "root" + if supervisor == "writable": + git(ctx.cwd, "init", "-b", "main") + git( + ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--allow-empty", + "-m", + "initial", + ) + ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + "explore": ctx.extras["agents"].get("explore"), + } + ) + approvals = [] + ctx.approval_callback = lambda *args: approvals.append(args) + seen_questions = [] + denials = [] + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + denials.extend( + m["content"] + for m in messages + if m.get("name") == "workers" + and m["role"] == "tool" + and "denied" in m.get("content", "") + ) + assert not denials, denials + outstanding = set() + for message in messages: + if message["role"] == "tool": + outstanding.remove(message["tool_call_id"]) + else: + assert not outstanding, "question delivery broke the tool protocol" + outstanding.update(c["id"] for c in message.get("tool_calls", [])) + assert not outstanding + prompt = messages[1]["content"] + questions = [ + m["content"] + for m in messages + if m["role"] == "user" and " asks] " in str(m["content"]) + ] + answered = any( + m.get("name") == "workers" and "queued" in m.get("content", "") + for m in messages + if m["role"] == "tool" + ) + if questions and not answered: + seen_questions.append(prompt) + worker_id = questions[-1].split()[1] + entry = { + "tool_calls": [ + { + "id": "answer", + "name": "workers", + "arguments": json.dumps( + {"action": "send", "id": worker_id, "text": "use blue"} + ), + } + ] + } + elif prompt == "leaf": + if len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "ask", + "name": "workers", + "arguments": '{"action":"question","text":"which color?"}', + } + ] + } + else: + assert any(m["content"] == "use blue" for m in messages) + entry = {"text": "blue result"} + elif len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "delegate", + "name": "task", + "arguments": json.dumps( + { + "prompt": "parent" if prompt == "root" and nested else "leaf", + "agent": "writer" + if prompt == "root" and supervisor == "writable" + else "explore", + "run_in_background": background, + } + ), + } + ] + } + elif any("blue result" in str(m.get("content")) for m in messages): + entry = {"text": "reviewed blue result"} + else: + entry = {"text": "premature final"} + async for event in self._stream(entry): + yield event + + provider = Provider([]) + store.append_message(session, {"role": "user", "content": "root"}) + runner = AgentRunner(provider, ctx.extras["registry"], ctx, session=session, store=store) + try: + async with asyncio.timeout(3): + result = await runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)] + ) + assert result.final_text == "reviewed blue result" + assert seen_questions == ["parent" if nested else "root"] + assert approvals == [] + assert all(worker.state == "completed" for worker in manager.list()) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["completed", "failed", "stopped", "queued_stop"]) +async def test_worker_dispatch_lifecycle_hooks_use_parent_context_once(setup, outcome): + manager, ctx, _, _, session = setup + log = ctx.cwd / "hooks.jsonl" + script = ctx.cwd / "hook.py" + script.write_text( + "import json, sys\n" + f"with open({str(log)!r}, 'a') as out:\n" + " out.write(json.dumps(json.load(sys.stdin)) + '\\n')\n" + "print(json.dumps({'verdict': 'deny'}))\n" + ) + command = f"{shlex.quote(sys.executable)} {shlex.quote(str(script))}" + ctx.config.hooks = {event: [command] for event in ("SubagentStart", "SubagentEnd")} + ctx.extras["hooks"], _ = dispatcher_from_config(ctx.config, ctx.cwd, session=session) + provider = ( + GatedProvider() + if outcome == "stopped" + else FakeProvider( + [{"error": RuntimeError("broken")} if outcome == "failed" else {"text": "done"}] + ) + ) + ctx.extras["provider"] = provider + try: + if outcome == "completed": + _, tool_result = await ctx.extras["registry"].dispatch_result( + "delegate", "task", '{"prompt":"work"}', ctx + ) + assert not tool_result.is_error + worker = manager.get(tool_result.metadata["worker_id"]) + else: + worker = await manager.start(ctx, agent="explore", prompt="work") + if outcome == "stopped": + async with asyncio.timeout(2): + await provider.started.get() + await manager.stop(worker.id) + elif outcome == "queued_stop": + await manager.stop(worker.id) + elif outcome == "failed": + with pytest.raises(SubagentError, match="broken"): + await manager.wait(worker.id) + else: + await manager.wait(worker.id) + assert log.exists() + events = [json.loads(line) for line in log.read_text().splitlines()] + assert [event["event"] for event in events] == ["SubagentStart", "SubagentEnd"] + for event in events: + assert event["agent"] == "explore" + assert event["session"]["id"] == session.id + assert event["cwd"] == str(ctx.cwd) + assert event["worker"]["id"] == worker.id + assert event["worker"]["dispatch_id"] == worker.dispatch_id + assert event["worker"]["parent_id"] is None + assert events[-1]["result"]["is_error"] == (outcome != "completed") + if outcome == "queued_stop": + assert provider.requests == [] + if outcome == "completed": + first_dispatch = worker.dispatch_id + await manager.send(worker.id, "follow up") + await manager.wait(worker.id) + events = [json.loads(line) for line in log.read_text().splitlines()] + assert [event["event"] for event in events] == ["SubagentStart", "SubagentEnd"] * 2 + assert [event["worker"]["dispatch_id"] for event in events] == [ + first_dispatch, + first_dispatch, + worker.dispatch_id, + worker.dispatch_id, + ] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_notification_ack_failure_cannot_lose_or_duplicate_delivery(setup, monkeypatch): + manager, ctx, _, store, session = setup + try: + worker = await manager.start(ctx, agent="explore", prompt="work", background=True) + await manager.wait(worker.id) + append = store.append_event + + def fail_ack(session, kind, data): + if kind == "worker_notification_ack": + raise OSError("ack disk failure") + return append(session, kind, data) + + monkeypatch.setattr(store, "append_event", fail_ack) + history = [] + with pytest.raises(OSError, match="ack disk failure"): + manager.consume(None, history) + assert history == store.load_for_model(session) + assert len(history) == 1 + monkeypatch.setattr(store, "append_event", append) + assert manager.consume(None, history) == [] + assert manager.drain_notifications() == [] + assert len(store.load_for_model(session)) == 1 + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["send", "resume"]) +async def test_write_followup_requires_workspace_guard_before_provider(setup, action): + manager, ctx, provider, _, _ = setup + # Missing guard must still fail closed even though production supplies one. + manager.workspace_guard = None + git(ctx.cwd, "init", "-b", "main") + git( + ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--allow-empty", + "-m", + "initial", + ) + ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + } + ) + try: + worker = await manager.start(ctx, agent="writer", prompt="initial") + await manager.wait(worker.id) + if action == "send": + await manager.send(worker.id, "follow up") + else: + await manager.stop(worker.id) + await manager.resume(worker.id, "follow up") + with pytest.raises(SubagentError, match="workspace guard"): + await manager.wait(worker.id) + assert len(provider.requests) == 1 + assert [item["text"] for item in manager.pending(worker.id)] == ["follow up"] + guarded = [] + + async def guard(candidate): + guarded.append(candidate.id) + wm = await WorktreeManager.discover(ctx.cwd) + inspection = await wm.reconcile(candidate.worktree.name, recreate=False) + assert inspection.present and not inspection.merge_in_progress + + manager.workspace_guard = guard + await manager.resume(worker.id) + assert (await manager.wait(worker.id)).final_text == "second" + assert guarded == [worker.id] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_waiting_worker_control_yields_after_read_and_reacquires_at_cap(setup): + manager, ctx, _, store, _ = setup + git(ctx.cwd, "init", "-b", "main") + git( + ctx.cwd, + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--allow-empty", + "-m", + "initial", + ) + end_started = ctx.cwd / "end-started" + release_end = ctx.cwd / "release-end" + script = ctx.cwd / "end.py" + script.write_text( + "import json, sys, time\nfrom pathlib import Path\n" + "event = json.load(sys.stdin)\n" + "if event['agent'] == 'explore':\n" + f" Path({str(end_started)!r}).touch()\n" + f" while not Path({str(release_end)!r}).exists(): time.sleep(0.01)\n" + ) + ctx.config.hooks = { + "SubagentEnd": [f"{shlex.quote(sys.executable)} {shlex.quote(str(script))}"] + } + ctx.extras["hooks"], _ = dispatcher_from_config(ctx.config, ctx.cwd, session=ctx.session) + ctx.extras["agents"] = AgentRegistry( + { + "writer": AgentDefinition("writer", "write", "", mode="subagent"), + "explore": ctx.extras["agents"].get("explore"), + } + ) + manager.confirm = lambda _: True + blockers_ready = asyncio.Event() + release_blockers = asyncio.Event() + blockers = 0 + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + nonlocal blockers + assert sum(w.state == "running" for w in manager.list()) <= 10 + prompt = messages[1]["content"] + if prompt == "blocker": + blockers += 1 + if blockers == 9: + blockers_ready.set() + await release_blockers.wait() + entry = {"text": "done"} + elif prompt == "parent" and len(messages) == 2: + await blockers_ready.wait() + entry = { + "tool_calls": [ + { + "id": "child", + "name": "task", + "arguments": '{"prompt":"never runs","run_in_background":true}', + } + ] + } + elif prompt == "parent" and not any(m.get("tool_call_id") == "stop" for m in messages): + parent = next(w for w in manager.list() if w.agent == "writer") + child = manager.children(parent.id)[0] + entry = { + "tool_calls": [ + { + "id": "stop", + "name": "workers", + "arguments": json.dumps({"action": "stop", "id": child.id}), + }, + {"id": "read", "name": "read", "arguments": '{"file_path":"missing"}'}, + ] + } + else: + assert prompt == "parent", "queued child must be stopped before its model runs" + entry = {"text": "reviewed"} + async for event in self._stream(entry): + yield event + + ctx.extras["provider"] = Provider([]) + try: + for _ in range(9): + await manager.start(ctx, agent="explore", prompt="blocker", origin="human") + parent = await manager.start(ctx, agent="writer", prompt="parent") + async with asyncio.timeout(3): + while not end_started.exists() or parent.state != "waiting": + await asyncio.sleep(0.01) + assert sum(w.state == "running" for w in manager.list()) == 9 + release_end.touch() + async with asyncio.timeout(3): + assert (await manager.wait(parent.id)).final_text == "reviewed" + assert [ + m["tool_call_id"] for m in store.load_for_model(parent.session) if m["role"] == "tool" + ] == ["child", "stop", "read"] + finally: + release_end.touch() + release_blockers.set() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_root_reviews_orphaned_foreground_descendant_without_waking_stopped_parent(setup): + manager, ctx, _, store, session = setup + child_started = asyncio.Event() + root_called = asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + prompt = messages[1]["content"] + if prompt == "parent": + entry = {"text": "parent done"} + elif prompt == "descendant": + child_started.set() + await root_called.wait() + entry = {"text": "orphan result"} + else: + root_called.set() + entry = { + "text": "reviewed" + if any("orphan result" in str(m.get("content")) for m in messages) + else "premature" + } + async for event in self._stream(entry): + yield event + + provider = Provider([]) + ctx.extras["provider"] = provider + try: + parent = await manager.start(ctx, agent="explore", prompt="parent") + await manager.wait(parent.id) + await manager.stop(parent.id) + nested = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: parent.id}) + await manager.start(nested, agent="explore", prompt="descendant") + await child_started.wait() + store.append_message(session, {"role": "user", "content": "root"}) + runner = AgentRunner(provider, ctx.extras["registry"], ctx, session=session, store=store) + async with asyncio.timeout(2): + result = await runner.run( + [{"role": "system", "content": "root"}, *store.load_for_model(session)] + ) + assert result.final_text == "reviewed" + assert parent.state == "stopped" + assert manager.pending_notifications(parent.id) == [] + finally: + root_called.set() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_worker_remains_active_until_lifecycle_cleanup_finishes(setup): + manager, ctx, _, _, session = setup + started = ctx.cwd / "end-started" + release = ctx.cwd / "end-release" + script = ctx.cwd / "end.py" + script.write_text( + "import time\nfrom pathlib import Path\n" + f"Path({str(started)!r}).touch()\n" + f"while not Path({str(release)!r}).exists(): time.sleep(0.01)\n" + ) + ctx.config.hooks = { + "SubagentEnd": [f"{shlex.quote(sys.executable)} {shlex.quote(str(script))}"] + } + ctx.extras["hooks"], _ = dispatcher_from_config(ctx.config, ctx.cwd, session=session) + try: + worker = await manager.start(ctx, agent="explore", prompt="work") + async with asyncio.timeout(2): + while not started.exists(): + await asyncio.sleep(0.01) + assert worker.is_active + with pytest.raises(RuntimeError, match="active"): + manager.attach(session) + release.touch() + await manager.wait(worker.id) + assert not worker.is_active + finally: + release.touch() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_unanswered_question_escalates_when_its_parent_is_stopped(setup): + manager, ctx, _, _, _ = setup + ctx.extras["provider"] = GatedProvider() + try: + parent = await manager.start(ctx, agent="explore", prompt="parent") + nested = replace(ctx, extras={**ctx.extras, WORKER_CURRENT_EXTRA: parent.id}) + child = await manager.start(nested, agent="explore", prompt="child") + manager.ask_parent(child.id, "which color?") + parent_history = [] + manager.consume(parent.id, parent_history) + assert "which color?" in parent_history[-1]["content"] + await manager.stop(parent.id) + root_history = [] + manager.consume(None, root_history) + assert any("which color?" in m["content"] for m in root_history) + assert manager.consume(None, root_history) == [] + await manager.send(child.id, "use blue") + assert parent.state == "stopped" + assert manager.questions(child.id) == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_parent_questions_yield_all_ten_child_leases(setup): + manager, ctx, _, _, _ = setup + spare_started = asyncio.Event() + release_spare = asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + if messages[1]["content"] == "spare": + spare_started.set() + await release_spare.wait() + entry = {"text": "done"} + elif len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "question", + "name": "workers", + "arguments": '{"action":"question","text":"need an answer"}', + } + ] + } + else: + assert messages[-1]["content"] == "answer" + entry = {"text": "done"} + assert sum(w.state == "running" for w in manager.list()) <= 10 + async for event in self._stream(entry): + yield event + + ctx.extras["provider"] = Provider([]) + try: + children = [await manager.start(ctx, agent="explore", prompt="ask") for _ in range(10)] + spare = await manager.start(ctx, agent="explore", prompt="spare", origin="human") + async with asyncio.timeout(2): + await spare_started.wait() + while any(w.state != "waiting" for w in children): + await asyncio.sleep(0) + for child in children: + await manager.send(child.id, "answer") + release_spare.set() + async with asyncio.timeout(2): + await asyncio.gather(*(manager.wait(w.id) for w in [*children, spare])) + assert all(w.state == "completed" for w in children) + finally: + release_spare.set() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_batch_persists_completed_delegation_once_across_resume(setup): + manager, ctx, _, store, _ = setup + slow_started = asyncio.Event() + release = asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + prompt = messages[1]["content"] + if prompt == "parent" and len(messages) == 2: + entry = { + "tool_calls": [ + {"id": "fast", "name": "task", "arguments": '{"prompt":"fast"}'}, + {"id": "slow", "name": "task", "arguments": '{"prompt":"slow"}'}, + ] + } + elif prompt == "slow": + slow_started.set() + await release.wait() + entry = {"text": "slow result"} + elif prompt == "fast": + entry = {"text": "fast result"} + else: + assert sum("fast result" in str(m.get("content")) for m in messages) == 1 + entry = {"text": "reviewed"} + async for event in self._stream(entry): + yield event + + ctx.extras["provider"] = Provider([]) + try: + parent = await manager.start(ctx, agent="explore", prompt="parent") + async with asyncio.timeout(2): + await slow_started.wait() + while not any( + w.state == "completed" and not w.is_active for w in manager.children(parent.id) + ): + await asyncio.sleep(0) + # The fast tool can return through its supervisor's reacquisition. + await asyncio.sleep(0.01) + await manager.stop(parent.id) + assert [ + m["tool_call_id"] for m in store.load_for_model(parent.session) if m["role"] == "tool" + ] == ["fast"] + release.set() + await manager.resume(parent.id) + async with asyncio.timeout(2): + assert (await manager.wait(parent.id)).final_text == "reviewed" + finally: + release.set() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_wait_includes_followup_queued_during_completion_notification(setup): + manager, ctx, _, _, _ = setup + + async def notify(note): + if note["content"] == "first": + await manager.send(note["worker_id"], "follow up") + + manager.notify = notify + try: + worker = await manager.start(ctx, agent="explore", prompt="work", background=True) + assert (await manager.wait(worker.id)).final_text == "second" + assert not worker.is_active + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_depth_two_approvals_keep_requester_identity_fifo_and_child_grants(setup): + from lecode.config.models import PermissionRule + from lecode.permission import AllowAlways, AllowOnce + from lecode.tui.permission import ApprovalPrompt + + manager, ctx, _, store, _ = setup + ctx.config.permissions.rules.ask["read"] = [PermissionRule(pattern="*")] + fifo = ApprovalPrompt() + requested = asyncio.Queue() + + async def approve(name, args, reason, *, worker, conversation): + future = fifo.request( + name, args["file_path"], reason, worker=worker, conversation=conversation + ) + requested.put_nowait((worker, conversation)) + return await future + + ctx.approval_callback = approve + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + if len(messages) > 2: + entry = {"text": "done"} + elif messages[1]["content"] == "parent": + entry = { + "tool_calls": [ + {"id": "read", "name": "read", "arguments": '{"file_path":"parent.txt"}'}, + {"id": "spawn", "name": "task", "arguments": '{"prompt":"child"}'}, + ] + } + else: + entry = { + "tool_calls": [ + {"id": "one", "name": "read", "arguments": '{"file_path":"one.txt"}'}, + {"id": "two", "name": "read", "arguments": '{"file_path":"two.txt"}'}, + ] + } + async for event in self._stream(entry): + yield event + + ctx.extras["provider"] = Provider([]) + try: + parent = await manager.start(ctx, agent="explore", prompt="parent") + async with asyncio.timeout(2): + identities = [await requested.get() for _ in range(3)] + child = manager.children(parent.id)[0] + assert child.depth == 2 + assert identities == [ + (parent.id, parent.session.name), + (child.id, child.session.name), + (child.id, child.session.name), + ] + for identity, decision in zip( + identities, [AllowOnce(), AllowAlways("one.txt"), AllowOnce()], strict=True + ): + assert (fifo.pending.worker, fifo.pending.conversation) == identity + fifo.resolve(decision) + await manager.wait(parent.id) + assert store.load_grants(child.session) == [("read", "one.txt")] + assert store.load_grants(parent.session) == [] + assert ctx.session_perms.grants == [] + assert not fifo.is_pending + finally: + fifo.cancel() + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("running_supervisor", [False, True]) +@pytest.mark.parametrize("validation", [False, True]) +async def test_nested_integration_reserves_destination_execution( + workflow, monkeypatch, running_supervisor, validation +): + from tests.test_worker_controls import commit, dispatch, start, supervisor_runtime + + from lecode.agent.tools.bash import BashTool + + ctx = workflow.ctx + manager = ctx.extras["workers"] + entered, release = asyncio.Event(), asyncio.Event() + questions = [] + + async def confirm(question): + assert isinstance(question, str) + questions.append(question) + entered.set() + await release.wait() + return True + + manager.confirm = confirm + if validation: + manager.config.worktree.validation = ["true"] + original = BashTool.run + + async def validate(tool, args, context): + entered.set() + await release.wait() + return await original(tool, args, context) + + monkeypatch.setattr(BashTool, "run", validate) + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + if messages[1]["content"] == "child": + entry = {"text": "child done"} + elif len(messages) == 2: + entry = { + "tool_calls": [ + { + "id": "spawn", + "name": "task", + "arguments": '{"agent":"writer","prompt":"child"}', + } + ] + } + elif not any(m.get("tool_call_id") == "integrate" for m in messages): + parent = manager.children(None)[0] + child = manager.children(parent.id)[0] + entry = { + "tool_calls": [ + { + "id": "integrate", + "name": "workers", + "arguments": json.dumps( + { + "action": "integrate", + "id": child.id, + "reviewed_head": commit(child.cwd, "nested change\n"), + "reviewed_parent_head": git(parent.cwd, "rev-parse", "HEAD"), + } + ), + } + ] + } + else: + result = next(m for m in messages if m.get("tool_call_id") == "integrate") + assert "error" not in result["content"], result["content"] + entry = {"text": "integrated"} + async for event in self._stream(entry): + yield event + + integration = None + if running_supervisor: + ctx.extras["provider"] = Provider([]) + parent = await manager.start(ctx, agent="writer", prompt="parent") + else: + parent = await start(workflow) + nested = supervisor_runtime(workflow, parent) + child = await start(nested) + head = commit(child.cwd, "nested change\n") + integration = asyncio.create_task( + dispatch( + nested, + "integrate", + id=child.id, + reviewed_head=head, + reviewed_parent_head=git(parent.cwd, "rev-parse", "HEAD"), + ) + ) + try: + async with asyncio.timeout(3): + await entered.wait() + child = manager.children(parent.id)[0] + for worker in (parent, child): + with pytest.raises(WorktreeError, match="maintenance"): + await manager.send(worker.id, "race") + with pytest.raises(WorktreeError, match="maintenance"): + await manager.resume(worker.id, "race") + with pytest.raises(WorktreeError, match="maintenance"): + await manager.reconcile_workspace(worker) + assert manager.pending(worker.id) == [] + if not validation: + assert child.id in questions[0] and str(child.cwd) in questions[0] + release.set() + if integration: + result = await integration + assert not result.is_error, result.content + else: + assert (await manager.wait(parent.id)).final_text == "integrated" + assert (parent.cwd / "change.txt").read_text() == "nested change\n" + assert (ctx.cwd / "change.txt").read_text() == "base\n" + assert not manager._maintenance + finally: + release.set() + if integration: + await integration + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["cleanup", "recover"]) +async def test_parent_workspace_maintenance_reserves_human_origin_shared_subtree( + workflow, monkeypatch, action +): + import shutil + + from tests.test_worker_controls import dispatch, start, supervisor_runtime + + ctx = workflow.ctx + manager = ctx.extras["workers"] + ctx.extras["agents"] = AgentRegistry( + { + "writer": ctx.extras["agents"].get("writer"), + "reader": AgentDefinition( + "reader", "read", "", mode="subagent", overlay=AgentOverlay(mode="readonly") + ), + } + ) + parent = await start(workflow) + nested = supervisor_runtime(workflow, parent) + child = await manager.start(nested.ctx, agent="reader", prompt="read", origin="human") + await manager.wait(child.id) + assert child.cwd == parent.cwd and child.worktree is None + entered, release = asyncio.Event(), asyncio.Event() + + async def confirm(question): + assert isinstance(question, str) + assert parent.id in question and str(parent.cwd) in question + entered.set() + await release.wait() + return True + + if action == "recover": + shutil.rmtree(parent.cwd) + manager.confirm = confirm + else: + original = WorktreeManager.cleanup_worker + + async def cleanup(worktrees, name): + entered.set() + await release.wait() + return await original(worktrees, name) + + monkeypatch.setattr(WorktreeManager, "cleanup_worker", cleanup) + + maintenance = asyncio.create_task(dispatch(workflow, action, id=parent.id)) + try: + async with asyncio.timeout(3): + await entered.wait() + for worker in (parent, child): + with pytest.raises(WorktreeError, match="maintenance"): + await manager.send(worker.id, "race") + with pytest.raises(WorktreeError, match="maintenance"): + await manager.resume(worker.id, "race") + with pytest.raises(WorktreeError, match="maintenance"): + await manager.start(nested.ctx, agent="reader", prompt="race") + release.set() + result = await maintenance + assert not result.is_error, result.content + assert parent.cwd.exists() is (action == "recover") + assert manager.pending(parent.id) == manager.pending(child.id) == [] + assert not manager._maintenance + finally: + release.set() + await maintenance + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", [False, True], ids=["root", "parent-worker"]) +@pytest.mark.parametrize("rewrite", [False, True], ids=["integrate", "rewritten-inspect"]) +@pytest.mark.parametrize("validation", [False, True], ids=["confirmation", "validation"]) +async def test_integration_refuses_unfinished_sibling_bash( + workflow, monkeypatch, nested, rewrite, validation +): + from tests.test_worker_controls import commit, start, supervisor_runtime + + from lecode.agent.tools.base import ToolRegistry + from lecode.agent.tools.bash import BashTool + from lecode.hooks import apply_hooks + from lecode.hooks.runner import HookDispatcher, HookHandler + + manager = workflow.ctx.extras["workers"] + parent = await start(workflow) if nested else None + runtime = supervisor_runtime(workflow, parent) if nested else workflow + child = await start(runtime) + args = { + "action": "integrate", + "id": child.id, + "reviewed_head": commit(child.cwd), + "reviewed_parent_head": git(runtime.ctx.cwd, "rev-parse", "HEAD"), + } + release_bash, bash_finished = asyncio.Event(), asyncio.Event() + mutation_reservations, checks, controls, bash_results = [], [], [], [] + owner = parent.id if parent else None + + def assert_lease(): + assert len(manager._leases) <= 10 + if parent: + assert parent.id in manager._leases + assert parent.state == "running" + + inspect = WorktreeManager.inspect + + async def inspect_workspace(worktrees, name): + if name == child.worktree.name: + checks.append("inspect") + assert_lease() + # On the buggy path, let the sibling mutate the reserved destination. + release_bash.set() + await bash_finished.wait() + return await inspect(worktrees, name) + + monkeypatch.setattr(WorktreeManager, "inspect", inspect_workspace) + control = runtime.registry.get("workers") + run_control = control.run + + async def observed_control(arguments, context): + assert asyncio.current_task() in manager._tool_owners + assert manager._tool_owners[asyncio.current_task()] == owner + try: + result = await run_control(arguments, context) + controls.append((result, list(checks))) + return result + finally: + release_bash.set() + + monkeypatch.setattr(control, "run", observed_control) + if rewrite: + output = json.dumps({"verdict": "allow", "rewritten_input": args}) + command = f"{shlex.quote(sys.executable)} -c {shlex.quote(f'print({output!r})')}" + apply_hooks( + ToolRegistry([control]), + HookDispatcher({"PreToolUse": [HookHandler(command)]}, runtime.ctx.cwd), + ) + + bash = runtime.registry.get("bash") + run_bash = bash.run + + async def mutate_destination(arguments, context): + await release_bash.wait() + assert_lease() + mutation_reservations.append(bool(manager._maintenance)) + try: + result = await run_bash(arguments, context) + bash_results.append(result) + return result + finally: + bash_finished.set() + + monkeypatch.setattr(bash, "run", mutate_destination) + + async def confirm(question): + checks.append("confirmation") + assert bash_finished.is_set() + assert manager._maintenance + return True + + manager.confirm = confirm + if validation: + manager.config.worktree.validation = ["true"] + original = BashTool.run + + async def validate(tool, arguments, context): + checks.append("validation") + assert bash_finished.is_set() + assert manager._maintenance + assert_lease() + return await original(tool, arguments, context) + + monkeypatch.setattr(BashTool, "run", validate) + + integration = { + "id": "integrate", + "name": "workers", + "arguments": json.dumps({"action": "inspect", "id": child.id} if rewrite else args), + } + provider = FakeProvider( + [ + { + "tool_calls": [ + integration, + { + "id": "mutate", + "name": "bash", + "arguments": json.dumps( + {"command": "rm change.txt && git restore -- change.txt"} + ), + }, + ] + }, + {"tool_calls": [{**integration, "id": "retry"}]}, + {"text": "integrated"}, + ] + ) + runtime.ctx.extras["provider"] = provider + async with asyncio.timeout(5): + if parent: + await manager.resume(parent.id, "integrate child") + result = await manager.wait(parent.id) + else: + runner = AgentRunner( + provider, + runtime.registry, + runtime.ctx, + session=runtime.ctx.session, + store=manager.store, + ) + result = await runner.run([{"role": "user", "content": "integrate child"}]) + assert result.final_text == "integrated" + assert len(bash_results) == 1 + assert not bash_results[0].is_error, bash_results[0].content + assert mutation_reservations == [False], "sibling mutated a reserved destination" + refused, early_checks = controls[0] + assert refused.is_error + assert ( + "workspace maintenance must run separately after sibling tools complete" in refused.content + ) + assert early_checks == [], "mixed batch reached inspection, confirmation or validation" + assert not controls[1][0].is_error, controls[1][0].content + assert ("validation" if validation else "confirmation") in checks + assert (runtime.ctx.cwd / "change.txt").read_text() == "worker change\n" + assert not manager._maintenance + assert not manager._tool_owners + + +@pytest.mark.asyncio +async def test_independent_maintenance_cannot_borrow_running_supervisor(workflow): + from tests.test_worker_controls import start, supervisor_runtime + + manager = workflow.ctx.extras["workers"] + started, release = asyncio.Event(), asyncio.Event() + + class Provider(FakeProvider): + async def stream_chat(self, messages, **kwargs): + if any(m.get("content") == "active" for m in messages): + started.set() + await release.wait() + async for event in self._stream({"text": "done"}): + yield event + + workflow.ctx.extras["provider"] = Provider([]) + parent = await start(workflow) + child = await start(supervisor_runtime(workflow, parent)) + try: + await manager.resume(parent.id, "active") + async with asyncio.timeout(2): + await started.wait() + with pytest.raises(WorktreeError, match="idle"): + async with manager.maintain_workspace(child.id): + pytest.fail("independent execution entered the running parent's workspace") + release.set() + await manager.wait(parent.id) + assert not manager._maintenance + finally: + release.set() diff --git a/tests/test_worktree.py b/tests/test_worktree.py index 5abd936..28a1a1e 100644 --- a/tests/test_worktree.py +++ b/tests/test_worktree.py @@ -7,15 +7,19 @@ from __future__ import annotations +import asyncio +import shutil import subprocess +import sys from os.path import realpath +from pathlib import Path import pytest from tests.test_tui_app import FakeTui, _name_prompt, make_app from typer.testing import CliRunner from lecode.cli import app as cli_app -from lecode.extras.proc import run_proc +from lecode.extras.proc import ProcResult, run_proc from lecode.extras.worktree import WorktreeError, WorktreeManager runner = CliRunner() @@ -237,6 +241,9 @@ async def test_create_worker_pins_base_and_writes_sidecar(tmp_path): "dest_path": str(repo), "dest_branch": "main", "dest_common_dir": str(repo / ".git"), + "repo_identity": manager._identity(repo / ".git"), + "dest_identity": manager._identity(repo), + "checkout_identity": manager._identity(info.path), "integrated_at": None, "integrated_head": None, } @@ -286,6 +293,658 @@ async def test_inspect_reports_state(tmp_path): assert absent.merge_in_progress is False +async def worker_repo(tmp_path): + repo = await make_repo(tmp_path / "repo") + manager = WorktreeManager(repo) + info = await manager.create_worker( + "worker", base_commit="HEAD", dest_path=repo, dest_branch="main" + ) + (info.path / "feature.txt").write_text("feature\n") + await commit_all(info.path, "worker feature") + return repo, manager, info + + +async def validate_git(cmd, cwd): + assert cmd == "check" + return await run_proc(["git", "diff", "--exit-code", "HEAD"], cwd=cwd) + + +async def integrate_head(manager, info, *, reviewed_parent_head=None, **kwargs): + if reviewed_parent_head is None: + reviewed_parent_head = await git( + manager.read_sidecar(info.name)["dest_path"], "rev-parse", "HEAD" + ) + return await manager.integrate( + info.name, + reviewed_head=await git(info.path, "rev-parse", "HEAD"), + reviewed_parent_head=reviewed_parent_head, + validation=["check"], + validation_runner=validate_git, + **kwargs, + ) + + +async def test_discover_unrelated_linked_relative_subdir(tmp_path, monkeypatch): + repo = await make_repo(tmp_path / "repo") + linked = tmp_path / "linked" + await git(repo, "worktree", "add", "-b", "parent", str(linked)) + (linked / "sub").mkdir() + monkeypatch.chdir(tmp_path) + manager = await WorktreeManager.discover(Path("linked/sub")) + assert manager.repo_root == linked + info = await manager.create_worker( + "child", base_commit="HEAD", dest_path="linked/sub", dest_branch="parent" + ) + assert manager.read_sidecar("child")["dest_path"] == str(linked) + assert await git(linked, "status", "--porcelain") == "" + (info.path / "feature").write_text("child") + await commit_all(info.path, "child") + assert (await WorktreeManager.discover(info.path)).repo_root == linked + main_head = await git(repo, "rev-parse", "HEAD") + assert (await integrate_head(manager, info)).merged + assert await git(repo, "rev-parse", "HEAD") == main_head + assert (linked / "feature").read_text() == "child" + + +async def test_reconcile_merges_parent_and_retains_dirty_worker(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + worker_head = await git(info.path, "rev-parse", "HEAD") + (repo / "parent.txt").write_text("parent") + await commit_all(repo, "parent progress") + parent_head = await git(repo, "rev-parse", "HEAD") + (info.path / "dirty.txt").write_text("unfinished") + assert (await manager.reconcile(info.name)).dirty + assert await git(info.path, "rev-parse", "HEAD") == worker_head + (info.path / "dirty.txt").unlink() + # Parent dirt is not copied into the worker and must not be lost either. + (repo / "uncommitted").write_text("parent unfinished") + state = await manager.reconcile(info.name) + assert not state.dirty + assert await git(info.path, "rev-parse", "HEAD^1") == worker_head + assert await git(info.path, "rev-parse", "HEAD^2") == parent_head + assert not (info.path / "uncommitted").exists() + assert (repo / "uncommitted").read_text() == "parent unfinished" + + +async def test_reconcile_conflict_left_visible(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + (repo / "file.txt").write_text("parent") + await commit_all(repo, "parent edit") + (info.path / "file.txt").write_text("worker") + await commit_all(info.path, "worker edit") + head = await git(info.path, "rev-parse", "HEAD") + with pytest.raises(WorktreeError, match="merge"): + await manager.reconcile(info.name) + assert (await manager.reconcile(info.name)).merge_in_progress + assert await git(info.path, "rev-parse", "HEAD") == head + assert "<<<<<<<" in (info.path / "file.txt").read_text() + assert (repo / "file.txt").read_text() == "parent" + + +@pytest.mark.parametrize("keep_branch", [True, False]) +async def test_missing_worker_explicit_recovery(tmp_path, keep_branch): + repo, manager, info = await worker_repo(tmp_path) + head = await git(info.path, "rev-parse", "HEAD") + sidecar = manager.read_sidecar(info.name) + (info.path / "lost").write_text("uncommitted") + shutil.rmtree(info.path) + if not keep_branch: + await git(repo, "update-ref", "-d", f"refs/heads/{info.branch}") + with pytest.raises(WorktreeError, match="Missing uncommitted content cannot be recovered"): + await manager.reconcile(info.name) + assert not info.path.exists() + assert manager.read_sidecar(info.name) == sidecar + state = await manager.reconcile(info.name, recreate=True) + assert state.present and not state.dirty + expected = head if keep_branch else manager.read_sidecar(info.name)["base_commit"] + assert await git(info.path, "rev-parse", "HEAD") == expected + assert not (info.path / "lost").exists() + assert manager.read_sidecar(info.name) == { + **sidecar, + "checkout_identity": manager._identity(info.path), + } + assert await WorktreeManager(repo).attach(info.name) == info + + +@pytest.mark.parametrize("change", ["directory", "repo", "branch", "detached"]) +async def test_reconcile_refuses_replaced_worker(tmp_path, change): + repo, manager, info = await worker_repo(tmp_path) + if change in {"directory", "repo"}: + await git(repo, "worktree", "remove", str(info.path)) + if change == "repo": + await make_repo(info.path) + else: + info.path.mkdir() + elif change == "branch": + await git(info.path, "switch", "-c", "unrelated") + assert (await WorktreeManager.discover(info.path)).repo_root == repo + else: + await git(info.path, "switch", "--detach") + with pytest.raises(WorktreeError): + await manager.reconcile(info.name, recreate=True) + with pytest.raises(WorktreeError): + await manager.cleanup_worker(info.name, discard=True) + assert info.path.is_dir() + + +async def test_explicit_recreation_repins_checkout_inode(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + sidecar = manager.read_sidecar(info.name) + (info.path / "unfinished").write_text("retained original work") + saved = tmp_path / "saved-worker" + info.path.rename(saved) + manager = WorktreeManager(repo) + assert not (await manager.inspect(info.name)).present + with pytest.raises(WorktreeError, match="Missing uncommitted content cannot be recovered"): + await manager.reconcile(info.name) + assert manager.read_sidecar(info.name) == sidecar + assert not info.path.exists() + recovered = await manager.reconcile(info.name, recreate=True) + identity = manager._identity(info.path) + assert recovered.present and not recovered.dirty + assert identity != sidecar["checkout_identity"] + assert recovered.sidecar == {**sidecar, "checkout_identity": identity} + assert (saved / "unfinished").read_text() == "retained original work" + assert not (info.path / "unfinished").exists() + assert await manager.attach(info.name) == info + assert (await integrate_head(manager, info)).merged + await manager.cleanup_worker(info.name) + assert manager.read_sidecar(info.name)["checkout_identity"] == identity + + +@pytest.mark.parametrize("action", ["inspect", "attach", "reconcile", "integrate", "cleanup"]) +async def test_worker_inode_replacement_with_same_git_pointer_refused(tmp_path, action): + repo, manager, info = await worker_repo(tmp_path) + sidecar = manager.read_sidecar(info.name) + (info.path / "unfinished").write_text("irreplaceable work") + saved = tmp_path / "saved-worker" + info.path.rename(saved) + info.path.mkdir() + (info.path / ".git").write_text((saved / ".git").read_text()) + await git(info.path, "restore", ".") + assert await git(info.path, "status", "--porcelain") == "" + assert await git(info.path, "branch", "--show-current") == info.branch + manager = WorktreeManager(repo) + with pytest.raises(WorktreeError, match="checkout identity changed"): + if action == "integrate": + await integrate_head(manager, info) + elif action == "cleanup": + await manager.cleanup_worker(info.name, discard=True) + elif action == "reconcile": + await manager.reconcile(info.name, recreate=True) + else: + await getattr(manager, action)(info.name) + assert info.path.is_dir() + assert (saved / "unfinished").read_text() == "irreplaceable work" + assert manager.read_sidecar(info.name) == sidecar + + +async def test_integrate_rejects_parent_rewind_before_validation(tmp_path): + repo = await make_repo(tmp_path / "repo") + base = await git(repo, "rev-parse", "HEAD") + (repo / "unreviewed-parent-change").write_text("parent work") + await commit_all(repo, "parent progress") + reviewed_parent = await git(repo, "rev-parse", "HEAD") + manager = WorktreeManager(repo) + info = await manager.create_worker( + "worker", base_commit="HEAD", dest_path=repo, dest_branch="main" + ) + (info.path / "feature").write_text("reviewed feature") + await commit_all(info.path, "worker feature") + candidate = await git(info.path, "rev-parse", "HEAD") + await git(repo, "reset", "--hard", base) + + async def never_run(cmd, cwd): + pytest.fail("stale parent review must be rejected before validation") + + with pytest.raises(WorktreeError, match="re-review"): + await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=reviewed_parent, + validation=["check"], + validation_runner=never_run, + ) + assert await git(repo, "rev-parse", "HEAD") == base + assert not (repo / "unreviewed-parent-change").exists() + + +@pytest.mark.parametrize("already_merged", [False, True]) +async def test_integrate_parent_advance_always_requires_rereview(tmp_path, already_merged): + repo, manager, info = await worker_repo(tmp_path) + reviewed_parent = await git(repo, "rev-parse", "HEAD") + (repo / "parent-progress").write_text("parent progress") + await commit_all(repo, "parent progress") + target = await git(repo, "rev-parse", "HEAD") + if already_merged: + await manager.reconcile(info.name) + candidate = await git(info.path, "rev-parse", "HEAD") + + async def never_run(cmd, cwd): + pytest.fail("a mismatched parent review must never reach validation") + + with pytest.raises(WorktreeError, match="re-review"): + await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=reviewed_parent, + validation=["check"], + validation_runner=never_run, + ) + assert await git(repo, "rev-parse", "HEAD") == target + assert await manager._ancestor(target, await git(info.path, "rev-parse", "HEAD")) + assert manager.read_sidecar(info.name)["integrated_head"] is None + assert (await integrate_head(manager, info)).merged + + +@pytest.mark.parametrize("reviewed_parent", ["HEAD", "a" * 39, "A" * 40, "g" * 64, None]) +async def test_integrate_requires_exact_parent_hash(tmp_path, reviewed_parent): + _repo, manager, info = await worker_repo(tmp_path) + with pytest.raises(WorktreeError, match="reviewed_parent_head must be the exact full"): + await manager.integrate( + info.name, + reviewed_head=await git(info.path, "rev-parse", "HEAD"), + reviewed_parent_head=reviewed_parent, + validation=["check"], + validation_runner=validate_git, + ) + + +@pytest.mark.parametrize("replace", ["worker", "parent"]) +async def test_validation_rejects_checkout_replacement_before_next_command(tmp_path, replace): + repo = await make_repo(tmp_path / "repo") + parent = tmp_path / "parent" + await git(repo, "worktree", "add", "-b", "parent", str(parent)) + manager = WorktreeManager(repo) + info = await manager.create_worker( + "worker", base_commit="HEAD", dest_path=parent, dest_branch="parent" + ) + (info.path / "feature").write_text("feature") + await commit_all(info.path, "worker feature") + candidate = await git(info.path, "rev-parse", "HEAD") + parent_head = await git(parent, "rev-parse", "HEAD") + sidecar = manager.read_sidecar(info.name) + path = info.path if replace == "worker" else parent + saved = tmp_path / "saved-checkout" + calls = [] + + async def validate(cmd, cwd): + calls.append(cmd) + assert cmd == "first", "must stop before validating a replacement checkout" + (path / "unfinished").write_text("retained work") + path.rename(saved) + path.mkdir() + (path / ".git").write_text((saved / ".git").read_text()) + await git(path, "restore", ".") + return ProcResult(0, "", "") + + with pytest.raises(WorktreeError, match="checkout identity changed"): + await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=parent_head, + validation=["first", "second"], + validation_runner=validate, + ) + assert calls == ["first"] + assert await git(parent, "rev-parse", "HEAD") == parent_head + assert (saved / "unfinished").read_text() == "retained work" + assert manager.read_sidecar(info.name) == sidecar + + +@pytest.mark.parametrize("change", ["dirty", "missing", "replaced", "branch", "detached"]) +async def test_integrate_refuses_changed_parent(tmp_path, change): + repo = await make_repo(tmp_path / "repo") + parent = tmp_path / "parent" + await git(repo, "worktree", "add", "-b", "parent", str(parent)) + manager = WorktreeManager(repo) + info = await manager.create_worker( + "worker", base_commit="HEAD", dest_path=parent, dest_branch="parent" + ) + (info.path / "feature").write_text("feature") + await commit_all(info.path, "feature") + target = await git(parent, "rev-parse", "HEAD") + if change == "dirty": + (parent / "file.txt").write_text("uncommitted") + elif change in {"missing", "replaced"}: + # Keep the original inode alive so replacement cannot reuse it. + parent.rename(tmp_path / "old-parent") + if change == "replaced": + parent.mkdir() + (parent / ".git").write_text((tmp_path / "old-parent" / ".git").read_text()) + await git(parent, "restore", ".") + elif change == "branch": + await git(parent, "switch", "-c", "different") + else: + await git(parent, "switch", "--detach") + with pytest.raises(WorktreeError): + await integrate_head(manager, info, reviewed_parent_head=target) + assert await git(repo, "rev-parse", "parent") == target + + +async def test_integrate_rejects_stale_review_and_requires_merge_rereview(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + old_head = await git(info.path, "rev-parse", "HEAD") + (info.path / "extra").write_text("extra") + await commit_all(info.path, "extra work") + with pytest.raises(WorktreeError, match="re-review"): + await manager.integrate( + info.name, + reviewed_head=old_head, + reviewed_parent_head=await git(repo, "rev-parse", "HEAD"), + validation=["check"], + validation_runner=validate_git, + ) + (repo / "parent").write_text("progress") + await commit_all(repo, "parent progress") + target = await git(repo, "rev-parse", "HEAD") + with pytest.raises(WorktreeError, match=r"destination merged.*re-review"): + await integrate_head(manager, info) + assert await git(repo, "rev-parse", "HEAD") == target + assert await git(info.path, "rev-parse", "HEAD^2") == target + candidate = await git(info.path, "rev-parse", "HEAD") + assert (await integrate_head(manager, info)).merged + assert await git(repo, "rev-parse", "HEAD") == candidate + assert manager.read_sidecar(info.name)["integrated_head"] == candidate + + +@pytest.mark.parametrize( + "change", + [ + "fail", + "timeout", + "worker-dirty", + "worker-head", + "target-head", + "target-dirty", + "target-switch", + "merge-state", + ], +) +async def test_integration_validation_cannot_change_candidate_or_target(tmp_path, change): + repo, manager, info = await worker_repo(tmp_path) + before = await git(repo, "rev-parse", "HEAD") + candidate = await git(info.path, "rev-parse", "HEAD") + calls = [] + + async def validate(cmd, cwd): + calls.append((cmd, cwd)) + if len(calls) > 1: + return ProcResult(0, "", "") + if change == "fail": + return ProcResult(1, "", "failed test") + if change == "timeout": + return ProcResult(0, "", "timeout", timed_out=True) + if change.startswith("worker"): + (cwd / "mutation").write_text("validation mutation") + if change == "worker-head": + await commit_all(cwd, "mutation") + elif change in {"target-head", "target-dirty"}: + (repo / "mutation").write_text("parent mutation") + if change == "target-head": + await commit_all(repo, "parent mutation") + elif change == "target-switch": + await git(repo, "switch", "-c", "other") + else: + # A real --no-commit merge, with a clean index but MERGE_HEAD set. + await git(repo, "switch", "-c", "other") + await git(repo, *_COMMIT, "--allow-empty", "-m", "diverge") + await git(repo, "switch", "main") + await git(cwd, "merge", "--no-ff", "--no-commit", "other") + return ProcResult(0, "", "") + + with pytest.raises(WorktreeError): + await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=before, + validation=["first", "second"], + validation_runner=validate, + ) + assert calls[0] == ("first", info.path) + assert len(calls) == 1 + assert not await manager._ancestor(candidate, await git(repo, "rev-parse", "main")) + if change != "target-head": + assert await git(repo, "rev-parse", "main") == before + assert manager.read_sidecar(info.name)["integrated_head"] is None + + +async def test_integration_empty_checks_require_explicit_human_gate(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + candidate = await git(info.path, "rev-parse", "HEAD") + parent_head = await git(repo, "rev-parse", "HEAD") + + async def never_run(cmd, cwd): + pytest.fail("empty validation must not call a runner") + + for checks in ([], [" "], [1]): + with pytest.raises(WorktreeError): + await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=parent_head, + validation=checks, + validation_runner=never_run, + ) + result = await manager.integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=parent_head, + validation=[], + validation_runner=never_run, + allow_unvalidated=True, + ) + assert result.merged + assert await git(repo, "rev-parse", "HEAD") == candidate + + +async def test_nested_integration_cleanup_checks_current_pinned_target(tmp_path): + repo, manager, parent = await worker_repo(tmp_path) + root_head = await git(repo, "rev-parse", "HEAD") + child = await manager.create_worker( + "child", base_commit=parent.branch, dest_path=parent.path, dest_branch=parent.branch + ) + (child.path / "nested").write_text("nested") + await commit_all(child.path, "nested work") + with pytest.raises(WorktreeError, match="not integrated"): + await manager.cleanup_worker(child.name) + candidate = await git(child.path, "rev-parse", "HEAD") + await integrate_head(manager, child) + assert await git(parent.path, "rev-parse", "HEAD") == candidate + assert await git(repo, "rev-parse", "HEAD") == root_head + (child.path / "after").write_text("extra commit after integration") + await commit_all(child.path, "after integration") + with pytest.raises(WorktreeError, match="not integrated"): + await manager.cleanup_worker(child.name) + assert child.path.exists() + await integrate_head(manager, child) + sidecar = manager.read_sidecar(child.name) + await manager.cleanup_worker(child.name) + assert not child.path.exists() + assert manager.read_sidecar(child.name) == sidecar + assert (parent.path / "after").read_text() == "extra commit after integration" + + +async def test_cleanup_dirty_requires_explicit_discard(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + await integrate_head(manager, info) + (info.path / "unfinished").write_text("dirty") + with pytest.raises(WorktreeError, match="uncommitted"): + await manager.cleanup_worker(info.name) + await manager.cleanup_worker(info.name, discard=True) + assert not info.path.exists() + assert manager.read_sidecar(info.name) is not None + assert (repo / "feature.txt").read_text() == "feature\n" + + +async def test_integration_lock_serializes_tasks_and_cancellation(tmp_path): + repo, _manager, info = await worker_repo(tmp_path) + candidate = await git(info.path, "rev-parse", "HEAD") + parent_head = await git(repo, "rev-parse", "HEAD") + entered = asyncio.Event() + hold = asyncio.Event() + + async def validate(cmd, cwd): + entered.set() + await hold.wait() + return await validate_git(cmd, cwd) + + async def integrate(): + return await WorktreeManager(repo).integrate( + info.name, + reviewed_head=candidate, + reviewed_parent_head=parent_head, + validation=["check"], + validation_runner=validate, + ) + + first = asyncio.create_task(integrate()) + await asyncio.wait_for(entered.wait(), 5) + (lock_path,) = (repo / ".git" / "lecode-integration-locks").iterdir() + inode = lock_path.stat().st_ino + pending = asyncio.create_task(integrate()) + await asyncio.sleep(0.15) + assert not pending.done() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + assert lock_path.stat().st_ino == inode + assert await git(repo, "rev-parse", "HEAD") != candidate + hold.set() + assert (await asyncio.wait_for(integrate(), 5)).merged + assert lock_path.stat().st_ino == inode + assert await git(repo, "rev-parse", "HEAD") == candidate + + +async def test_integration_lock_serializes_processes(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + candidate = await git(info.path, "rev-parse", "HEAD") + parent_head = await git(repo, "rev-parse", "HEAD") + script = """ +import asyncio, sys +from lecode.extras.worktree import WorktreeManager +from lecode.extras.proc import ProcResult +async def validate(cmd, cwd): + print("ready", flush=True) + await asyncio.to_thread(sys.stdin.readline) + return ProcResult(0, "", "") +asyncio.run(WorktreeManager(sys.argv[1]).integrate( + "worker", reviewed_head=sys.argv[2], reviewed_parent_head=sys.argv[3], + validation=["check"], validation_runner=validate)) +""" + proc = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + script, + str(repo), + candidate, + parent_head, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + task = None + try: + assert await asyncio.wait_for(proc.stdout.readline(), 10) == b"ready\n" + (lock_path,) = (repo / ".git" / "lecode-integration-locks").iterdir() + inode = lock_path.stat().st_ino + task = asyncio.create_task(integrate_head(manager, info)) + await asyncio.sleep(0.15) + assert not task.done() + assert await git(repo, "rev-parse", "HEAD") != candidate + proc.stdin.write(b"continue\n") + await proc.stdin.drain() + _, stderr = await asyncio.wait_for(proc.communicate(), 10) + assert proc.returncode == 0, stderr + with pytest.raises(WorktreeError, match="re-review"): + await asyncio.wait_for(task, 10) + assert (await integrate_head(manager, info)).merged + assert lock_path.stat().st_ino == inode + assert await git(repo, "rev-parse", "HEAD") == candidate + finally: + if proc.returncode is None: + proc.kill() + await proc.wait() + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +async def test_sibling_integrations_serialize_by_destination(tmp_path): + repo, manager, first = await worker_repo(tmp_path) + second = await manager.create_worker( + "second", base_commit="HEAD", dest_path=repo, dest_branch="main" + ) + (second.path / "second").write_text("second feature") + await commit_all(second.path, "second feature") + first_head = await git(first.path, "rev-parse", "HEAD") + second_head = await git(second.path, "rev-parse", "HEAD") + entered = asyncio.Event() + release = asyncio.Event() + + async def validate(cmd, cwd): + entered.set() + await release.wait() + return await validate_git(cmd, cwd) + + first_task = asyncio.create_task( + manager.integrate( + first.name, + reviewed_head=first_head, + reviewed_parent_head=await git(repo, "rev-parse", "HEAD"), + validation=["check"], + validation_runner=validate, + ) + ) + second_task = None + try: + await asyncio.wait_for(entered.wait(), 5) + second_task = asyncio.create_task(integrate_head(WorktreeManager(repo), second)) + await asyncio.sleep(0.15) + assert not second_task.done() + release.set() + assert (await asyncio.wait_for(first_task, 5)).merged + with pytest.raises(WorktreeError, match="re-review"): + await asyncio.wait_for(second_task, 5) + assert await git(repo, "rev-parse", "HEAD") == first_head + assert await git(second.path, "rev-parse", "HEAD^1") == second_head + assert await git(second.path, "rev-parse", "HEAD^2") == first_head + assert (await integrate_head(manager, second)).merged + assert (repo / "second").read_text() == "second feature" + assert (repo / "feature.txt").read_text() == "feature\n" + finally: + for task in (first_task, second_task): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +async def test_cleanup_rejects_switched_pinned_parent_even_if_other_branch_contains_worker( + tmp_path, +): + repo, manager, info = await worker_repo(tmp_path) + await integrate_head(manager, info) + await git(repo, "switch", "-c", "other") + with pytest.raises(WorktreeError, match="branch changed"): + await manager.cleanup_worker(info.name) + assert info.path.exists() + await git(repo, "switch", "main") + await manager.cleanup_worker(info.name) + assert not info.path.exists() + + +async def test_missing_switched_registration_is_not_recreated(tmp_path): + repo, manager, info = await worker_repo(tmp_path) + await git(info.path, "switch", "-c", "unrelated") + shutil.rmtree(info.path) + with pytest.raises(WorktreeError, match="registration has changed branch"): + await manager.reconcile(info.name, recreate=True) + assert "unrelated" in await git(repo, "worktree", "list", "--porcelain") + + # -- /worktree /wt-merge /wt-exit commands ------------------------------------------ From 26680d156107aad6aee97e051f874eb9bf099597 Mon Sep 17 00:00:00 2001 From: Mouhand-Kaddo Date: Thu, 17 Sep 2026 08:57:21 +0400 Subject: [PATCH 8/8] test: record hook payloads portably --- tests/test_tui_app.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_tui_app.py b/tests/test_tui_app.py index 505e0c9..c03a6b9 100644 --- a/tests/test_tui_app.py +++ b/tests/test_tui_app.py @@ -288,7 +288,7 @@ async def test_notes_hook_checks_composed_prompt_before_persistence(tmp_path, mo config = Config() config.hooks = { "UserPromptSubmit": [ - 'payload=$(cat); echo "$payload" >> prompts.jsonl; ' + 'payload=$(cat); printf "%s\\n" "$payload" >> prompts.jsonl; ' 'case "$payload" in *BLOCKED*) echo \'{"verdict":"deny"}\' ;; esac' ] } @@ -322,7 +322,7 @@ async def test_generated_prompts_are_guarded_at_each_submission( config = Config() config.hooks = { "UserPromptSubmit": [ - 'payload=$(cat); echo "$payload" >> prompts.jsonl; ' + 'payload=$(cat); printf "%s\\n" "$payload" >> prompts.jsonl; ' 'case "$payload" in *BLOCKED*) echo \'{"verdict":"deny"}\' ;; esac' ] }