pcode 0.1.0__py3-none-any.whl

pcode/stream_events.py

@@ -0,0 +1,437 @@
+"""Translate one agent run's Pydantic AI events into pcode's runtime events.
+
+`agent.run_stream_events` yields typed upstream events: model parts streaming
+in, tool calls and results, and the harness's shell, delegation and file
+events. `EventTranslator.translate` turns each into the plain dataclasses in
+`pcode.runtime` that the terminal renders. It holds the per-run bookkeeping
+that pairs a tool's start with its end.
+"""
+
+import re
+from collections.abc import AsyncIterator, Iterator
+from dataclasses import replace
+from datetime import datetime, timezone
+from time import monotonic
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai import (
+    AgentRunResultEvent,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    PartDeltaEvent,
+    PartEndEvent,
+    PartStartEvent,
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+    ThinkingPartDelta,
+)
+from pydantic_ai.messages import (
+    ModelResponse,
+    NativeToolCallPart,
+    NativeToolReturnPart,
+    RetryPromptPart,
+    ToolCallPart,
+)
+from pydantic_ai_harness.shell import (
+    CommandFinishedEvent,
+    CommandOutputEvent,
+    CommandStartedEvent,
+)
+from pydantic_ai_harness.subagents import DelegationEndEvent, DelegationStartEvent
+
+from pcode.cache_warnings import CacheBustEvent
+from pcode.delegation import ChildActivity, ChildOutput
+from pcode.edit_preview import StreamingEditPreview
+from pcode.filesystem import FileChangeEvent
+from pcode.inspection import capture
+from pcode.plan_preview import StreamingPlanPreview
+from pcode.runtime import (
+    CacheBust,
+    ChildPlan,
+    ChildText,
+    Event,
+    Message,
+    PlanUpdated,
+    RunStatus,
+    TextDelta,
+    Thinking,
+    ThinkingDelta,
+    ToolStarted,
+    ToolSummary,
+)
+from pcode.shell import ShellPreview, result_projection
+from pcode.tool_display import (
+    COMMAND_TOOLS,
+    assignment,
+    command_error,
+    delegation_detail,
+    execution_mode,
+    invocation,
+    job_status,
+    label,
+    native_result_detail,
+    native_result_projection,
+    result_detail,
+    spilled_call_id,
+    stated_purpose,
+    subject,
+    target,
+)
+from pcode.turn import TurnContext
+
+if TYPE_CHECKING:
+    from pcode.live import AgentRuntime
+
+
+class EventTranslator:
+    """Per-run state and one handler per upstream event kind.
+
+    Besides translating, `on_run_result` records the finished run: it settles
+    the turn's history and counts the turn on the runtime.
+    """
+
+    def __init__(
+        self,
+        runtime: "AgentRuntime",
+        context: TurnContext,
+        plan_items: list[dict],
+        *,
+        plan_preview: StreamingPlanPreview | None,
+        edit_preview: StreamingEditPreview | None,
+    ) -> None:
+        self.runtime = runtime
+        self.context = context
+        self.run_id = context.run_id
+        # The plan as last published; a settled tool republishes it on change.
+        self.plan_items = plan_items
+        self.plan_preview = plan_preview
+        self.edit_preview = edit_preview
+        self.emitted_text = False
+        # Running tools by call id: (name, args, monotonic start).
+        self.tools: dict[str, tuple[str, dict, float]] = {}
+        # Every tool this run called, by call id, settled or not: a later
+        # `read_tool_result` names the call whose spilled output it reads.
+        self.called: dict[str, str] = {}
+        self.delegates: dict[str, ToolStarted] = {}
+        self.child_tools: dict[str, ToolStarted] = {}
+        self.delegation_ends: dict[str, DelegationEndEvent] = {}
+        self.shell_preview = ShellPreview()
+        self.shell_ends: dict[str, CommandFinishedEvent] = {}
+
+    async def translate(self, event: Any) -> AsyncIterator[Event]:
+        """Yield the pcode events for one upstream event, in display order."""
+        if self.edit_preview is not None:
+            for edit_update in self.edit_preview.update(event):
+                yield edit_update
+        if isinstance(event, FunctionToolResultEvent):
+            # The only handler that awaits: it rereads the plan store.
+            async for out in self.on_tool_result(event):
+                yield out
+        else:
+            for out in self._dispatch(event):
+                yield out
+        if self.plan_preview is not None:
+            if (update := self.plan_preview.update(event, self.plan_items)) is not None:
+                yield update
+
+    def _dispatch(self, event: Any) -> Iterator[Event]:
+        match event:
+            case CacheBustEvent():
+                return self.on_cache_bust(event)
+            case FileChangeEvent():
+                return self.on_file_change(event)
+            case DelegationStartEvent():
+                return self.on_delegation_start(event)
+            case DelegationEndEvent():
+                return self.on_delegation_end(event)
+            case ChildOutput():
+                return self.on_child_output(event)
+            case ChildActivity():
+                return self.on_child_activity(event)
+            case CommandStartedEvent() | CommandOutputEvent() | CommandFinishedEvent():
+                return self.on_shell(event)
+            case PartStartEvent():
+                return self.on_part_start(event)
+            case PartDeltaEvent():
+                return self.on_part_delta(event)
+            case PartEndEvent():
+                return self.on_part_end(event)
+            case FunctionToolCallEvent():
+                return self.on_tool_call(event)
+            case AgentRunResultEvent():
+                return self.on_run_result(event)
+            case _:
+                return iter(())
+
+    def activity(self) -> RunStatus:
+        tools = self.tools
+        if not tools:
+            return RunStatus("Waiting for model…")
+        if len(tools) == 1:
+            name, args, _ = next(iter(tools.values()))
+            where = target(name, args, self.spill_source(name, args))
+            title = label(name) if name == "delegate_task" else name
+            return RunStatus(f"Running {title}" + (f" · {where}" if where else "") + "…")
+        names = ", ".join(label(item[0]) for item in list(tools.values())[:3])
+        return RunStatus(f"Running {len(tools)} tools · {names}…")
+
+    def spill_source(self, name: str, args: dict) -> str:
+        """The tool whose stored output a `read_tool_result` call reads, or "".
+
+        This run's calls are not in the turn's history until it settles;
+        earlier turns', including a resumed session's, are. Either answer is
+        remembered, so a call's start, status, and result scan history once.
+        """
+        call_id = spilled_call_id(args) if name == "read_tool_result" else ""
+        if not call_id:
+            return ""
+        if call_id not in self.called:
+            self.called[call_id] = next(
+                (
+                    part.tool_name
+                    for message in reversed(self.context.history)
+                    if isinstance(message, ModelResponse)
+                    for part in message.parts
+                    if isinstance(part, ToolCallPart) and part.tool_call_id == call_id
+                ),
+                "",
+            )
+        return self.called[call_id]
+
+    # Harness and pcode capability events
+
+    def on_cache_bust(self, event: CacheBustEvent) -> Iterator[Event]:
+        yield CacheBust(event.text)
+
+    def on_file_change(self, event: FileChangeEvent) -> Iterator[Event]:
+        yield event.change
+
+    def on_delegation_start(self, event: DelegationStartEvent) -> Iterator[Event]:
+        if start := self.delegates.get(event.tool_call_id):
+            start = replace(start, activity="Waiting for model")
+            self.delegates[event.tool_call_id] = start
+            yield start
+
+    def on_delegation_end(self, event: DelegationEndEvent) -> Iterator[Event]:
+        self.delegation_ends[event.tool_call_id] = event
+        # Deliberately not added to session totals: a child's *tokens*
+        # already reach `result.usage` even under its own budget, so
+        # adding `event.usage` here counts them twice. Only its
+        # request count stays isolated, which is the point of the
+        # budget. Tokens from an interrupted turn are a separate gap.
+        # Timeouts/budget stops can leave a child's tool without a
+        # result. Settle it before the parent resumes its tool loop.
+        for call_id, child in list(self.child_tools.items()):
+            if child.parent_call_id == event.tool_call_id:
+                yield ToolSummary(
+                    child.name,
+                    child.detail + " → Interrupted",
+                    failed=True,
+                    call_id=call_id,
+                    run_id=self.run_id,
+                    outcome="interrupted",
+                    parent_call_id=child.parent_call_id,
+                )
+                del self.child_tools[call_id]
+
+    def on_child_output(self, event: ChildOutput) -> Iterator[Event]:
+        if event.tool_call_id in self.delegates:
+            yield ChildText(
+                event.tool_call_id,
+                event.text,
+                thinking=event.thinking,
+                start=event.start,
+            )
+
+    def on_child_activity(self, event: ChildActivity) -> Iterator[Event]:
+        if start := self.delegates.get(event.tool_call_id):
+            if event.activity != start.activity:
+                start = replace(start, activity=event.activity)
+                self.delegates[event.tool_call_id] = start
+                yield start
+            if event.plan is not None:
+                yield ChildPlan(event.tool_call_id, event.plan)
+            if event.child is not None:
+                if isinstance(event.child, ToolStarted):
+                    self.child_tools[event.child.call_id] = event.child
+                else:
+                    self.child_tools.pop(event.child.call_id, None)
+                yield event.child
+
+    def on_shell(
+        self, event: CommandStartedEvent | CommandOutputEvent | CommandFinishedEvent
+    ) -> Iterator[Event]:
+        if isinstance(event, CommandFinishedEvent):
+            self.shell_ends[event.tool_call_id] = event
+        if (output := self.shell_preview.update(event)) is not None:
+            yield output
+
+    # Model response parts
+
+    def on_part_start(self, event: PartStartEvent) -> Iterator[Event]:
+        if isinstance(event.part, TextPart):
+            yield TextDelta(event.part.content)
+        elif isinstance(event.part, ThinkingPart):
+            if event.part.content:
+                yield ThinkingDelta(event.part.content)
+            yield RunStatus("Thinking…")
+        elif isinstance(event.part, NativeToolReturnPart):
+            yield from self.on_native_tool_return(event.part)
+
+    def on_native_tool_return(self, part: NativeToolReturnPart) -> Iterator[Event]:
+        # Provider-executed tools (native web search/fetch) return
+        # inside the response stream; there is no function event.
+        name, args, started = self.tools.pop(part.tool_call_id, (part.tool_name, {}, monotonic()))
+        detail, failed = native_result_detail(name, args, part.content, part.outcome)
+        yield ToolSummary(
+            name,
+            detail,
+            failed=failed,
+            call_id=part.tool_call_id,
+            result=capture(native_result_projection(part.content)),
+            run_id=self.run_id,
+            outcome=part.outcome if not failed else "error",
+            elapsed_seconds=max(0, monotonic() - started),
+        )
+        yield self.activity()
+
+    def on_part_delta(self, event: PartDeltaEvent) -> Iterator[Event]:
+        if isinstance(event.delta, TextPartDelta):
+            yield TextDelta(event.delta.content_delta)
+        elif isinstance(event.delta, ThinkingPartDelta):
+            if event.delta.content_delta:
+                yield ThinkingDelta(event.delta.content_delta)
+            yield RunStatus("Thinking…")
+
+    def on_part_end(self, event: PartEndEvent) -> Iterator[Event]:
+        if isinstance(event.part, TextPart) and event.part.content:
+            self.emitted_text = True
+            yield Message(event.part.content)
+        elif isinstance(event.part, ThinkingPart) and event.part.content:
+            yield Thinking(event.part.content)
+        elif isinstance(event.part, NativeToolCallPart):
+            yield from self.on_native_tool_call(event.part)
+
+    def on_native_tool_call(self, part: NativeToolCallPart) -> Iterator[Event]:
+        try:
+            args = part.args_as_dict()
+        except (ValueError, TypeError):
+            args = {}
+        self.tools[part.tool_call_id] = (part.tool_name, args, monotonic())
+        agent, task = assignment(part.tool_name, args)
+        yield ToolStarted(
+            part.tool_name,
+            target(part.tool_name, args),
+            part.tool_call_id,
+            arguments=capture(args if args else part.args),
+            run_id=self.run_id,
+            started_at=datetime.now(timezone.utc).isoformat(),
+            agent=agent,
+            task=task,
+        )
+        yield self.activity()
+
+    # Function tools and the run's end
+
+    def on_tool_call(self, event: FunctionToolCallEvent) -> Iterator[Event]:
+        try:
+            args = event.part.args_as_dict()
+        except (ValueError, TypeError):
+            args = {}
+        self.tools[event.part.tool_call_id] = (event.part.tool_name, args, monotonic())
+        self.called[event.part.tool_call_id] = event.part.tool_name
+        agent, task = assignment(event.part.tool_name, args)
+        command, purpose = subject(event.part.tool_name, args, self.runtime.jobs)
+        start = ToolStarted(
+            event.part.tool_name,
+            target(event.part.tool_name, args, self.spill_source(event.part.tool_name, args)),
+            event.part.tool_call_id,
+            arguments=capture(args if args else event.part.args),
+            run_id=self.run_id,
+            started_at=datetime.now(timezone.utc).isoformat(),
+            process_id=capture(args.get("command_id", "")),
+            command=command,
+            purpose=purpose,
+            execution=execution_mode(event.part.tool_name, args),
+            agent=agent,
+            task=task,
+        )
+        if event.part.tool_name == "delegate_task":
+            self.delegates[event.part.tool_call_id] = start
+        yield start
+        yield self.activity()
+
+    async def on_tool_result(self, event: FunctionToolResultEvent) -> AsyncIterator[Event]:
+        name, args, started = self.tools.pop(
+            event.tool_call_id,
+            (event.part.tool_name or "tool", {}, monotonic()),
+        )
+        outcome = "retry" if isinstance(event.part, RetryPromptPart) else event.part.outcome
+        detail, failed = result_detail(
+            name, args, event.part.content, outcome, self.spill_source(name, args)
+        )
+        shell_end = self.shell_ends.pop(event.tool_call_id, None)
+        if name == "shell" and shell_end is not None and outcome == "success":
+            # The result's own job marker is authoritative: it is
+            # written after the wait ends, while the event is a
+            # snapshot the command can finish just after.
+            status, failed = job_status(event.part.content)
+            detail = target(name, args) + (f" → {status}" if status else "")
+            if shell_end.truncated:
+                detail += " · preview capped"
+        self.delegates.pop(event.tool_call_id, None)
+        end = self.delegation_ends.pop(event.tool_call_id, None)
+        if end is not None:
+            outcome = end.outcome
+            detail, failed = delegation_detail(args, outcome)
+        elif name == "delegate_task" and outcome == "success":
+            # Harness returns max_calls refusals as normal strings,
+            # with no lifecycle events because no child was launched.
+            outcome = "not_started"
+            detail = target(name, args) + " → Not started"
+            failed = True
+        # Read the store after every settled tool: covers granular,
+        # batched, and future plan mutations without parsing results.
+        items = [item.model_dump(mode="json") for item in await self.context.plan_store.get_items()]
+        if items != self.plan_items:
+            self.plan_items = items
+            yield PlanUpdated(items)
+        display_content = (
+            result_projection(event.part.content, shell_end)
+            if name == "shell"
+            else event.part.content
+        )
+        yield ToolSummary(
+            name,
+            detail,
+            failed=failed,
+            call_id=event.tool_call_id,
+            result=capture(display_content),
+            run_id=self.run_id,
+            outcome=outcome,
+            process_id=(
+                str(shell_end.pid)
+                if shell_end is not None
+                else match[1]
+                if name == "start_command"
+                and isinstance(event.part.content, str)
+                and (match := re.search(r"^ID: (\w+)$", event.part.content, re.MULTILINE))
+                else capture(args.get("command_id", ""))
+            ),
+            elapsed_seconds=max(0, monotonic() - started),
+            command=invocation(name, args),
+            purpose=stated_purpose(args),
+            error=command_error(display_content) if failed and name in COMMAND_TOOLS else "",
+        )
+        yield self.activity()
+
+    def on_run_result(self, event: AgentRunResultEvent) -> Iterator[Event]:
+        result = event.result
+        if result.output and not self.emitted_text:
+            yield Message(str(result.output))
+        # Full successful history. The outer persistence wrapper also
+        # recovers settled tool-boundary snapshots after failures.
+        self.context.history = result.all_messages()
+        self.context.pending_shell = []
+        self.runtime.turns += 1

pcode/strict_tools.py

@@ -0,0 +1,130 @@
+"""Grammar-constrained tool arguments for Anthropic models.
+
+Claude sometimes emits a nested tool argument as an escaped JSON *string* rather
+than the array the schema asks for, and botches the escaping while doing it. In
+~2k recorded `edit_file` calls the Anthropic share of `replacements` arrivals
+failed validation about one time in eight, while the OpenAI share never did --
+Pydantic AI turns OpenAI's strict mode on by itself, and strict mode is
+constrained decoding.
+
+Anthropic exposes the same thing as `tools[].strict`, but two gaps keep it off:
+`ToolDefinition.strict` defaults to `None`, which Pydantic AI reads as "off" for
+Anthropic (only OpenAI infers it), and the Anthropic profile carries no
+`json_schema_transformer`, so the schema Pydantic generates is rejected with
+`For 'object' type, 'additionalProperties' must be explicitly set to false`.
+This module closes the object nodes and sets the flag.
+
+It does not close the hole. With `strict: true` confirmed on the wire, Anthropic
+still delivers `replacements` as a truncated JSON string a few times in ten when
+the model is nudged toward that shape (claude-opus-5, subscription endpoint,
+pcode's schema and the plainest `{"type": "array"}` alike), while a scalar
+property in the same schema is held to its type. Production moved from 11.5% to
+9.6% of `replacements` arrivals failing. Left on because nothing measured got
+worse, not because it works; the retry budget is what actually recovers the
+edit.
+"""
+
+from dataclasses import replace
+
+from pydantic_ai.capabilities import PrepareTools
+from pydantic_ai.tools import ToolDefinition
+
+from pcode.preferences import SETTINGS, load_preferences
+
+# Only `edit_file` has a measured failure: its `replacements` array is the one
+# tool argument with a nested object shape. Keeping the set small keeps the
+# blast radius small too -- Anthropic validates the schema of tools marked
+# strict and 400s the whole request, so MCP and extension tools carrying
+# arbitrary schemas must stay out of it.
+STRICT_TOOLS = frozenset({"edit_file"})
+
+
+def strict_tools_enabled() -> bool:
+    preferences = load_preferences()
+    return preferences.get("strict_tools", SETTINGS["strict_tools"].default) == "on"
+
+
+def _closed(schema: object) -> object:
+    """Copy `schema` with every object node explicitly closed."""
+    if isinstance(schema, dict):
+        node = {key: _closed(value) for key, value in schema.items()}
+        if node.get("type") == "object":
+            node["additionalProperties"] = False
+        return node
+    if isinstance(schema, list):
+        return [_closed(item) for item in schema]
+    return schema
+
+
+# Anthropic's strict mode accepts a subset of JSON Schema and rejects the whole
+# request on anything outside it, so an upstream schema change could break every
+# edit rather than one. An allow-list fails the safe way: an unrecognized keyword
+# declines strict and leaves the tool exactly as it behaves today.
+SUPPORTED_KEYWORDS = frozenset(
+    {
+        "$defs",
+        "$ref",
+        "additionalProperties",
+        "anyOf",
+        "const",
+        "default",
+        "description",
+        "enum",
+        "items",
+        "properties",
+        "required",
+        "title",
+        "type",
+    }
+)
+
+# Keyed by name rather than by keyword, so their keys are not schema vocabulary.
+_SCHEMA_MAPS = frozenset({"$defs", "properties"})
+
+
+def _constrainable(node: object) -> bool:
+    """Whether `node` is safe to close and send as strict, and faithful once closed."""
+    if isinstance(node, list):
+        return all(_constrainable(item) for item in node)
+    if not isinstance(node, dict):
+        return True
+    # A map-valued object: closing it would forbid the keys it exists to carry.
+    if node.get("type") == "object" and not node.get("properties"):
+        return False
+    if node.keys() - SUPPORTED_KEYWORDS:
+        return False
+    return all(
+        all(_constrainable(child) for child in value.values())
+        if key in _SCHEMA_MAPS and isinstance(value, dict)
+        else _constrainable(value)
+        for key, value in node.items()
+    )
+
+
+def prepare_strict_tools(ctx, tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
+    """Mark the allow-listed tools strict when the model can honor it."""
+    model = getattr(ctx, "model", None)
+    if getattr(model, "system", None) != "anthropic":
+        return tool_defs
+    profile = getattr(model, "profile", None)
+    # The same gate `AnthropicModel` applies before it will send `strict` at all,
+    # so an unsupported model silently keeps the schema it already had.
+    if not (profile and profile.get("supports_json_schema_output", False)):
+        return tool_defs
+    return [
+        replace(
+            tool_def,
+            parameters_json_schema=_closed(tool_def.parameters_json_schema),
+            strict=True,
+        )
+        if tool_def.name in STRICT_TOOLS and _constrainable(tool_def.parameters_json_schema)
+        else tool_def
+        for tool_def in tool_defs
+    ]
+
+
+def create_strict_tools() -> PrepareTools | None:
+    """Return the capability when enabled, or None to leave tool schemas alone."""
+    if not strict_tools_enabled():
+        return None
+    return PrepareTools(prepare_strict_tools)

pcode/syntax.py

@@ -0,0 +1,43 @@
+"""Syntax highlighting rendered on the terminal's own background."""
+
+from functools import cache
+
+from pygments.token import _TokenType
+from rich.style import Style
+from rich.syntax import Syntax, SyntaxTheme
+
+
+class TransparentSyntaxTheme(SyntaxTheme):
+    """A Pygments style with its backgrounds dropped rather than painted over.
+
+    Passing ``background_color="default"`` to :class:`~rich.syntax.Syntax`
+    paints every token background, not just the block's, and some styles carry
+    meaning there: gruvbox, the default for both palettes, colours diff +/−
+    lines by background and sets their foreground to the style's own background
+    colour. Flattened that way a diff loses every marker colour and reads as
+    unhighlighted text. Fold a meaningful background into the foreground
+    instead, and drop the backgrounds that only repeat the block's.
+    """
+
+    def __init__(self, code_theme: str) -> None:
+        self.theme = Syntax.get_theme(code_theme)
+        self.background = self.theme.get_background_style().bgcolor
+
+    def get_background_style(self) -> Style:
+        return Style()
+
+    def get_style_for_token(self, token_type: _TokenType) -> Style:
+        style = self.theme.get_style_for_token(token_type)
+        color = style.color if style.bgcolor in (None, self.background) else style.bgcolor
+        return Style(
+            color=color,
+            bold=style.bold,
+            italic=style.italic,
+            underline=style.underline,
+        )
+
+
+@cache
+def transparent_theme(code_theme: str) -> TransparentSyntaxTheme:
+    """The named style, resolved once, with its backgrounds left transparent."""
+    return TransparentSyntaxTheme(code_theme)