whiletrue 0.1.0__py3-none-any.whl

whiletrue/claude_code/plugin/mcp_server.py

@@ -0,0 +1,98 @@
+"""Stdio MCP server exposing whiletrue's agent tools.
+
+Passed to Claude Code with `--mcp-config`. It holds no logic of its own: the tool list and every
+call are forwarded to the whiletrue process, which validates results against the step's output
+caps. Standard library only.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import urllib.request
+from collections.abc import Mapping
+from typing import IO, Any
+
+PROTOCOL_FALLBACK = "2024-11-05"
+
+
+def post_json(url: str, path: str, body: Any, token: str, session: str) -> Any:
+    request = urllib.request.Request(
+        url + path,
+        data=json.dumps(body).encode(),
+        headers={"X-WT-Token": token, "X-WT-Session": session, "Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(request, timeout=60) as response:
+        return json.loads(response.read() or b"{}")
+
+
+def _reply(message_id: Any, result: Any) -> dict[str, Any]:
+    return {"jsonrpc": "2.0", "id": message_id, "result": result}
+
+
+def handle(message: dict[str, Any], env: Mapping[str, str]) -> dict[str, Any] | None:
+    """Answers one JSON-RPC message (None for notifications)."""
+    method, message_id = message.get("method"), message.get("id")
+    if message_id is None:
+        return None
+    url, token, session = (
+        env.get("WT_ENDPOINT", ""),
+        env.get("WT_TOKEN", ""),
+        env.get("WT_SESSION", ""),
+    )
+    params: dict[str, Any] = message.get("params") or {}
+    if method == "initialize":
+        return _reply(
+            message_id,
+            {
+                "protocolVersion": params.get("protocolVersion", PROTOCOL_FALLBACK),
+                "capabilities": {"tools": {}},
+                "serverInfo": {"name": "whiletrue", "version": "0.1.0"},
+            },
+        )
+    if method == "ping":
+        return _reply(message_id, {})
+    if method == "tools/list":
+        try:
+            tools = post_json(url, "/tools/list", {}, token, session).get("tools", [])
+        except Exception:
+            tools = []
+        return _reply(message_id, {"tools": tools})
+    if method == "tools/call":
+        try:
+            answer = post_json(
+                url,
+                "/tools/call",
+                {"name": params.get("name"), "arguments": params.get("arguments") or {}},
+                token,
+                session,
+            )
+            text, is_error = str(answer.get("text", "")), bool(answer.get("is_error"))
+        except Exception as error:
+            text, is_error = f"whiletrue is unreachable: {error}", True
+        return _reply(
+            message_id, {"content": [{"type": "text", "text": text}], "isError": is_error}
+        )
+    return {
+        "jsonrpc": "2.0",
+        "id": message_id,
+        "error": {"code": -32601, "message": f"Method not found: {method}"},
+    }
+
+
+def serve(stdin: IO[str], stdout: IO[str], env: Mapping[str, str]) -> None:
+    for line in stdin:
+        if not line.strip():
+            continue
+        try:
+            message = json.loads(line)
+        except ValueError:
+            continue
+        response = handle(message, env) if isinstance(message, dict) else None  # pyright: ignore[reportUnknownArgumentType]
+        if response is not None:
+            stdout.write(json.dumps(response) + "\n")
+            stdout.flush()
+
+
+if __name__ == "__main__":  # pragma: no cover
+    serve(sys.stdin, sys.stdout, __import__("os").environ)

whiletrue/claude_code/session.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from whiletrue.claude_code.endpoint import ToolResult
+from whiletrue.claude_code.learnings import LearningsStore
+from whiletrue.claude_code.transcript import Price, Usage, read_usage
+from whiletrue.core import AgentLimits, CapValidation
+
+#: Why a session was parked rather than finished.
+ParkReason = Literal["needs_human", "budget_exceeded", "turns_exceeded", "wall_exceeded"]
+
+SUBMIT_TOOL = "whiletrue_submit"
+LEARN_TOOL = "whiletrue_learn"
+ASK_TOOL = "whiletrue_ask"
+CONTEXT_TOOL = "whiletrue_context"
+
+_NUDGE = (
+    "You have not submitted a result yet. Call the whiletrue_submit tool now with your result "
+    "(it is how this step finishes). If you are blocked, call whiletrue_ask with your question."
+)
+
+
+@dataclass
+class SessionOutcome:
+    #: `submitted`: a result arrived; `parked`: the step needs a person (see `reason`); `ended`:
+    #: the session closed without a result.
+    status: Literal["submitted", "parked", "ended"]
+    value: Any = None
+    #: The agent's last message.
+    text: str = ""
+    reason: ParkReason | None = None
+    detail: str | None = None
+    session_id: str | None = None
+    steered: bool = False
+    turns: int = 0
+    cost_usd: float = 0.0
+    steered_cost_usd: float = 0.0
+
+
+@dataclass(kw_only=True)
+class SessionSpec:
+    #: The step key the pane is registered under.
+    key: str
+    item: dict[str, Any] = field(default_factory=lambda: {})
+    #: Validates a submitted result against the step's output cap (None: accept anything).
+    validate: Callable[[Any], CapValidation] | None = None
+    result_cap: str | None = None
+    result_schema: dict[str, Any] | None = None
+    limits: AgentLimits | None = None
+    learnings: LearningsStore | None = None
+    #: Extra fields stored with every learning (run, item, step).
+    provenance: dict[str, Any] = field(default_factory=lambda: {})
+    #: Extra context returned by `whiletrue_context` (for example spec excerpts).
+    extra_context: dict[str, Any] = field(default_factory=lambda: {})
+    max_nudges: int = 3
+    prices: dict[str, Price] | None = None
+    clock: Callable[[], float] = time.time
+    on_event: Callable[[dict[str, Any]], None] | None = None
+    on_attention: Callable[[str], None] | None = None
+    on_attention_cleared: Callable[[], None] | None = None
+    #: Called once when a limit is hit, so the runner can end the session.
+    on_limit: Callable[[ParkReason], None] | None = None
+
+
+class StepSession:
+    """The in-process half of one interactive agent session.
+
+    The plugin's hooks and MCP tools land here (through the endpoint). It decides when the step
+    is done, parks it when it needs a person or hits a limit, and records a steer when a human
+    types into the pane.
+    """
+
+    def __init__(self, spec: SessionSpec) -> None:
+        self._spec = spec
+        self._started = spec.clock()
+        self.outcome: asyncio.Future[SessionOutcome] = asyncio.get_running_loop().create_future()
+        self.session_id: str | None = None
+        self.transcript_path: str | None = None
+        self.submitted = False
+        self._value: Any = None
+        self._question: str | None = None
+        self._last_message = ""
+        self._nudges = 0
+        self._prompts = 0
+        self._steered_at: float | None = None
+        self.tool_calls = 0
+        self.needs_attention = False
+        self.limit_hit: ParkReason | None = None
+
+    @property
+    def steered(self) -> bool:
+        return self._steered_at is not None
+
+    def _event(self, kind: str, **data: Any) -> None:
+        if self._spec.on_event is not None:
+            self._spec.on_event({"kind": kind, **data})
+
+    # ---- tools -----------------------------------------------------------------------------
+
+    def tools(self) -> list[dict[str, Any]]:
+        spec = self._spec
+        cap = f" The result must match the {spec.result_cap} schema." if spec.result_cap else ""
+        result_schema: dict[str, Any] = spec.result_schema if spec.result_schema else {}
+        return [
+            {
+                "name": SUBMIT_TOOL,
+                "description": "Finish this step by submitting your result. Call it when the "
+                f"work is complete.{cap} Invalid results are returned with the errors.",
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {"result": result_schema},
+                    "required": ["result"],
+                },
+            },
+            {
+                "name": LEARN_TOOL,
+                "description": "Record a short, reusable lesson about this project.",
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {"note": {"type": "string"}},
+                    "required": ["note"],
+                },
+            },
+            {
+                "name": ASK_TOOL,
+                "description": "Ask the human a question and stop. The step parks until answered.",
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {"question": {"type": "string"}},
+                    "required": ["question"],
+                },
+            },
+            {
+                "name": CONTEXT_TOOL,
+                "description": "Get this item's details and recent project learnings.",
+                "inputSchema": {"type": "object", "properties": {}},
+            },
+        ]
+
+    async def on_tool(self, name: str, args: dict[str, Any]) -> ToolResult:
+        if name == SUBMIT_TOOL:
+            return self._submit(args)
+        if name == LEARN_TOOL:
+            return self._learn(args)
+        if name == ASK_TOOL:
+            question = args.get("question")
+            if not isinstance(question, str) or not question.strip():
+                return ToolResult("`question` is required", is_error=True)
+            self._question = question
+            self._event("ask", question=question)
+            return ToolResult("Recorded. Stop now; a person will answer and resume the work.")
+        if name == CONTEXT_TOOL:
+            learnings = self._spec.learnings.recent(20) if self._spec.learnings else []
+            body = {"item": self._spec.item, "learnings": learnings, **self._spec.extra_context}
+            return ToolResult(json.dumps(body))
+        return ToolResult(f'Unknown tool "{name}"', is_error=True)
+
+    def _submit(self, args: dict[str, Any]) -> ToolResult:
+        if "result" not in args:
+            return ToolResult("`result` is required", is_error=True)
+        value: Any = args["result"]
+        if self._spec.validate is not None:
+            check = self._spec.validate(value)
+            if not check.ok:
+                self._event("submit_rejected", errors=check.errors)
+                return ToolResult(
+                    "Result rejected, fix these and call whiletrue_submit again:\n"
+                    + "\n".join(f"- {e}" for e in check.errors),
+                    is_error=True,
+                )
+            value = check.value
+        self._value = value
+        self.submitted = True
+        self._event("submit")
+        return ToolResult("Result accepted.")
+
+    def _learn(self, args: dict[str, Any]) -> ToolResult:
+        note = args.get("note")
+        if not isinstance(note, str) or not note.strip():
+            return ToolResult("`note` is required", is_error=True)
+        if self._spec.learnings is None:
+            return ToolResult("Learnings are not enabled for this run", is_error=True)
+        self._spec.learnings.append(note.strip(), **self._spec.provenance)
+        self._event("learn")
+        return ToolResult("Noted.")
+
+    # ---- hooks -----------------------------------------------------------------------------
+
+    async def on_hook(self, event: str, payload: dict[str, Any]) -> dict[str, Any] | None:
+        match event:
+            case "SessionStart":
+                self.session_id = payload.get("session_id") or self.session_id
+                self.transcript_path = payload.get("transcript_path") or self.transcript_path
+                self._event("session_start", session_id=self.session_id)
+            case "UserPromptSubmit":
+                self._on_prompt(str(payload.get("prompt", "")))
+            case "PreToolUse":
+                return self._on_pre_tool(payload)
+            case "PostToolUse":
+                self.tool_calls += 1
+                self._clear_attention()
+                self._event("tool_use", name=payload.get("tool_name", ""))
+            case "Notification":
+                message = str(payload.get("message", ""))
+                self.needs_attention = True
+                self._event("attention", message=message)
+                if self._spec.on_attention is not None:
+                    self._spec.on_attention(message)
+            case "Stop":
+                return self._on_stop(payload)
+            case "SessionEnd":
+                self._on_end()
+            case _:
+                pass
+        return None
+
+    def _clear_attention(self) -> None:
+        if self.needs_attention:
+            self.needs_attention = False
+            if self._spec.on_attention_cleared is not None:
+                self._spec.on_attention_cleared()
+
+    def _on_prompt(self, prompt: str) -> None:
+        self._prompts += 1
+        self._clear_attention()
+        if self._prompts > 1 and self._steered_at is None:
+            # The prompt we launched with is the first; anything later was typed by a person.
+            self._steered_at = self._spec.clock()
+        if self._prompts > 1:
+            self._event("steer", prompt=prompt)
+
+    def _usage(self) -> Usage:
+        return read_usage(self.transcript_path, self._spec.prices, steered_since=self._steered_at)
+
+    def _check_limits(self) -> ParkReason | None:
+        limits = self._spec.limits
+        if limits is None or self.steered:
+            return None  # a person driving the session lifts every limit
+        if (
+            limits.wall_ms is not None
+            and (self._spec.clock() - self._started) * 1000 >= limits.wall_ms
+        ):
+            return "wall_exceeded"
+        if limits.budget_usd is None and limits.turns is None:
+            return None
+        usage = self._usage()
+        if limits.budget_usd is not None and usage.cost_usd >= limits.budget_usd:
+            return "budget_exceeded"
+        if limits.turns is not None and usage.turns > limits.turns:
+            return "turns_exceeded"
+        return None
+
+    def _on_pre_tool(self, payload: dict[str, Any]) -> dict[str, Any] | None:
+        if self.limit_hit is None:
+            hit = self._check_limits()
+            if hit is None:
+                return None
+            self.limit_hit = hit
+            self._event("limit", reason=hit)
+            if self._spec.on_limit is not None:
+                self._spec.on_limit(hit)
+        elif self.steered:
+            return None
+        label = {
+            "budget_exceeded": "the budget estimate was exceeded",
+            "turns_exceeded": "the turn limit was reached",
+            "wall_exceeded": "the wall-time limit was reached",
+        }.get(self.limit_hit, "a limit was reached")
+        return {
+            "hookSpecificOutput": {
+                "hookEventName": "PreToolUse",
+                "permissionDecision": "deny",
+                "permissionDecisionReason": f"whiletrue: {label}; this step is being parked.",
+            }
+        }
+
+    def _on_stop(self, payload: dict[str, Any]) -> dict[str, Any] | None:
+        if self.outcome.done():
+            return None
+        self._last_message = str(payload.get("last_assistant_message", self._last_message))
+        if self.limit_hit is not None and not self.steered:
+            self._finish("parked", reason=self.limit_hit, detail=f"limit hit: {self.limit_hit}")
+        elif self._question is not None:
+            self._finish("parked", reason="needs_human", detail=self._question)
+        elif self.submitted:
+            self._finish("submitted")
+        elif self.steered:
+            return None  # a person is driving: waiting for them is not a failure
+        elif self._nudges < self._spec.max_nudges:
+            self._nudges += 1
+            self._event("nudge", count=self._nudges)
+            return {"decision": "block", "reason": _NUDGE}
+        else:
+            self._finish(
+                "parked",
+                reason="needs_human",
+                detail="the agent stopped without submitting a result after "
+                f"{self._nudges} reminder(s)",
+            )
+        return None
+
+    def park(self, reason: ParkReason, detail: str) -> None:
+        """Parks the step from outside (the runner's wall-time timer, a dead pane, ...)."""
+        if self.outcome.done():
+            return
+        self.limit_hit = reason if reason != "needs_human" else self.limit_hit
+        self._finish("parked", reason=reason, detail=detail)
+
+    def _on_end(self) -> None:
+        if self.outcome.done():
+            return
+        if self.submitted:
+            self._finish("submitted")
+        else:
+            self._finish("ended", detail="the session ended without submitting a result")
+
+    def _finish(
+        self,
+        status: Literal["submitted", "parked", "ended"],
+        *,
+        reason: ParkReason | None = None,
+        detail: str | None = None,
+    ) -> None:
+        usage = self._usage()
+        self._event("done", status=status)
+        self.outcome.set_result(
+            SessionOutcome(
+                status=status,
+                value=self._value,
+                text=self._last_message,
+                reason=reason,
+                detail=detail,
+                session_id=self.session_id,
+                steered=self.steered,
+                turns=usage.turns,
+                cost_usd=usage.cost_usd,
+                steered_cost_usd=usage.steered_cost_usd,
+            )
+        )

whiletrue/claude_code/transcript.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, cast
+
+
+@dataclass(frozen=True)
+class Price:
+    """USD per million tokens."""
+
+    input: float
+    output: float
+    cache_read: float
+    cache_write: float
+
+
+#: Rough list prices by model family, matched by substring of the model id. These drive the
+#: *estimate* that enforces budgets in interactive sessions; they are not billing data.
+DEFAULT_PRICES: dict[str, Price] = {
+    "opus": Price(input=15.0, output=75.0, cache_read=1.5, cache_write=18.75),
+    "sonnet": Price(input=3.0, output=15.0, cache_read=0.3, cache_write=3.75),
+    "haiku": Price(input=1.0, output=5.0, cache_read=0.1, cache_write=1.25),
+}
+
+
+@dataclass
+class Usage:
+    #: Assistant messages (model calls) so far.
+    turns: int = 0
+    cost_usd: float = 0.0
+    #: The part of `cost_usd` spent after a human steered the session.
+    steered_cost_usd: float = 0.0
+
+    @property
+    def agent_cost_usd(self) -> float:
+        return self.cost_usd - self.steered_cost_usd
+
+
+def _price(model: str, prices: dict[str, Price]) -> Price:
+    for family, price in prices.items():
+        if family in model:
+            return price
+    # Unknown model: assume the dearest family so a budget is never silently under-counted.
+    return max(prices.values(), key=lambda p: p.output)
+
+
+def _epoch(stamp: Any) -> float | None:
+    if not isinstance(stamp, str):
+        return None
+    try:
+        return datetime.fromisoformat(stamp.replace("Z", "+00:00")).timestamp()
+    except ValueError:
+        return None
+
+
+def _loads(raw: str) -> Any:
+    try:
+        return json.loads(raw)
+    except ValueError:
+        return None
+
+
+def _obj(value: Any) -> dict[str, Any] | None:
+    return cast(dict[str, Any], value) if isinstance(value, dict) else None
+
+
+def _tokens(usage: dict[str, Any], key: str) -> int:
+    value = usage.get(key)
+    return value if isinstance(value, int) else 0
+
+
+def read_usage(
+    path: str | None,
+    prices: dict[str, Price] | None = None,
+    *,
+    steered_since: float | None = None,
+) -> Usage:
+    """Estimates turns and spend from a Claude Code session transcript (JSONL).
+
+    A message is logged once per content block, so messages are deduplicated by id. Spend after
+    `steered_since` (epoch seconds) is reported separately.
+    """
+    table = prices if prices is not None else DEFAULT_PRICES
+    usage = Usage()
+    if path is None:
+        return usage
+    try:
+        text = Path(path).read_text(encoding="utf-8")
+    except OSError:
+        return usage
+    seen: set[str] = set()
+    for raw in text.splitlines():
+        entry = _obj(_loads(raw))
+        if entry is None or entry.get("type") != "assistant":
+            continue
+        message = _obj(entry.get("message"))
+        if message is None:
+            continue
+        message_id = message.get("id")
+        if isinstance(message_id, str):
+            if message_id in seen:
+                continue
+            seen.add(message_id)
+        usage.turns += 1
+        counts = _obj(message.get("usage"))
+        if counts is None:
+            continue
+        model = message.get("model")
+        price = _price(model if isinstance(model, str) else "", table)
+        cost = (
+            _tokens(counts, "input_tokens") * price.input
+            + _tokens(counts, "output_tokens") * price.output
+            + _tokens(counts, "cache_read_input_tokens") * price.cache_read
+            + _tokens(counts, "cache_creation_input_tokens") * price.cache_write
+        ) / 1_000_000
+        usage.cost_usd += cost
+        when = _epoch(entry.get("timestamp"))
+        if steered_since is not None and when is not None and when >= steered_since:
+            usage.steered_cost_usd += cost
+    return usage

whiletrue/cli/__init__.py

@@ -0,0 +1,15 @@
+"""The `whiletrue` command-line interface."""
+
+from whiletrue.cli.app import entrypoint, main
+from whiletrue.cli.common import EXIT, Io
+from whiletrue.cli.registry import RegistryOptions, create_default_registry, load_plugins
+
+__all__ = [
+    "EXIT",
+    "Io",
+    "RegistryOptions",
+    "create_default_registry",
+    "entrypoint",
+    "load_plugins",
+    "main",
+]

whiletrue/cli/__main__.py

@@ -0,0 +1,3 @@
+from whiletrue.cli.app import entrypoint
+
+entrypoint()