baton-proxy 0.1.0__py3-none-any.whl

baton_proxy/__init__.py

@@ -0,0 +1,9 @@
+"""baton-proxy — subprocess-wrap MCP proxy.
+
+Wraps a stdio MCP server, injects an annotation tool into the handshake,
+and emits friction events to a baton-console.
+
+See README.md for usage.
+"""
+
+__version__ = "0.1.0"

baton_proxy/__main__.py

@@ -0,0 +1,8 @@
+"""Entry point for `python -m baton_proxy` and the `baton-proxy` console script."""
+
+from __future__ import annotations
+
+from baton_proxy.proxy import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

baton_proxy/config.py

@@ -0,0 +1,100 @@
+"""Runtime configuration — read from environment variables once at startup.
+
+Subprocess-wrap deployment is 1-process-per-user (Claude Desktop / Claude Code
+spawns one proxy per MCP server entry), so a static per-process token model is
+fine. Hosted-HTTP deployment will need a per-request resolver; not in scope here.
+
+Zero-config defaults
+--------------------
+
+The proxy is meant to be install-and-play: add ``baton-proxy --`` in front of
+any MCP server, restart, and you get a stream of friction events in
+``/tmp/baton-proxy.jsonl`` (and on stderr). No env vars required. The
+defaults are deliberately placeholder-flavoured (``"local"``) so that the
+upgrade to a real Console is forced to be explicit.
+
+When ``BATON_EVENT_SINK`` resolves to an http(s):// sink, the emitter
+refuses to start while ``BATON_CONSENT_TOKEN`` is still the placeholder —
+placeholder-tagged events must never leak to a remote collector.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from dataclasses import dataclass
+
+# Zero-config defaults. Multi-sink (stderr + local file) so the events are
+# immediately visible both as a live stream and as a persistent log; tenant
+# and consent default to a sentinel ``"local"`` to make it obvious in any
+# downstream system that the install hasn't been wired to a real Console yet.
+DEFAULT_EVENT_SINK = "stderr:,file:///tmp/baton-proxy.jsonl"
+DEFAULT_TENANT_ID = "local"
+DEFAULT_CONSENT_TOKEN = "local"
+
+
+@dataclass(frozen=True)
+class Config:
+    """All runtime knobs. Created via Config.from_env()."""
+
+    # Process-lifetime session identifier per SPEC §11.4. Every event the proxy
+    # emits during this process shares this session_id.
+    session_id: str
+
+    # Where emitted events go. A URL whose scheme selects the sink:
+    #   https://console.example.com  -> HTTP POST to {url}/v0/events
+    #   file:///tmp/events.jsonl     -> append-JSONL to the local path
+    #   stderr:                      -> JSONL to stderr
+    # Comma-separated values fan out (MultiSink). Defaults via from_env() to
+    # ``DEFAULT_EVENT_SINK`` (stderr + local file). None disables emission —
+    # only test code sets this to None directly; ``from_env()`` always
+    # returns a populated value.
+    event_sink: str | None
+    tenant_id: str | None
+    # Only required for http(s) sinks; ignored for file/stderr sinks. The
+    # HTTP sink raises at startup if event_sink is http(s):// and this is None.
+    api_key: str | None
+    consent_token: str | None
+
+    # Vendor identifier surfaced in proxy logs and the annotation tool's
+    # namespace prefix (``{vendor_id}_annotate``). Optional; falls back to
+    # the generic ``vendor_annotate`` name.
+    vendor_id: str | None
+
+    # Where the proxy writes its own operational log. Stderr by default;
+    # override with BATON_PROXY_LOG_FILE for persistent debugging.
+    log_file: str | None
+
+    @property
+    def emission_enabled(self) -> bool:
+        """True when the envelope-essential fields are populated. With
+        ``from_env()`` defaults this is always True; only test code that
+        passes ``event_sink=None`` etc. directly will see False."""
+        return all(
+            v is not None
+            for v in (self.event_sink, self.tenant_id, self.consent_token)
+        )
+
+    @property
+    def using_placeholder_consent(self) -> bool:
+        """True when consent_token is still the install-time placeholder.
+        Emitter refuses to start an http(s) sink while this is True — a
+        placeholder consent token must never reach a remote collector."""
+        return self.consent_token == DEFAULT_CONSENT_TOKEN
+
+    @classmethod
+    def from_env(cls) -> Config:
+        return cls(
+            session_id=str(uuid.uuid4()),
+            event_sink=_env("BATON_EVENT_SINK") or DEFAULT_EVENT_SINK,
+            tenant_id=_env("BATON_TENANT_ID") or DEFAULT_TENANT_ID,
+            api_key=_env("BATON_API_KEY"),
+            consent_token=_env("BATON_CONSENT_TOKEN") or DEFAULT_CONSENT_TOKEN,
+            vendor_id=_env("BATON_VENDOR_ID"),
+            log_file=_env("BATON_PROXY_LOG_FILE"),
+        )
+
+
+def _env(name: str) -> str | None:
+    v = os.environ.get(name)
+    return v if v else None

baton_proxy/emitter.py

@@ -0,0 +1,306 @@
+"""Async friction-event emitter.
+
+The proxy intercepts MCP traffic on the hot path (every `tools/call`). Doing
+a synchronous network call from that thread would add the full ingest
+round-trip (~50-200ms) to every tool call. Trust pattern: sub-ms overhead.
+So emission is queued and drained on a background thread; the hot path
+only pays an `enqueue()`.
+
+Failure mode: the background thread logs and drops on sink failures. A
+backed-up or dead emitter must NEVER block proxy I/O — that's the
+fail-open contract. Queue is bounded; overflow drops the oldest event
+and logs once per 100 drops.
+
+Where events go is the Sink's job (sinks.py). The Emitter just enqueues,
+drains, and hands each event to ``self._sink.write(event)``. Sink is built
+once at start() from ``BATON_EVENT_SINK`` (URL-driven, comma-separated
+list builds a MultiSink); misconfig (unsupported scheme, http without
+api_key) raises at start() — never a silent no-emit.
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+import time
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any
+
+from baton_proxy.config import Config
+from baton_proxy.sinks import Sink, make_sink
+
+logger = logging.getLogger(__name__)
+
+# Bounded queue — backed-up emitter shouldn't accumulate unbounded memory.
+# 1000 events buys a decent buffer for typical 5-10 RPS tool-call workloads.
+_QUEUE_MAXSIZE = 1000
+
+_SDK_VERSION = "baton-proxy/0.0.1"
+_AGENT_RUNTIME = "mcp-proxy"
+
+
+@dataclass(frozen=True)
+class _Event:
+    """Wire envelope, mirrors baton-console IncomingEvent shape.
+
+    Schemas are mirrored rather than imported so the proxy isn't lock-stepped
+    to a baton-console release. The console accepts `spec_version: str = "0.1"`
+    with a default and `extra="forbid"` on everything else.
+    """
+
+    event_id: str
+    event_type: str
+    session_id: str
+    sequence_number: int
+    captured_at: str
+    tenant_id: str
+    consent_token: str
+    sdk_version: str
+    agent_runtime: str
+    payload: dict[str, Any]
+    runtime_meta: dict[str, Any] | None = None
+
+    def to_json(self) -> dict[str, Any]:
+        d: dict[str, Any] = {
+            "event_id": self.event_id,
+            "event_type": self.event_type,
+            "session_id": self.session_id,
+            "sequence_number": self.sequence_number,
+            "captured_at": self.captured_at,
+            "tenant_id": self.tenant_id,
+            "consent_token": self.consent_token,
+            "sdk_version": self.sdk_version,
+            "agent_runtime": self.agent_runtime,
+            "payload": self.payload,
+        }
+        if self.runtime_meta is not None:
+            d["runtime_meta"] = self.runtime_meta
+        return d
+
+
+class Emitter:
+    """Background-thread emitter. Construct, call .start(), enqueue from any
+    thread, and call .stop() at shutdown.
+
+    When `config.emission_enabled` is False, .start() / .enqueue_*() are no-ops
+    so callers don't need to branch.
+    """
+
+    def __init__(self, config: Config) -> None:
+        self._config = config
+        self._queue: queue.Queue[_Event | None] = queue.Queue(maxsize=_QUEUE_MAXSIZE)
+        self._thread: threading.Thread | None = None
+        self._seq = 0
+        self._seq_lock = threading.Lock()
+        # Serialises put_nowait across producers. queue.Queue's internal mutex
+        # guards individual operations but not a get+put pair, so an unguarded
+        # drop-oldest sequence has a window where another producer can refill
+        # the queue between our get and put.
+        self._enqueue_lock = threading.Lock()
+        self._drop_count = 0
+        # Sink set up in start(); None until then.
+        self._sink: Sink | None = None
+
+    def start(self) -> None:
+        if not self._config.emission_enabled:
+            return
+        if self._thread is not None:
+            return
+        assert self._config.event_sink is not None  # emission_enabled gates this
+        self._guard_remote_consent()
+        self._sink = make_sink(self._config.event_sink, api_key=self._config.api_key)
+        self._thread = threading.Thread(target=self._drain, name="baton-proxy-emitter", daemon=True)
+        self._thread.start()
+
+    def _guard_remote_consent(self) -> None:
+        """Refuse to ship events to a remote sink while the consent token is
+        still the install-time placeholder. Local file/stderr sinks are
+        always OK — the placeholder just marks "this install hasn't been
+        wired to a real Console yet". The check runs before sink
+        construction so a misconfigured install fails loudly at startup
+        instead of silently leaking placeholder-tagged events.
+        """
+        if not self._config.using_placeholder_consent:
+            return
+        assert self._config.event_sink is not None
+        parts = [p.strip() for p in self._config.event_sink.split(",") if p.strip()]
+        if any(p.startswith(("http://", "https://")) for p in parts):
+            raise ValueError(
+                "Refusing to ship events to an http(s) sink with placeholder "
+                "BATON_CONSENT_TOKEN='local' — set BATON_CONSENT_TOKEN to the "
+                "real per-install consent token before pointing at a Console."
+            )
+
+    def stop(self, timeout: float = 2.0) -> None:
+        if self._thread is None:
+            return
+        # Blocking put with timeout — if the queue is full, put_nowait would
+        # silently drop the sentinel and the drain thread would loop until
+        # daemon-killed at process exit (losing buffered events). put() waits
+        # for the drain thread to free a slot, which it does once per second.
+        try:
+            self._queue.put(None, timeout=timeout)
+        except queue.Full:
+            # Drain thread is dead or wedged; nothing more we can do here.
+            pass
+        self._thread.join(timeout=timeout)
+        self._thread = None
+        if self._sink is not None:
+            self._sink.close()
+            self._sink = None
+
+    def enqueue_tool_call_start(
+        self,
+        *,
+        tool_name: str,
+        params: Mapping[str, Any] | None,
+        runtime_meta: Mapping[str, Any] | None = None,
+    ) -> None:
+        self._enqueue(
+            event_type="tool_call_start",
+            payload={"tool_name": tool_name, "params": dict(params) if params else {}},
+            runtime_meta=dict(runtime_meta) if runtime_meta else None,
+        )
+
+    def enqueue_tool_call_end(
+        self,
+        *,
+        tool_name: str,
+        result: Any,
+        duration_ms: int,
+        runtime_meta: Mapping[str, Any] | None = None,
+    ) -> None:
+        self._enqueue(
+            event_type="tool_call_end",
+            payload={"tool_name": tool_name, "result": result, "duration_ms": duration_ms},
+            runtime_meta=dict(runtime_meta) if runtime_meta else None,
+        )
+
+    def enqueue_tool_call_error(
+        self,
+        *,
+        tool_name: str,
+        error_type: str,
+        error_body: str,
+        duration_ms: int,
+        runtime_meta: Mapping[str, Any] | None = None,
+    ) -> None:
+        self._enqueue(
+            event_type="tool_call_error",
+            payload={
+                "tool_name": tool_name,
+                "error_type": error_type,
+                "error_body": error_body,
+                "duration_ms": duration_ms,
+            },
+            runtime_meta=dict(runtime_meta) if runtime_meta else None,
+        )
+
+    def enqueue_annotation(
+        self,
+        *,
+        signal_type: str | None,
+        intent: str | None,
+        suggested_improvement: str | None,
+        expected_outcome: str | None = None,
+        workflow: str | None = None,
+        context: Mapping[str, Any] | None = None,
+        runtime_meta: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Annotation event per SPEC §11.4; nullable keys omitted when None."""
+        candidates: dict[str, Any] = {
+            "signal_type": signal_type,
+            "intent": intent,
+            "suggested_improvement": suggested_improvement,
+            "expected_outcome": expected_outcome,
+            "workflow": workflow,
+            "context": dict(context) if context is not None else None,
+        }
+        payload = {k: v for k, v in candidates.items() if v is not None}
+        self._enqueue(
+            event_type="annotation",
+            payload=payload,
+            runtime_meta=dict(runtime_meta) if runtime_meta else None,
+        )
+
+    def _enqueue(
+        self,
+        *,
+        event_type: str,
+        payload: dict[str, Any],
+        runtime_meta: dict[str, Any] | None,
+    ) -> None:
+        if not self._config.emission_enabled or self._thread is None:
+            return
+
+        with self._seq_lock:
+            seq = self._seq
+            self._seq += 1
+
+        event = _Event(
+            event_id=str(uuid.uuid4()),
+            event_type=event_type,
+            session_id=self._config.session_id,
+            sequence_number=seq,
+            captured_at=datetime.now(UTC).isoformat().replace("+00:00", "Z"),
+            tenant_id=self._config.tenant_id,  # type: ignore[arg-type]
+            consent_token=self._config.consent_token,  # type: ignore[arg-type]
+            sdk_version=_SDK_VERSION,
+            agent_runtime=_AGENT_RUNTIME,
+            payload=payload,
+            runtime_meta=runtime_meta,
+        )
+
+        with self._enqueue_lock:
+            try:
+                self._queue.put_nowait(event)
+            except queue.Full:
+                # Drop-oldest. Held under _enqueue_lock so the get+put pair
+                # is atomic w.r.t. other producers; without it a concurrent
+                # put_nowait could refill the slot between our get and put
+                # and silently drop the new event instead of the oldest.
+                self._drop_count += 1
+                try:
+                    self._queue.get_nowait()
+                except queue.Empty:
+                    pass
+                try:
+                    self._queue.put_nowait(event)
+                except queue.Full:
+                    pass
+                if self._drop_count % 100 == 1:
+                    logger.warning(
+                        "baton-proxy emitter queue full, dropped %d events", self._drop_count
+                    )
+
+    def _drain(self) -> None:
+        while True:
+            try:
+                event = self._queue.get(timeout=1.0)
+            except queue.Empty:
+                continue
+            if event is None:
+                return
+            self._deliver(event)
+
+    def _deliver(self, event: _Event) -> None:
+        """Hand one event to the sink. Any failure is logged and dropped —
+        fail-open contract: a broken sink must not stall the drain loop or
+        propagate exceptions that would kill the daemon thread."""
+        assert self._sink is not None  # start() built it
+        try:
+            self._sink.write(event.to_json())
+        except Exception as e:  # noqa: BLE001 — fail-open at delivery boundary
+            logger.warning(
+                "baton-proxy emit %s -> %s: %s", event.event_type, type(e).__name__, e
+            )
+
+
+def utc_now_ms() -> int:
+    """Monotonic-ish millisecond clock for duration math. time.monotonic()
+    gives a relative clock; multiply to ms."""
+    return int(time.monotonic() * 1000)