traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/remote/__init__.py

@@ -0,0 +1,33 @@
+"""Hosted mode: the model runs on a traceforge server, tools run on the user's machine.
+
+Protocol (JSON over one WebSocket per turn, ``/v1/agent``):
+
+client -> server  {"type": "turn", "protocol": 2, "mode": "agent", "effort": "high",
+                   "model": "...optional, one of GET /v1/models...",
+                   "messages": [...history...], "content": [...new user content blocks...],
+                   "tools": [...tool definitions...], "parallel_safe": [...read-only tool names...],
+                   "instructions": "...project instructions..."}
+server -> client  {"type": "event", "kind": "text" | "thinking" | "status" | "usage" | "error"
+                   | "notice", "data": ...}
+server -> client  {"type": "checkpoint", "start": n, "messages": [...history from index n...]}
+server -> client  {"type": "tool_request", "id": "...", "name": "...", "input": {...}}
+client -> server  {"type": "tool_result", "id": "...", "content": "...", "is_error": false}
+client -> server  {"type": "cancel"}
+server -> client  {"type": "done", "messages": [...full history...], "usage": {...}, "text": "...",
+                   "stop_reason": "...", "refusal": null}
+server -> client  {"type": "error", "code": "...", "message": "...", "retryable": false}
+
+The server is stateless: the conversation history lives with the client. A
+``checkpoint`` frame follows every model step (protocol 2), so a client whose
+connection drops mid-turn resumes from the last checkpoint with the tool results
+it already has, instead of starting over. Tool requests for tools listed in
+``parallel_safe`` may arrive several at a time. Protocol-1 clients get no
+checkpoints and no parallel requests.
+
+HTTP endpoints (same ``Authorization: Bearer`` key): ``GET /health``,
+``GET /v1/models`` (models a client may request) and ``GET /v1/me`` (key name,
+tokens used this month, limit and reset date).
+"""
+
+PROTOCOL_VERSION = 2
+SUPPORTED_PROTOCOLS = (1, 2)

traceforge/remote/client.py

@@ -0,0 +1,422 @@
+"""Client side of hosted mode: runs tools locally for a remote agent.
+
+Durability: the server sends a ``checkpoint`` frame with the new history after
+every model step, and the client saves it in the thread. If the connection drops
+in the middle of a turn (a server restart, a proxy timeout, a laptop changing
+networks), the client reconnects with backoff and resumes from the last
+checkpoint: tool results it already computed are sent back instead of running the
+tools again, and the agent is told to continue where it left off. Failures that
+retrying cannot fix (a revoked key, a used-up quota, an invalid request) stop at once.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from traceforge.agent import (
+    AgentError,
+    AgentEvent,
+    EventSink,
+    LocalExecutor,
+    TurnResult,
+    Usage,
+    _close_dangling_tool_uses,
+    _emit,
+)
+from traceforge.approval import ApprovalMode, ApprovalPolicy, Approver
+from traceforge.config import Settings
+from traceforge.hooks import HookedExecutor, HookRunner
+from traceforge.remote import PROTOCOL_VERSION
+from traceforge.runtime import TurnLoop, project_instructions, registry_for_mode
+from traceforge.sessions import SessionStore, Thread
+from traceforge.tools import ToolContext
+from traceforge.tools.base import QuestionHandler
+
+MAX_ATTEMPTS = 4
+RESUME_TEXT = (
+    "The connection to the traceforge server dropped in the middle of this turn and has been "
+    "restored. The history above is everything that completed. Continue the task where you left "
+    "off; do not repeat steps that already succeeded."
+)
+
+
+def websocket_url(base: str) -> str:
+    base = base.rstrip("/")
+    if base.startswith("https://"):
+        return "wss://" + base[len("https://") :] + "/v1/agent"
+    if base.startswith("http://"):
+        return "ws://" + base[len("http://") :] + "/v1/agent"
+    return base + "/v1/agent"
+
+
+def http_url(base: str, path: str) -> str:
+    base = base.rstrip("/")
+    if base.startswith("wss://"):
+        base = "https://" + base[len("wss://") :]
+    elif base.startswith("ws://"):
+        base = "http://" + base[len("ws://") :]
+    return base + path
+
+
+def backoff_schedule() -> list[float]:
+    raw = os.environ.get("TRACEFORGE_RETRY_BACKOFF", "2,5,10,20")
+    try:
+        values = [max(0.0, float(item)) for item in raw.split(",") if item.strip()]
+    except ValueError:
+        values = []
+    return values or [2.0, 5.0, 10.0, 20.0]
+
+
+def idle_timeout() -> float | None:
+    """Seconds the server may stay silent mid-turn (``TRACEFORGE_IDLE_TIMEOUT``, 0 = forever)."""
+    try:
+        value = float(os.environ.get("TRACEFORGE_IDLE_TIMEOUT", "900"))
+    except ValueError:
+        value = 900.0
+    return value if value > 0 else None
+
+
+class RemoteError(AgentError):
+    def __init__(self, message: str, *, kind: str = "server", retryable: bool = False):
+        super().__init__(message)
+        self.kind = kind
+        self.retryable = retryable
+
+
+def classify(exc: BaseException) -> RemoteError:
+    """Map a transport failure onto (kind, retryable)."""
+    if isinstance(exc, RemoteError):
+        return exc
+    try:
+        from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus
+    except ImportError:  # pragma: no cover
+        ConnectionClosed = InvalidHandshake = InvalidStatus = ()  # type: ignore[assignment]
+    if isinstance(exc, InvalidStatus):
+        status = exc.response.status_code
+        message = f"server refused the connection: HTTP {status}"
+        if status in (401, 403):
+            return RemoteError(message, kind="authentication")
+        if status == 429 or status >= 500:
+            return RemoteError(message, kind="server", retryable=True)
+        return RemoteError(message, kind="invalid_request")
+    if isinstance(exc, ConnectionClosed):
+        received = getattr(exc, "rcvd", None)
+        if received is not None and received.code in (4401, 4403):
+            return RemoteError(
+                f"server rejected the request: {received.reason or received.code}",
+                kind="authentication",
+            )
+        if received is not None and received.code == 4429:
+            return RemoteError(
+                f"server rejected the request: {received.reason or 'quota'}", kind="quota"
+            )
+        if received is not None and received.code in (4400, 1008):
+            return RemoteError(
+                f"server rejected the request: {received.reason}", kind="invalid_request"
+            )
+        return RemoteError(f"lost connection: {exc}", kind="connection", retryable=True)
+    if isinstance(exc, InvalidHandshake):
+        return RemoteError(f"handshake failed: {exc}", kind="connection", retryable=True)
+    if isinstance(exc, TimeoutError):
+        return RemoteError("the server stopped responding", kind="timeout", retryable=True)
+    if isinstance(exc, OSError):
+        return RemoteError(f"could not reach the server: {exc}", kind="network", retryable=True)
+    return RemoteError(f"{type(exc).__name__}: {exc}", kind="unknown")
+
+
+_FATAL_CODES = {"unauthorized": "authentication", "quota": "quota", "protocol": "protocol"}
+
+
+@dataclass
+class _Attempt:
+    """State carried across reconnects within one exchange."""
+
+    working: list[dict[str, Any]]
+    results: dict[str, tuple[str, bool]] = field(default_factory=dict)
+    progressed: bool = False
+
+
+class RemoteRuntime(TurnLoop):
+    """Same interface as :class:`traceforge.runtime.Runtime`, backed by a server."""
+
+    def __init__(
+        self,
+        workspace: Path,
+        settings: Settings,
+        mode: str = "agent",
+        approver: Approver | None = None,
+        ask_user: QuestionHandler | None = None,
+        store: SessionStore | None = None,
+        hooks: HookRunner | None = None,
+    ):
+        if not settings.server_url:
+            raise ValueError("server_url is not configured")
+        self.workspace = Path(workspace).resolve()
+        self.settings = settings
+        self.mode = mode
+        self.store = store or SessionStore()
+        self.hooks = hooks or HookRunner(self.workspace, settings.hooks)
+        approval = (
+            ApprovalMode.READ_ONLY if mode in {"ask", "review"} else ApprovalMode(settings.approval)
+        )
+        self.policy = ApprovalPolicy(approval, approver)
+        self.context = ToolContext(
+            self.workspace,
+            ask_user=ask_user,
+            read_only=mode in {"ask", "review", "plan"},
+            state={"fab": settings.fab},
+        )
+        self.registry = registry_for_mode(mode)
+        self.executor = HookedExecutor(
+            LocalExecutor(self.registry, self.context, self.policy), self.hooks
+        )
+        self.usage = Usage()
+        self._socket: Any = None
+
+    @property
+    def agent(self) -> RemoteRuntime:  # `runtime.agent.usage` works for both runtimes
+        return self
+
+    async def cancel(self) -> None:
+        if self._socket is not None:
+            try:
+                await self._socket.send(json.dumps({"type": "cancel"}))
+            except Exception:  # noqa: BLE001 - the socket may already be gone
+                pass
+
+    def _start_frame(self, messages: list[dict[str, Any]], content: list[dict[str, Any]]) -> dict:
+        frame: dict[str, Any] = {
+            "type": "turn",
+            "protocol": PROTOCOL_VERSION,
+            "mode": self.mode,
+            "effort": self.settings.effort,
+            "messages": messages,
+            "content": content,
+            "tools": self.executor.definitions(),
+            "parallel_safe": sorted(
+                name for name in self.registry.names() if self.executor.parallel_safe(name)
+            ),
+            "instructions": project_instructions(self.workspace, self.settings)[:20_000],
+        }
+        if self.settings.model and self.settings.model_explicit:
+            frame["model"] = self.settings.model
+        return frame
+
+    async def _exchange(
+        self, thread: Thread, content: list[dict[str, Any]], sink: EventSink | None
+    ) -> TurnResult:
+        try:
+            from websockets.asyncio.client import connect  # noqa: F401
+        except ImportError as exc:  # pragma: no cover
+            raise AgentError(
+                "hosted mode needs the 'remote' extra: pip install 'traceforge-pcb[remote]'"
+            ) from exc
+        self.executor.env = self.hook_env(thread)
+        attempt = _Attempt(list(thread.messages))
+        request_content = content
+        delays = backoff_schedule()
+        for number in range(1, MAX_ATTEMPTS + 1):
+            try:
+                return await self._attempt(thread, attempt, request_content, sink)
+            except asyncio.CancelledError:
+                await self.cancel()
+                self._save_partial(thread, attempt)
+                raise
+            except BaseException as exc:  # noqa: BLE001 - classified below
+                if isinstance(exc, KeyboardInterrupt):
+                    raise
+                error = classify(exc)
+                if not error.retryable or number == MAX_ATTEMPTS:
+                    self._save_partial(thread, attempt)
+                    raise error from exc
+                delay = delays[min(number - 1, len(delays) - 1)]
+                await _emit(
+                    sink,
+                    "notice",
+                    {
+                        "kind": "retry",
+                        "attempt": number + 1,
+                        "max_attempts": MAX_ATTEMPTS,
+                        "delay": delay,
+                        "error": error.kind,
+                        "message": f"{error}; reconnecting in {delay:g}s "
+                        f"(attempt {number + 1}/{MAX_ATTEMPTS})",
+                    },
+                )
+                await asyncio.sleep(delay)
+                if attempt.progressed:
+                    request_content = self._resume_point(attempt)
+                    self._save_partial(thread, attempt, close=False)
+        raise AgentError("the server closed the connection before finishing the turn")
+
+    def _resume_point(self, attempt: _Attempt) -> list[dict[str, Any]]:
+        """Close the checkpointed history with known tool results; return the resume prompt."""
+        history = attempt.working
+        if history and history[-1]["role"] == "assistant":
+            pending = [
+                block
+                for block in history[-1]["content"]
+                if isinstance(block, dict) and block.get("type") == "tool_use"
+            ]
+            if pending:
+                results = []
+                for block in pending:
+                    text, is_error = attempt.results.get(
+                        block["id"], ("not run: the connection dropped before this tool ran", True)
+                    )
+                    result = {"type": "tool_result", "tool_use_id": block["id"], "content": text}
+                    if is_error:
+                        result["is_error"] = True
+                    results.append(result)
+                history.append({"role": "user", "content": results})
+        return [{"type": "text", "text": RESUME_TEXT}]
+
+    def _save_partial(self, thread: Thread, attempt: _Attempt, *, close: bool = True) -> None:
+        if not attempt.progressed:
+            return
+        messages = [dict(m) for m in attempt.working]
+        if close:
+            _close_dangling_tool_uses(messages)
+        thread.messages = messages
+        if self.store is not None:
+            self.store.save(thread)
+
+    async def _attempt(
+        self,
+        thread: Thread,
+        attempt: _Attempt,
+        content: list[dict[str, Any]],
+        sink: EventSink | None,
+    ) -> TurnResult:
+        from websockets.asyncio.client import connect
+
+        assert self.settings.server_url
+        headers = {"Authorization": f"Bearer {self.settings.server_key or ''}"}
+        timeout = idle_timeout()
+        send_lock = asyncio.Lock()
+        tool_tasks: set[asyncio.Task] = set()
+        async with connect(
+            websocket_url(self.settings.server_url),
+            additional_headers=headers,
+            max_size=16 * 1024 * 1024,
+            open_timeout=30,
+        ) as socket:
+            self._socket = socket
+            try:
+                await socket.send(json.dumps(self._start_frame(attempt.working, content)))
+                while True:
+                    raw = await asyncio.wait_for(socket.recv(), timeout)
+                    message = json.loads(raw)
+                    kind = message.get("type")
+                    if kind == "event":
+                        await _emit(sink, message.get("kind", "status"), message.get("data"))
+                    elif kind == "checkpoint":
+                        start = int(message.get("start") or 0)
+                        new = message.get("messages") or []
+                        if 0 <= start <= len(attempt.working):
+                            attempt.working[start:] = new
+                            attempt.progressed = True
+                            thread.messages = [dict(m) for m in attempt.working]
+                            if self.store is not None:
+                                self.store.save(thread)
+                    elif kind == "tool_request":
+                        task = asyncio.create_task(
+                            self._run_tool(message, socket, send_lock, attempt, sink)
+                        )
+                        tool_tasks.add(task)
+                        task.add_done_callback(tool_tasks.discard)
+                    elif kind == "done":
+                        thread.messages = message.get("messages") or attempt.working
+                        usage = Usage(
+                            **{
+                                k: int(v)
+                                for k, v in (message.get("usage") or {}).items()
+                                if k in Usage.__dataclass_fields__
+                            }
+                        )
+                        for name in Usage.__dataclass_fields__:
+                            setattr(
+                                self.usage, name, getattr(self.usage, name) + getattr(usage, name)
+                            )
+                        thread.usage = self.usage.to_dict()
+                        if self.store is not None:
+                            self.store.save(thread)
+                        await _emit(
+                            sink,
+                            "done",
+                            {"stop_reason": message.get("stop_reason"), "usage": usage.to_dict()},
+                        )
+                        return TurnResult(
+                            message.get("text", ""),
+                            message.get("stop_reason"),
+                            usage,
+                            refusal=message.get("refusal"),
+                        )
+                    elif kind == "error":
+                        code = str(message.get("code") or "")
+                        raise RemoteError(
+                            f"server: {message.get('message')} ({code})",
+                            kind=_FATAL_CODES.get(code, "server"),
+                            retryable=bool(message.get("retryable")),
+                        )
+            finally:
+                self._socket = None
+                for task in tool_tasks:
+                    task.cancel()
+
+    async def _run_tool(
+        self,
+        message: dict[str, Any],
+        socket: Any,
+        lock: asyncio.Lock,
+        attempt: _Attempt,
+        sink: EventSink | None,
+    ) -> None:
+        tool_id = str(message.get("id"))
+        name, arguments = message.get("name", ""), message.get("input") or {}
+        if tool_id in attempt.results:  # a resumed server asking again: do not re-run
+            text, is_error = attempt.results[tool_id]
+        else:
+            await _emit(sink, "tool_call", {"id": tool_id, "name": name, "input": arguments})
+            text, is_error = await self.executor.execute(name, arguments, tool_id)
+            attempt.results[tool_id] = (text, is_error)
+            await _emit(
+                sink,
+                "tool_result",
+                {"id": tool_id, "name": name, "content": text, "is_error": is_error},
+            )
+        async with lock:
+            await socket.send(
+                json.dumps(
+                    {"type": "tool_result", "id": tool_id, "content": text, "is_error": is_error}
+                )
+            )
+
+
+# ------------------------------------------------------------------ account API
+
+
+async def server_get(settings: Settings, path: str, timeout: float = 10.0) -> dict[str, Any]:
+    """GET a JSON endpoint of the configured server with the saved key."""
+    import httpx
+
+    if not settings.server_url:
+        raise AgentError("no hosted server configured; run `traceforge login`")
+    async with httpx.AsyncClient(timeout=timeout, trust_env=True) as client:
+        response = await client.get(
+            http_url(settings.server_url, path),
+            headers={"Authorization": f"Bearer {settings.server_key or ''}"},
+        )
+    if response.status_code in (401, 403):
+        raise RemoteError("the server rejected the API key", kind="authentication")
+    response.raise_for_status()
+    value = response.json()
+    return value if isinstance(value, dict) else {"value": value}
+
+
+__all__ = ["RemoteRuntime", "AgentEvent", "classify", "server_get"]