pythinker-code 2.0.0__py3-none-any.whl

pythinker_code/subagents/runner.py

@@ -0,0 +1,428 @@
+from __future__ import annotations
+
+import asyncio
+import uuid
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pythinker_core.chat_provider import APIStatusError, ChatProviderError
+from pythinker_core.tooling import ToolError, ToolOk, ToolReturnValue
+
+from pythinker_code.approval_runtime import (
+    ApprovalSource,
+    reset_current_approval_source,
+    set_current_approval_source,
+)
+from pythinker_code.soul import MaxStepsReached, RunCancelled, UILoopFn, get_wire_or_none, run_soul
+from pythinker_code.soul.pythinkersoul import PythinkerSoul
+from pythinker_code.soul.toolset import get_current_tool_call_or_none
+from pythinker_code.subagents.builder import SubagentBuilder
+from pythinker_code.subagents.core import SubagentRunSpec, prepare_soul
+from pythinker_code.subagents.models import AgentInstanceRecord, AgentLaunchSpec
+from pythinker_code.subagents.output import SubagentOutputWriter
+from pythinker_code.subagents.store import SubagentStore
+from pythinker_code.utils.logging import logger
+from pythinker_code.wire import Wire
+from pythinker_code.wire.file import WireFile
+from pythinker_code.wire.types import (
+    ApprovalRequest,
+    ApprovalResponse,
+    HookRequest,
+    QuestionRequest,
+    SubagentEvent,
+    ToolCallRequest,
+)
+
+if TYPE_CHECKING:
+    from pythinker_code.soul.agent import Runtime
+
+SUMMARY_MIN_LENGTH = 200
+SUMMARY_CONTINUATION_ATTEMPTS = 1
+SUMMARY_CONTINUATION_PROMPT = """
+Your previous response was too brief. Please provide a more comprehensive summary that includes:
+
+1. Specific technical details and implementations
+2. Detailed findings and analysis
+3. All important information that the parent agent should know
+""".strip()
+
+
+# ---------------------------------------------------------------------------
+# Shared result types and execution helpers (used by both foreground and
+# background runners).
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SoulRunFailure:
+    """Describes why a soul run did not produce a usable result."""
+
+    message: str
+    brief: str
+
+
+async def run_soul_checked(
+    soul: PythinkerSoul,
+    prompt: str,
+    ui_loop_fn: UILoopFn,
+    wire_path: Path,
+    phase: str,
+) -> SoulRunFailure | None:
+    """Run a single soul turn and validate the result.
+
+    Returns a ``SoulRunFailure`` if the run failed or produced an invalid
+    result, or ``None`` on success.  Most exceptions (``MaxStepsReached``,
+    ``ChatProviderError``, generic ``Exception``) are converted to failures.
+    Only ``CancelledError`` and ``RunCancelled`` are re-raised.
+    """
+    try:
+        await run_soul(
+            soul,
+            prompt,
+            ui_loop_fn,
+            asyncio.Event(),
+            wire_file=WireFile(wire_path),
+            runtime=soul.runtime,
+        )
+    except MaxStepsReached as exc:
+        logger.warning(
+            "Subagent max steps reached ({n_steps}) when {phase}",
+            n_steps=exc.n_steps,
+            phase=phase,
+        )
+        return SoulRunFailure(
+            message=(
+                f"Max steps {exc.n_steps} reached when {phase}. "
+                "Please try splitting the task into smaller subtasks."
+            ),
+            brief="Max steps reached",
+        )
+    except RunCancelled:
+        raise
+    except asyncio.CancelledError:
+        raise
+    except APIStatusError as exc:
+        logger.warning(
+            "Subagent LLM API error (HTTP {status_code}) when {phase}: {error}",
+            status_code=exc.status_code,
+            phase=phase,
+            error=exc,
+        )
+        return SoulRunFailure(
+            message=f"LLM API error (HTTP {exc.status_code}) when {phase}: {exc}",
+            brief=f"API error ({exc.status_code})",
+        )
+    except ChatProviderError as exc:
+        logger.warning(
+            "Subagent LLM provider error when {phase}: {error}",
+            phase=phase,
+            error=exc,
+        )
+        return SoulRunFailure(
+            message=f"LLM provider error when {phase}: {exc}",
+            brief="LLM provider error",
+        )
+    except Exception as exc:
+        logger.exception("Subagent soul run failed when {phase}", phase=phase)
+        return SoulRunFailure(
+            message=f"Unexpected error when {phase}: {exc}",
+            brief="Agent run error",
+        )
+
+    context = soul.context
+    if not context.history or context.history[-1].role != "assistant":
+        return SoulRunFailure(
+            message="The agent did not produce a valid assistant response.",
+            brief="Invalid agent result",
+        )
+    return None
+
+
+async def run_with_summary_continuation(
+    soul: PythinkerSoul,
+    prompt: str,
+    ui_loop_fn: UILoopFn,
+    wire_path: Path,
+) -> tuple[str | None, SoulRunFailure | None]:
+    """Run soul, then optionally extend the summary if it is too short.
+
+    Returns ``(final_response, failure)``.  On success ``failure`` is
+    ``None`` and ``final_response`` contains the agent's output text.
+    On failure ``final_response`` is ``None``.
+    """
+    failure = await run_soul_checked(soul, prompt, ui_loop_fn, wire_path, "running agent")
+    if failure is not None:
+        return None, failure
+
+    final_response = soul.context.history[-1].extract_text(sep="\n")
+    remaining = SUMMARY_CONTINUATION_ATTEMPTS
+    while remaining > 0 and len(final_response) < SUMMARY_MIN_LENGTH:
+        remaining -= 1
+        failure = await run_soul_checked(
+            soul,
+            SUMMARY_CONTINUATION_PROMPT,
+            ui_loop_fn,
+            wire_path,
+            "continuing the agent summary",
+        )
+        if failure is not None:
+            return None, failure
+        final_response = soul.context.history[-1].extract_text(sep="\n")
+
+    return final_response, None
+
+
+# ---------------------------------------------------------------------------
+# Foreground runner
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ForegroundRunRequest:
+    description: str
+    prompt: str
+    requested_type: str
+    model: str | None
+    resume: str | None
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class PreparedInstance:
+    record: AgentInstanceRecord
+    actual_type: str
+    resumed: bool
+
+
+class ForegroundSubagentRunner:
+    def __init__(self, runtime: Runtime):
+        self._runtime = runtime
+        assert runtime.subagent_store is not None
+        self._store: SubagentStore = runtime.subagent_store
+        self._builder = SubagentBuilder(runtime)
+
+    async def run(self, req: ForegroundRunRequest) -> ToolReturnValue:
+        prepared = await self._prepare_instance(req)
+        agent_id = prepared.record.agent_id
+        actual_type = prepared.actual_type
+        resumed = prepared.resumed
+
+        type_def = self._runtime.labor_market.require_builtin_type(actual_type)
+        launch_spec = prepared.record.launch_spec
+        if req.model is not None:
+            launch_spec = replace(
+                launch_spec,
+                model_override=req.model,
+                effective_model=req.model,
+            )
+
+        output_writer = SubagentOutputWriter(self._store.output_path(agent_id))
+        output_writer.stage("runner_started")
+
+        spec = SubagentRunSpec(
+            agent_id=agent_id,
+            type_def=type_def,
+            launch_spec=launch_spec,
+            prompt=req.prompt,
+            resumed=resumed,
+        )
+        soul, prompt = await prepare_soul(
+            spec,
+            self._runtime,
+            self._builder,
+            self._store,
+            on_stage=output_writer.stage,
+        )
+
+        self._store.update_instance(
+            agent_id,
+            status="running_foreground",
+            description=req.description.strip(),
+        )
+        approval_source: ApprovalSource | None = None
+        approval_source_token = None
+        try:
+            # Propagate hook engine from parent runtime to subagent soul
+            if self._runtime.hook_engine is not None:
+                soul.set_hook_engine(self._runtime.hook_engine)
+            tool_call = get_current_tool_call_or_none()
+            ui_loop_fn = self._make_ui_loop_fn(
+                parent_tool_call_id=tool_call.id if tool_call is not None else None,
+                agent_id=agent_id,
+                subagent_type=actual_type,
+                output_writer=output_writer,
+            )
+
+            # Use a single stable ApprovalSource for the entire run (including summary
+            # continuation).  This ensures cancel_by_source can reliably cancel all
+            # pending approval requests belonging to this foreground subagent execution.
+            approval_source = ApprovalSource(
+                kind="foreground_turn",
+                id=uuid.uuid4().hex,
+                agent_id=agent_id,
+                subagent_type=actual_type,
+            )
+            approval_source_token = set_current_approval_source(approval_source)
+
+            # --- SubagentStart hook ---
+            hook_engine = soul.hook_engine
+            from pythinker_code.hooks import events as hook_events
+
+            await hook_engine.trigger(
+                "SubagentStart",
+                matcher_value=actual_type,
+                input_data=hook_events.subagent_start(
+                    session_id=self._runtime.session.id,
+                    cwd=str(Path.cwd()),
+                    agent_name=actual_type,
+                    prompt=req.prompt[:500],
+                ),
+            )
+
+            output_writer.stage("run_soul_start")
+            final_response, failure = await run_with_summary_continuation(
+                soul,
+                prompt,
+                ui_loop_fn,
+                self._store.wire_path(agent_id),
+            )
+            if failure is not None:
+                self._store.update_instance(agent_id, status="failed")
+                output_writer.stage(f"failed: {failure.brief}")
+                return ToolError(message=failure.message, brief=failure.brief)
+            output_writer.stage("run_soul_finished")
+
+            # --- SubagentStop hook ---
+            # fire_and_forget_trigger keeps a strong reference to the task on
+            # the hook engine so it cannot be garbage-collected before it
+            # finishes. Without that, asyncio's WeakSet bookkeeping would let
+            # GC reap the still-pending task and trigger the broken
+            # "Exception None" loop-handler path.
+            hook_engine.fire_and_forget_trigger(
+                "SubagentStop",
+                matcher_value=actual_type,
+                input_data=hook_events.subagent_stop(
+                    session_id=self._runtime.session.id,
+                    cwd=str(Path.cwd()),
+                    agent_name=actual_type,
+                    response=(final_response or "")[:500],
+                ),
+            )
+        except asyncio.CancelledError:
+            self._store.update_instance(agent_id, status="killed")
+            output_writer.stage("cancelled")
+            raise
+        except RunCancelled as exc:
+            self._store.update_instance(agent_id, status="killed")
+            output_writer.stage("cancelled")
+            raise RunCancelled("Subagent run was cancelled.") from exc
+        except Exception:
+            self._store.update_instance(agent_id, status="failed")
+            output_writer.stage("failed_exception")
+            raise
+        finally:
+            if approval_source_token is not None:
+                reset_current_approval_source(approval_source_token)
+            if approval_source is not None and self._runtime.approval_runtime is not None:
+                self._runtime.approval_runtime.cancel_by_source(
+                    approval_source.kind, approval_source.id
+                )
+
+        if final_response is None:
+            self._store.update_instance(agent_id, status="failed")
+            output_writer.stage("failed: empty output")
+            return ToolError(
+                message="Agent completed but produced no output.",
+                brief="Empty agent output",
+            )
+        self._store.update_instance(agent_id, status="idle")
+        output_writer.summary(final_response)
+        lines = [
+            f"agent_id: {agent_id}",
+            "resumed: true" if resumed else "resumed: false",
+        ]
+        if resumed and req.requested_type and req.requested_type != actual_type:
+            lines.append(f"requested_subagent_type: {req.requested_type}")
+        lines.extend(
+            [
+                f"actual_subagent_type: {actual_type}",
+                "status: completed",
+                "",
+                "[summary]",
+                final_response,
+            ]
+        )
+        return ToolOk(output="\n".join(lines))
+
+    async def _prepare_instance(self, req: ForegroundRunRequest) -> PreparedInstance:
+        if req.resume:
+            record = self._store.require_instance(req.resume)
+            if record.status in {"running_foreground", "running_background"}:
+                raise RuntimeError(
+                    f"Agent instance {record.agent_id} is still {record.status} and cannot be "
+                    "resumed concurrently."
+                )
+            return PreparedInstance(
+                record=record,
+                actual_type=record.subagent_type,
+                resumed=True,
+            )
+
+        actual_type = req.requested_type or "coder"
+        type_def = self._runtime.labor_market.require_builtin_type(actual_type)
+        agent_id = f"a{uuid.uuid4().hex[:8]}"
+        record = self._store.create_instance(
+            agent_id=agent_id,
+            description=req.description.strip(),
+            launch_spec=AgentLaunchSpec(
+                agent_id=agent_id,
+                subagent_type=actual_type,
+                model_override=req.model,
+                effective_model=req.model or type_def.default_model,
+            ),
+        )
+        from pythinker_code.telemetry import track
+
+        track("subagent_created")
+        return PreparedInstance(
+            record=record,
+            actual_type=actual_type,
+            resumed=False,
+        )
+
+    @staticmethod
+    def _make_ui_loop_fn(
+        *,
+        parent_tool_call_id: str | None,
+        agent_id: str,
+        subagent_type: str,
+        output_writer: SubagentOutputWriter,
+    ):
+        super_wire = get_wire_or_none()
+
+        async def _ui_loop_fn(wire: Wire) -> None:
+            wire_ui = wire.ui_side(merge=True)
+            while True:
+                msg = await wire_ui.receive()
+                # Always write to output file regardless of wire availability.
+                output_writer.write_wire_message(msg)
+                if super_wire is None or parent_tool_call_id is None:
+                    continue
+                if isinstance(
+                    msg,
+                    ApprovalRequest | ApprovalResponse | ToolCallRequest | QuestionRequest,
+                ):
+                    super_wire.soul_side.send(msg)
+                    continue
+                if isinstance(msg, HookRequest):
+                    continue
+                super_wire.soul_side.send(
+                    SubagentEvent(
+                        parent_tool_call_id=parent_tool_call_id,
+                        agent_id=agent_id,
+                        subagent_type=subagent_type,
+                        event=msg,
+                    )
+                )
+
+        return _ui_loop_fn

pythinker_code/subagents/store.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import json
+import shutil
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, cast
+
+from pydantic import BaseModel, ValidationError
+
+from pythinker_code.session import Session
+from pythinker_code.subagents.models import AgentInstanceRecord, AgentLaunchSpec, SubagentStatus
+from pythinker_code.utils.io import atomic_json_write
+from pythinker_code.utils.logging import logger
+
+
+class _AgentLaunchSpecPayload(BaseModel):
+    agent_id: str
+    subagent_type: str
+    model_override: str | None
+    effective_model: str | None
+    created_at: float
+
+
+class _AgentInstanceRecordPayload(BaseModel):
+    agent_id: str
+    subagent_type: str
+    status: str
+    description: str
+    created_at: float
+    updated_at: float
+    last_task_id: str | None = None
+    launch_spec: _AgentLaunchSpecPayload
+
+
+_VALID_SUBAGENT_STATUSES = cast(
+    tuple[str, ...],
+    ("idle", "running_foreground", "running_background", "completed", "failed", "killed"),
+)
+
+
+def _record_from_dict(data: dict[str, Any]) -> AgentInstanceRecord:
+    payload = _AgentInstanceRecordPayload.model_validate(data)
+    if payload.status not in _VALID_SUBAGENT_STATUSES:
+        raise ValueError(f"Invalid subagent status: {payload.status!r}")
+    return AgentInstanceRecord(
+        agent_id=payload.agent_id,
+        subagent_type=payload.subagent_type,
+        status=cast(SubagentStatus, payload.status),
+        description=payload.description,
+        created_at=payload.created_at,
+        updated_at=payload.updated_at,
+        last_task_id=payload.last_task_id,
+        launch_spec=AgentLaunchSpec(
+            agent_id=payload.launch_spec.agent_id,
+            subagent_type=payload.launch_spec.subagent_type,
+            model_override=payload.launch_spec.model_override,
+            effective_model=payload.launch_spec.effective_model,
+            created_at=payload.launch_spec.created_at,
+        ),
+    )
+
+
+class SubagentStore:
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    @property
+    def root(self) -> Path:
+        return self._session.dir / "subagents"
+
+    def instance_dir(self, agent_id: str, *, create: bool = False) -> Path:
+        path = self.root / agent_id
+        if create:
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def context_path(self, agent_id: str) -> Path:
+        return self.instance_dir(agent_id) / "context.jsonl"
+
+    def wire_path(self, agent_id: str) -> Path:
+        return self.instance_dir(agent_id) / "wire.jsonl"
+
+    def meta_path(self, agent_id: str) -> Path:
+        return self.instance_dir(agent_id) / "meta.json"
+
+    def prompt_path(self, agent_id: str) -> Path:
+        return self.instance_dir(agent_id) / "prompt.txt"
+
+    def output_path(self, agent_id: str) -> Path:
+        return self.instance_dir(agent_id) / "output"
+
+    def create_instance(
+        self,
+        *,
+        agent_id: str,
+        description: str,
+        launch_spec: AgentLaunchSpec,
+    ) -> AgentInstanceRecord:
+        self._initialize_instance_files(agent_id)
+        record = AgentInstanceRecord(
+            agent_id=agent_id,
+            subagent_type=launch_spec.subagent_type,
+            status="idle",
+            description=description,
+            created_at=launch_spec.created_at,
+            updated_at=launch_spec.created_at,
+            last_task_id=None,
+            launch_spec=launch_spec,
+        )
+        self.write_instance(record)
+        return record
+
+    def write_instance(self, record: AgentInstanceRecord) -> None:
+        instance_dir = self.instance_dir(record.agent_id)
+        atomic_json_write(asdict(record), instance_dir / "meta.json")
+
+    def _initialize_instance_files(self, agent_id: str) -> None:
+        instance_dir = self.instance_dir(agent_id, create=True)
+        (instance_dir / "context.jsonl").touch(exist_ok=True)
+        (instance_dir / "wire.jsonl").touch(exist_ok=True)
+        (instance_dir / "prompt.txt").touch(exist_ok=True)
+        (instance_dir / "output").touch(exist_ok=True)
+
+    def get_instance(self, agent_id: str) -> AgentInstanceRecord | None:
+        meta = self.meta_path(agent_id)
+        if not meta.exists():
+            return None
+        return _load_instance_record(meta)
+
+    def require_instance(self, agent_id: str) -> AgentInstanceRecord:
+        record = self.get_instance(agent_id)
+        if record is None:
+            raise FileNotFoundError(f"Subagent instance not found: {agent_id}")
+        return record
+
+    def update_instance(
+        self,
+        agent_id: str,
+        *,
+        status: SubagentStatus | None = None,
+        description: str | None = None,
+        last_task_id: str | None | object = ...,
+    ) -> AgentInstanceRecord:
+        import time
+
+        current = self.require_instance(agent_id)
+        record = AgentInstanceRecord(
+            agent_id=current.agent_id,
+            subagent_type=current.subagent_type,
+            status=current.status if status is None else status,
+            description=current.description if description is None else description,
+            created_at=current.created_at,
+            updated_at=time.time(),
+            last_task_id=(
+                current.last_task_id if last_task_id is ... else cast(str | None, last_task_id)
+            ),
+            launch_spec=current.launch_spec,
+        )
+        self.write_instance(record)
+        return record
+
+    def list_instances(self) -> list[AgentInstanceRecord]:
+        records: list[AgentInstanceRecord] = []
+        if not self.root.exists():
+            return records
+        for path in self.root.iterdir():
+            if not path.is_dir():
+                continue
+            meta = path / "meta.json"
+            if not meta.exists():
+                continue
+            record = _load_instance_record(meta)
+            if record is None:
+                continue
+            records.append(record)
+        records.sort(key=lambda record: record.updated_at, reverse=True)
+        return records
+
+    def delete_instance(self, agent_id: str) -> None:
+        instance_dir = self.instance_dir(agent_id)
+        if not instance_dir.exists():
+            return
+        shutil.rmtree(instance_dir)
+
+
+def _load_instance_record(meta_path: Path) -> AgentInstanceRecord | None:
+    try:
+        return _record_from_dict(json.loads(meta_path.read_text(encoding="utf-8")))
+    except (OSError, json.JSONDecodeError, ValidationError, TypeError, ValueError) as exc:
+        logger.warning(
+            "Skipping invalid subagent metadata {path}: {error}",
+            path=meta_path,
+            error=exc,
+        )
+        return None