traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/__init__.py

@@ -0,0 +1,3 @@
+"""Traceforge: an AI agent that designs, edits and reviews KiCad circuit boards."""
+
+__version__ = "1.0.0"

traceforge/__main__.py

@@ -0,0 +1,3 @@
+from traceforge.cli import main
+
+raise SystemExit(main())

traceforge/acp_server.py

@@ -0,0 +1,570 @@
+"""Agent Client Protocol server: ``traceforge acp``.
+
+Lets an ACP client (Zed, Neovim plugins, ...) use traceforge as its agent over stdio. Each ACP
+session wraps a :class:`traceforge.chat.ChatSession`, so slash commands, modes, models, threads
+and the approval policy behave as they do in the terminal; this module only translates:
+
+* agent events to ``session/update`` notifications (message and thought chunks, tool calls and
+  their results, the plan from ``write_todos``, token usage);
+* approval requests to ``session/request_permission`` (allow once, allow always, reject);
+* ACP session modes to traceforge modes (ask, plan, agent), and config options to effort, model
+  and approvals;
+* ACP prompts (text, images, files the client attaches) to a traceforge prompt.
+
+stdout carries the protocol, so nothing here prints. Needs ``pip install agent-client-protocol``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+from urllib.parse import unquote, urlparse
+
+import acp
+from acp import schema
+
+from traceforge import __version__
+from traceforge.agent import AgentEvent
+from traceforge.approval import ApprovalAnswer, ApprovalRequest
+from traceforge.attachments import AttachmentError, ImageAttachment
+from traceforge.chat import APPROVALS, EFFORTS, MODE_COPY, MODES, ChatSession
+from traceforge.config import Settings
+from traceforge.sessions import SessionStore, prompt_text
+
+ALLOW_ONCE, ALLOW_ALWAYS, REJECT = "allow_once", "allow_always", "reject_once"
+_TOOL_KINDS = {"read": "read", "write": "edit", "shell": "execute", "interactive": "other"}
+_SEARCH_TOOLS = {"grep", "search_files", "glob", "find_parts", "search_parts", "web_search"}
+_STOP = {"end_turn": "end_turn", "max_tokens": "max_tokens", "refusal": "refusal"}
+_MODE_NAMES = {"ask": "Ask", "plan": "Plan", "agent": "Agent"}
+_TEXT_LIMIT = 6000
+
+
+def tool_kind(name: str, kind: str) -> str:
+    if name in _SEARCH_TOOLS:
+        return "search"
+    return _TOOL_KINDS.get(str(kind), "other")
+
+
+def _locations(arguments: dict[str, Any], workspace: Path) -> list[schema.ToolCallLocation]:
+    path = arguments.get("path") if isinstance(arguments, dict) else None
+    if not isinstance(path, str) or not path:
+        return []
+    target = Path(path)
+    return [
+        schema.ToolCallLocation(path=str(target if target.is_absolute() else workspace / target))
+    ]
+
+
+def file_path(uri: str) -> str:
+    parsed = urlparse(uri)
+    return unquote(parsed.path) if parsed.scheme == "file" else uri
+
+
+def prompt_to_text(blocks: list[Any]) -> tuple[str, list[ImageAttachment], list[str]]:
+    """The text, image attachments and problems of an ACP prompt's content blocks."""
+    parts: list[str] = []
+    images: list[ImageAttachment] = []
+    problems: list[str] = []
+    for block in blocks:
+        kind = getattr(block, "type", "")
+        if kind == "text":
+            parts.append(block.text)
+        elif kind == "image":
+            try:
+                if block.data:
+                    images.append(ImageAttachment.from_base64(block.data, "image"))
+                elif block.uri and file_path(block.uri).startswith("/"):
+                    images.append(ImageAttachment.from_path(Path(file_path(block.uri))))
+            except AttachmentError as exc:
+                problems.append(str(exc))
+        elif kind == "resource_link":
+            parts.append(f"@{file_path(block.uri)}")
+        elif kind == "resource":
+            resource = block.resource
+            body = getattr(resource, "text", None)
+            if body is not None:
+                parts.append(
+                    f"<attached_file path={file_path(resource.uri)!r}>\n{body}\n</attached_file>"
+                )
+    return "\n".join(p for p in parts if p).strip(), images, problems
+
+
+@dataclass
+class AcpSession:
+    id: str
+    chat: ChatSession
+    cancelled: bool = False
+    calls: dict[str, str] = field(default_factory=dict)  # tool name -> id of its latest call
+    pending: set[asyncio.Future] = field(default_factory=set)
+
+
+class AcpUI:
+    """:class:`traceforge.chat.ChatUI` that talks to an ACP client."""
+
+    def __init__(self, agent: TraceforgeAcpAgent):
+        self.agent = agent
+        self.session: AcpSession | None = None
+
+    @property
+    def conn(self) -> Any:
+        return self.agent.conn
+
+    async def _update(self, update: Any) -> None:
+        assert self.session is not None
+        await self.conn.session_update(session_id=self.session.id, update=update)
+
+    async def on_event(self, event: AgentEvent) -> None:
+        session = self.session
+        assert session is not None
+        kind, data = event.kind, event.data
+        if kind == "text":
+            await self._update(acp.update_agent_message_text(str(data)))
+        elif kind == "thinking":
+            await self._update(acp.update_agent_thought_text(str(data)))
+        elif kind == "tool_call":
+            arguments = data.get("input") or {}
+            name = str(data.get("name"))
+            session.calls[name] = str(data.get("id"))
+            kind_name = self.agent.kind_of(session, name)
+            await self._update(
+                acp.start_tool_call(
+                    str(data.get("id")),
+                    f"{name} {_brief(arguments)}".strip(),
+                    kind=kind_name,
+                    status="in_progress",
+                    locations=_locations(arguments, session.chat.workspace),
+                    raw_input=arguments,
+                )
+            )
+        elif kind == "tool_result":
+            content = str(data.get("content") or "")
+            await self._update(
+                acp.update_tool_call(
+                    str(data.get("id")),
+                    status="failed" if data.get("is_error") else "completed",
+                    content=[acp.tool_content(acp.text_block(content[:_TEXT_LIMIT]))]
+                    if content
+                    else None,
+                    raw_output=content[:_TEXT_LIMIT],
+                )
+            )
+        elif kind == "todos":
+            entries = [
+                acp.plan_entry(
+                    str(t.get("content") or t.get("task") or ""),
+                    status=t.get("status")
+                    if t.get("status") in {"pending", "in_progress", "completed"}
+                    else "pending",
+                )
+                for t in data or []
+                if isinstance(t, dict)
+            ]
+            await self._update(acp.update_plan(entries))
+        elif kind == "usage" and isinstance(data, dict):
+            used = int(data.get("input_tokens", 0)) + int(data.get("cache_read_input_tokens", 0))
+            from traceforge.models import context_window
+
+            await self._update(
+                schema.UsageUpdate(used=used, size=context_window(), session_update="usage_update")
+            )
+        elif kind in {"status", "notice"}:
+            message = data.get("message") if isinstance(data, dict) else data
+            await self._update(acp.update_agent_thought_text(f"{message}\n"))
+        elif kind == "error":
+            await self._update(acp.update_agent_message_text(f"\n**Error:** {data}\n"))
+
+    async def approve(self, request: ApprovalRequest) -> bool | ApprovalAnswer:
+        session = self.session
+        assert session is not None
+        call_id = session.calls.get(request.tool, request.tool)
+        options = [
+            schema.PermissionOption(option_id=ALLOW_ONCE, name="Allow", kind="allow_once"),
+            schema.PermissionOption(
+                option_id=ALLOW_ALWAYS, name=f"Always allow {request.tool}", kind="allow_always"
+            ),
+            schema.PermissionOption(option_id=REJECT, name="Reject", kind="reject_once"),
+        ]
+        tool_call = schema.ToolCallUpdate(
+            tool_call_id=call_id,
+            title=request.summary or f"{request.tool} {_brief(request.arguments)}",
+            kind=tool_kind(request.tool, request.kind),
+            status="pending",
+            raw_input=request.arguments,
+        )
+        outcome = await self._ask_client(tool_call, options)
+        if outcome == ALLOW_ONCE:
+            return ApprovalAnswer(True, "once")
+        if outcome == ALLOW_ALWAYS:
+            return ApprovalAnswer(True, "tool")
+        return ApprovalAnswer(False, "once", "the user rejected it in the editor")
+
+    async def _ask_client(self, tool_call: Any, options: list[Any]) -> str | None:
+        """The chosen option id, or None when the client cancelled."""
+        session = self.session
+        assert session is not None
+        future = asyncio.ensure_future(
+            self.conn.request_permission(
+                session_id=session.id, tool_call=tool_call, options=options
+            )
+        )
+        session.pending.add(future)
+        try:
+            response = await future
+        except asyncio.CancelledError:
+            return None
+        finally:
+            session.pending.discard(future)
+        outcome = response.outcome
+        return outcome.option_id if getattr(outcome, "outcome", "") == "selected" else None
+
+    async def ask(self, question: str, options: list[str]) -> str:
+        if not options:
+            return (
+                "(this editor cannot take a typed answer; continue with your best judgement and "
+                "state your assumption)"
+            )
+        listed = [
+            schema.PermissionOption(option_id=f"opt{i}", name=text[:80], kind="allow_once")
+            for i, text in enumerate(options)
+        ]
+        tool_call = schema.ToolCallUpdate(
+            tool_call_id=f"ask-{abs(hash(question)) % 10**8}", title=question, status="pending"
+        )
+        chosen = await self._ask_client(tool_call, listed)
+        if chosen and chosen.startswith("opt") and chosen[3:].isdigit():
+            return options[int(chosen[3:])]
+        return ""
+
+    async def choose(self, title: str, options: list[tuple[str, str]]) -> str | None:
+        listed = [
+            schema.PermissionOption(option_id=f"opt{i}", name=label[:80], kind="allow_once")
+            for i, (_v, label) in enumerate(options[:12])
+        ]
+        tool_call = schema.ToolCallUpdate(
+            tool_call_id=f"choose-{abs(hash(title)) % 10**8}", title=title, status="pending"
+        )
+        chosen = await self._ask_client(tool_call, listed)
+        if chosen and chosen[3:].isdigit() and int(chosen[3:]) < len(options):
+            return options[int(chosen[3:])][0]
+        return None
+
+    async def confirm(self, question: str) -> bool:
+        return await self.choose(question, [("yes", "Yes"), ("no", "No")]) == "yes"
+
+    def show(self, text: str, *, markdown: bool = False) -> None:
+        self.agent.spawn(self._update(acp.update_agent_message_text(f"{text}\n\n")))
+
+    def error(self, text: str) -> None:
+        self.agent.spawn(self._update(acp.update_agent_message_text(f"**Error:** {text}\n\n")))
+
+    async def clear(self) -> None:
+        return None
+
+    async def replay(self, thread: Any) -> None:
+        for message in thread.messages:
+            role = message.get("role")
+            content = message.get("content")
+            blocks = content if isinstance(content, list) else [{"type": "text", "text": content}]
+            text = "".join(
+                str(b.get("text", ""))
+                for b in blocks
+                if isinstance(b, dict)
+                and b.get("type") == "text"
+                and not str(b.get("text", "")).startswith("<workspace>")
+            ).strip()
+            if not text:
+                continue
+            if role == "assistant":
+                await self._update(acp.update_agent_message_text(text))
+            elif role == "user" and prompt_text(message):
+                await self._update(acp.update_user_message_text(prompt_text(message)))
+
+
+def _brief(arguments: dict[str, Any], limit: int = 70) -> str:
+    for key in ("path", "command", "query", "reference", "description"):
+        if isinstance(arguments.get(key), str):
+            first = (arguments[key].splitlines() or [""])[0]
+            return first[:limit]
+    text = json.dumps(arguments, default=str)
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+class TraceforgeAcpAgent:
+    """The ACP ``Agent`` for traceforge (see :class:`acp.Agent`)."""
+
+    def __init__(
+        self,
+        *,
+        client: Any = None,
+        store: SessionStore | None = None,
+        settings_factory: Any = None,
+    ):
+        self.model_client = client  # injectable model client (tests)
+        self.store = store
+        self.settings_factory = settings_factory or (lambda cwd: Settings.load(cwd))
+        self.sessions: dict[str, AcpSession] = {}
+        self.conn: Any = None
+        self._tasks: set[asyncio.Task] = set()
+        self.client_capabilities: Any = None
+
+    # ------------------------------------------------------------ plumbing
+
+    def on_connect(self, conn: Any) -> None:
+        self.conn = conn
+
+    def spawn(self, coroutine: Any) -> None:
+        task = asyncio.ensure_future(coroutine)
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    def kind_of(self, session: AcpSession, name: str) -> str:
+        tool = session.chat.runtime.registry.get(name)
+        return tool_kind(name, tool.kind.value if tool else "other")
+
+    def _make_session(self, cwd: str, mode: str = "agent") -> AcpSession:
+        workspace = Path(cwd).expanduser().resolve()
+        settings = self.settings_factory(workspace)
+        ui = AcpUI(self)
+        chat = ChatSession(
+            workspace,
+            settings,
+            ui,
+            mode=mode,
+            store=self.store,
+            client=self.model_client,
+            terminal_extras=False,  # the editor shows diffs and does its own review
+        )
+        session = AcpSession("", chat)
+        ui.session = session
+        return session
+
+    def _register(self, session: AcpSession, session_id: str) -> None:
+        session.id = session_id
+        self.sessions[session_id] = session
+
+    def _get(self, session_id: str) -> AcpSession:
+        try:
+            return self.sessions[session_id]
+        except KeyError:
+            raise acp.RequestError.invalid_params(
+                {"session_id": f"unknown session {session_id}"}
+            ) from None
+
+    def _modes(self, session: AcpSession) -> schema.SessionModeState:
+        return schema.SessionModeState(
+            current_mode_id=session.chat.mode,
+            available_modes=[
+                schema.SessionMode(id=m, name=_MODE_NAMES[m], description=MODE_COPY[m])
+                for m in MODES
+            ],
+        )
+
+    def _options(self, session: AcpSession) -> list[Any]:
+        from traceforge import models
+
+        chat = session.chat
+        model_options = models.picker_options(chat.settings.model) + [
+            (tier, f"{tier} tier") for tier in models.TIERS
+        ]
+        seen: set[str] = set()
+        unique = [(v, n) for v, n in model_options if not (v in seen or seen.add(v))]
+
+        def select(config_id: str, name: str, current: str, values: list[tuple[str, str]]):
+            return schema.SessionConfigOptionSelect(
+                id=config_id,
+                name=name,
+                type="select",
+                current_value=current,
+                options=[schema.SessionConfigSelectOption(value=v, name=n) for v, n in values],
+            )
+
+        return [
+            select("model", "Model", chat.settings.model, unique),
+            select("effort", "Effort", chat.settings.effort, [(e, e) for e in EFFORTS]),
+            select("approval", "Approvals", chat.settings.approval, [(a, a) for a in APPROVALS]),
+        ]
+
+    async def _announce_commands(self, session: AcpSession) -> None:
+        from traceforge.chat import command_list
+
+        commands = [
+            schema.AvailableCommand(name=c.name, description=c.summary) for c in command_list()
+        ] + [
+            schema.AvailableCommand(name=s.name, description=s.description)
+            for s in session.chat.skills.values()
+        ]
+        await self.conn.session_update(
+            session_id=session.id,
+            update=schema.AvailableCommandsUpdate(
+                available_commands=commands, session_update="available_commands_update"
+            ),
+        )
+
+    # ------------------------------------------------------------- protocol
+
+    async def initialize(
+        self,
+        protocol_version: int,
+        client_capabilities: Any = None,
+        client_info: Any = None,
+        **kwargs: Any,
+    ) -> schema.InitializeResponse:
+        self.client_capabilities = client_capabilities
+        return schema.InitializeResponse(
+            protocol_version=min(protocol_version, acp.PROTOCOL_VERSION),
+            agent_capabilities=schema.AgentCapabilities(
+                load_session=True,
+                prompt_capabilities=schema.PromptCapabilities(image=True, embedded_context=True),
+                session_capabilities=schema.SessionCapabilities(
+                    list=schema.SessionListCapabilities()
+                ),
+            ),
+            agent_info=schema.Implementation(
+                name="traceforge", title="Traceforge", version=__version__
+            ),
+            auth_methods=[],
+        )
+
+    async def new_session(
+        self, cwd: str, additional_directories: Any = None, mcp_servers: Any = None, **kwargs: Any
+    ) -> schema.NewSessionResponse:
+        session = self._make_session(cwd)
+        chat = session.chat
+        assert chat.store is not None
+        thread = chat.store.create(chat.workspace, "ACP session", chat.mode, meta={"client": "acp"})
+        chat.thread = thread
+        self._register(session, thread.id)
+        self.spawn(self._announce_commands(session))
+        return schema.NewSessionResponse(
+            session_id=thread.id, modes=self._modes(session), config_options=self._options(session)
+        )
+
+    async def load_session(
+        self,
+        cwd: str,
+        session_id: str,
+        mcp_servers: Any = None,
+        additional_directories: Any = None,
+        **kwargs: Any,
+    ) -> schema.LoadSessionResponse | None:
+        session = self._make_session(cwd)
+        try:
+            await session.chat.resume(session_id)  # replays the history to the client
+        except LookupError as exc:
+            raise acp.RequestError.invalid_params({"session_id": str(exc)}) from exc
+        assert session.chat.thread is not None
+        self._register(session, session_id)
+        self.spawn(self._announce_commands(session))
+        return schema.LoadSessionResponse(
+            modes=self._modes(session), config_options=self._options(session)
+        )
+
+    async def list_sessions(
+        self, cwd: str | None = None, cursor: str | None = None, **kwargs: Any
+    ) -> schema.ListSessionsResponse:
+        store = self.store or SessionStore()
+        threads = store.list(Path(cwd).resolve() if cwd else None, limit=50)
+        return schema.ListSessionsResponse(
+            sessions=[
+                schema.SessionInfo(session_id=t.id, cwd=t.workspace, title=t.title) for t in threads
+            ]
+        )
+
+    async def set_session_mode(
+        self, session_id: str, mode_id: str, **kwargs: Any
+    ) -> schema.SetSessionModeResponse | None:
+        session = self._get(session_id)
+        try:
+            await session.chat.set_mode(mode_id)
+        except ValueError as exc:
+            raise acp.RequestError.invalid_params({"mode_id": str(exc)}) from exc
+        return schema.SetSessionModeResponse()
+
+    async def set_config_option(
+        self, config_id: str, session_id: str, value: Any, **kwargs: Any
+    ) -> schema.SetSessionConfigOptionResponse | None:
+        chat = self._get(session_id).chat
+        try:
+            if config_id == "effort":
+                await chat.set_effort(str(value))
+            elif config_id == "model":
+                await chat.set_model(str(value))
+            elif config_id == "approval":
+                await chat.set_approval(str(value), confirm=False)
+            else:
+                raise ValueError(f"unknown option {config_id}")
+        except ValueError as exc:
+            raise acp.RequestError.invalid_params({"config_id": str(exc)}) from exc
+        return schema.SetSessionConfigOptionResponse(
+            config_options=self._options(self._get(session_id))
+        )
+
+    async def authenticate(self, method_id: str, **kwargs: Any) -> None:
+        return None
+
+    async def prompt(
+        self, session_id: str, prompt: list[Any], **kwargs: Any
+    ) -> schema.PromptResponse:
+        session = self._get(session_id)
+        chat = session.chat
+        if chat.busy:
+            raise acp.RequestError.invalid_request({"session_id": "a turn is already running"})
+        session.cancelled = False
+        text, images, problems = prompt_to_text(prompt)
+        for problem in problems:
+            chat.ui.error(problem)
+        chat.pending_images.extend(images)
+        if not text and not images:
+            return schema.PromptResponse(stop_reason="end_turn")
+        before = chat.usage.to_dict() if hasattr(chat.usage, "to_dict") else {}
+        result = await chat.submit(text or "(see the attached image)")
+        if session.cancelled:
+            return schema.PromptResponse(stop_reason="cancelled")
+        usage = None
+        if result is not None:
+            u = result.usage
+            usage = schema.Usage(
+                total_tokens=u.input_tokens + u.output_tokens,
+                input_tokens=u.input_tokens,
+                output_tokens=u.output_tokens,
+                cached_read_tokens=u.cache_read_input_tokens or None,
+            )
+        del before
+        stop = _STOP.get(result.stop_reason or "end_turn", "end_turn") if result else "end_turn"
+        return schema.PromptResponse(stop_reason=stop, usage=usage)
+
+    async def cancel(self, session_id: str, **kwargs: Any) -> None:
+        session = self.sessions.get(session_id)
+        if session is None:
+            return
+        session.cancelled = True
+        for future in list(session.pending):
+            future.cancel()
+        await session.chat.cancel()
+
+    async def close_session(self, session_id: str, **kwargs: Any) -> None:
+        session = self.sessions.pop(session_id, None)
+        if session is not None:
+            await session.chat.cancel()
+            await session.chat.processes.kill_all()
+
+    async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        raise acp.RequestError.method_not_found(method)
+
+    async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
+        return None
+
+
+def run_acp(workspace: Path, approval: str = "ask") -> None:
+    """Run the agent on stdin/stdout until the client disconnects (``traceforge acp``)."""
+
+    def settings_for(cwd: Path) -> Settings:
+        return Settings.load(cwd, {"approval": approval})
+
+    async def serve() -> None:
+        await acp.run_agent(TraceforgeAcpAgent(settings_factory=settings_for))
+
+    asyncio.run(serve())