feat(server): API key authentication for collector (#130)

AGENTS.md

@@ -48,10 +48,10 @@ python -m pytest tests/test_watch.py -v
 
 ## Version bumping
 
-Current version: `0.51.1` in `src/agent_trace/__init__.py`.
+Current version: `0.52.0` in `src/agent_trace/__init__.py`.
 
-- New feature (new command, new flag, new integration): bump minor (`0.51.1` → `0.52.0`)
-- Bug fix or small improvement: bump patch (`0.51.1` → `0.51.2`)
+- New feature (new command, new flag, new integration): bump minor (`0.52.0` → `0.53.0`)
+- Bug fix or small improvement: bump patch (`0.52.0` → `0.52.1`)
 - Breaking change to CLI or storage format: bump major — check with maintainer first
 
 ## docs/ structure

docs/commands.md

@@ -336,10 +336,20 @@ Export sessions as JSONL for eval datasets. Compatible with LangSmith, Braintrus
 
 ### `server`
 ```
-agent-strace server [--port N] [--storage DIR]
+agent-strace server [--port N] [--host HOST] [--storage DIR] [--auth-key KEY]
+agent-strace server keygen
 ```
 Start a server-side event collector. See [server.md](server.md).
 
+| Flag | Description |
+|---|---|
+| `--port N` | Port to listen on (default: 4317) |
+| `--host HOST` | Host to bind to (default: 0.0.0.0) |
+| `--storage DIR` | Trace storage directory (default: `$AGENT_STRACE_STORAGE` or `.agent-traces`) |
+| `--auth-key KEY` | Require `Authorization: Bearer KEY` on all requests (also read from `AGENT_STRACE_AUTH_KEY`) |
+
+`keygen` prints a new `ast_`-prefixed API key to stdout. Set `AGENT_STRACE_AUTH_KEY` on the client side to inject the header automatically into all outbound collector requests.
+
 ### `auto`
 ```
 agent-strace auto [--framework NAME] [--detect] -- <command>

docs/server.md

@@ -58,9 +58,40 @@ When multiple agents send to the same collector, sessions are linked via `parent
 
 ---
 
-## Security note
+## Authentication
 
-No authentication in v1 — intended for internal/private network use. Add a reverse proxy (nginx, Caddy) for auth and TLS.
+By default the server runs unauthenticated (local use). For any network-accessible deployment, enable API key auth:
+
+```bash
+# Generate a key
+agent-strace server keygen
+# → ast_a3f9b2c1d4e5f6a7b8c9d0e1f2a3b4c5
+
+# Start server with key enforcement
+agent-strace server --auth-key ast_a3f9b2c1d4e5f6a7b8c9d0e1f2a3b4c5
+
+# Or via environment variable
+AGENT_STRACE_AUTH_KEY=ast_a3f9b2c1d4e5f6a7b8c9d0e1f2a3b4c5 agent-strace server
+```
+
+Requests without a matching `Authorization: Bearer <key>` header receive `401 Unauthorized`.
+
+**Client side** — set `AGENT_STRACE_AUTH_KEY` alongside `AGENT_STRACE_ENDPOINT` and all outbound requests include the header automatically:
+
+```bash
+export AGENT_STRACE_ENDPOINT=https://collector.example.com
+export AGENT_STRACE_AUTH_KEY=ast_a3f9b2c1d4e5f6a7b8c9d0e1f2a3b4c5
+python my_agent.py
+```
+
+The `--stream-headers` flag on `agent-strace watch` also works for one-off overrides:
+
+```bash
+agent-strace watch --stream-to https://collector.example.com \
+  --stream-headers "Authorization=Bearer ast_..."
+```
+
+Key format: `ast_` prefix + 32 hex characters. Generated with `secrets.token_hex(16)` — no new dependencies.
 
 ---
 

src/agent_trace/__init__.py

@@ -1,3 +1,3 @@
 """agent-trace: strace for AI agents."""
 
-__version__ = "0.51.1"
+__version__ = "0.52.0"

src/agent_trace/cli.py

@@ -875,6 +875,11 @@ def build_parser() -> argparse.ArgumentParser:
     p_server.add_argument("--storage", metavar="DIR",
                           help="storage directory for traces "
                                "(default: $AGENT_STRACE_STORAGE or .agent-traces)")
+    p_server.add_argument("--auth-key", metavar="KEY", dest="auth_key",
+                          help="require Authorization: Bearer <KEY> on all requests "
+                               "(also read from AGENT_STRACE_AUTH_KEY env var)")
+    p_server_sub = p_server.add_subparsers(dest="server_subcommand")
+    p_server_sub.add_parser("keygen", help="generate a new ast_-prefixed API key")
 
     # sample
     p_sample = sub.add_parser(

src/agent_trace/server.py

@@ -14,11 +14,20 @@
 Usage:
     agent-strace server --port 4317 --storage ./traces
 
+    # With API key authentication
+    agent-strace server keygen                          # generate a key
+    agent-strace server --auth-key ast_<key>            # enforce auth
+    AGENT_STRACE_AUTH_KEY=ast_<key> agent-strace server # via env var
+
 Agents point to it via environment variable:
     AGENT_STRACE_ENDPOINT=http://collector:4317 python my_agent.py
 
-No authentication in v1 — intended for internal/private network use.
-Add a reverse proxy (nginx, Caddy) for auth.
+    # With auth key on the client side
+    AGENT_STRACE_ENDPOINT=https://collector.example.com
+    AGENT_STRACE_AUTH_KEY=ast_<key>
+
+Without --auth-key / AGENT_STRACE_AUTH_KEY the server runs unauthenticated
+(original behaviour, unchanged for local use).
 
 See ADR-0012 for architecture decisions.
 """
@@ -28,6 +37,7 @@
 import argparse
 import json
 import os
+import secrets
 import sys
 import threading
 import time
@@ -40,27 +50,46 @@
 from .store import TraceStore, DEFAULT_TRACE_DIR
 
 
+# ---------------------------------------------------------------------------
+# Key generation
+# ---------------------------------------------------------------------------
+
+KEY_PREFIX = "ast_"
+
+
+def generate_api_key() -> str:
+    """Return a new ``ast_``-prefixed API key using secrets.token_hex."""
+    return KEY_PREFIX + secrets.token_hex(16)
+
+
 # ---------------------------------------------------------------------------
 # Remote event sender (used by hooks when AGENT_STRACE_ENDPOINT is set)
 # ---------------------------------------------------------------------------
 
+def _auth_headers() -> dict[str, str]:
+    """Return Authorization header dict if AGENT_STRACE_AUTH_KEY is set."""
+    key = os.environ.get("AGENT_STRACE_AUTH_KEY", "").strip()
+    if key:
+        return {"Authorization": f"Bearer {key}"}
+    return {}
+
+
 def send_event_to_endpoint(event: TraceEvent, endpoint: str) -> bool:
     """POST a single event to a remote collector.
 
     Returns True on success. Failures are logged to stderr but never raise —
     the hook must not block the agent.
+
+    When AGENT_STRACE_AUTH_KEY is set, injects ``Authorization: Bearer``
+    automatically.
     """
     import urllib.request
     import urllib.error
 
     url = endpoint.rstrip("/") + "/events"
     body = (event.to_json() + "\n").encode("utf-8")
-    req = urllib.request.Request(
-        url,
-        data=body,
-        headers={"Content-Type": "application/x-ndjson"},
-        method="POST",
-    )
+    headers = {"Content-Type": "application/x-ndjson", **_auth_headers()}
+    req = urllib.request.Request(url, data=body, headers=headers, method="POST")
     try:
         with urllib.request.urlopen(req, timeout=5) as resp:
             return resp.status in (200, 202)
@@ -70,18 +99,18 @@ def send_event_to_endpoint(event: TraceEvent, endpoint: str) -> bool:
 
 
 def send_session_meta_to_endpoint(meta: SessionMeta, endpoint: str) -> bool:
-    """POST session metadata to a remote collector."""
+    """POST session metadata to a remote collector.
+
+    When AGENT_STRACE_AUTH_KEY is set, injects ``Authorization: Bearer``
+    automatically.
+    """
     import urllib.request
     import urllib.error
 
     url = endpoint.rstrip("/") + "/sessions"
     body = meta.to_json().encode("utf-8")
-    req = urllib.request.Request(
-        url,
-        data=body,
-        headers={"Content-Type": "application/json"},
-        method="POST",
-    )
+    headers = {"Content-Type": "application/json", **_auth_headers()}
+    req = urllib.request.Request(url, data=body, headers=headers, method="POST")
     try:
         with urllib.request.urlopen(req, timeout=5) as resp:
             return resp.status in (200, 202)
@@ -100,11 +129,19 @@ class CollectorHandler(BaseHTTPRequestHandler):
     # Injected by the server setup
     store: TraceStore
     _lock: threading.Lock
+    _auth_key: str  # empty string = no auth required
 
     def log_message(self, fmt: str, *args: Any) -> None:
         # Suppress default access log; write to stderr with our prefix
         sys.stderr.write(f"[server] {fmt % args}\n")
 
+    def _check_auth(self) -> bool:
+        """Return True if the request is authorised (or auth is disabled)."""
+        if not self._auth_key:
+            return True
+        auth = self.headers.get("Authorization", "")
+        return auth == f"Bearer {self._auth_key}"
+
     def _send_json(self, status: int, data: dict) -> None:
         body = json.dumps(data).encode("utf-8")
         self.send_response(status)
@@ -132,6 +169,9 @@ def _read_body(self) -> bytes:
     # ------------------------------------------------------------------
 
     def do_GET(self) -> None:
+        if not self._check_auth():
+            self._send_json(401, {"error": "Unauthorized"})
+            return
         path = urlparse(self.path).path.rstrip("/")
 
         if path == "/health":
@@ -180,6 +220,9 @@ def do_GET(self) -> None:
     # ------------------------------------------------------------------
 
     def do_POST(self) -> None:
+        if not self._check_auth():
+            self._send_json(401, {"error": "Unauthorized"})
+            return
         path = urlparse(self.path).path.rstrip("/")
         body = self._read_body()
 
@@ -248,12 +291,13 @@ def _handle_post_sessions(self, body: bytes) -> None:
             self._send_json(400, {"error": str(exc)})
 
 
-def _make_handler(store: TraceStore, lock: threading.Lock) -> type:
-    """Return a CollectorHandler subclass with store and lock injected."""
+def _make_handler(store: TraceStore, lock: threading.Lock, auth_key: str = "") -> type:
+    """Return a CollectorHandler subclass with store, lock, and auth key injected."""
     class Handler(CollectorHandler):
         pass
     Handler.store = store
     Handler._lock = lock
+    Handler._auth_key = auth_key
     return Handler
 
 
@@ -265,15 +309,21 @@ def run_server(
     port: int,
     storage_dir: str,
     host: str = "0.0.0.0",
+    auth_key: str = "",
 ) -> None:
-    """Start the collector server and block until interrupted."""
+    """Start the collector server and block until interrupted.
+
+    When *auth_key* is non-empty, all requests must include
+    ``Authorization: Bearer <auth_key>`` or receive 401.
+    """
     store = TraceStore(storage_dir)
     lock = threading.Lock()
-    handler_class = _make_handler(store, lock)
+    handler_class = _make_handler(store, lock, auth_key=auth_key)
 
     server = HTTPServer((host, port), handler_class)
+    auth_note = " (auth enabled)" if auth_key else " (no auth)"
     sys.stderr.write(
-        f"[agent-strace server] listening on {host}:{port}\n"
+        f"[agent-strace server] listening on {host}:{port}{auth_note}\n"
         f"[agent-strace server] storage: {Path(storage_dir).resolve()}\n"
         f"[agent-strace server] health: http://{host}:{port}/health\n"
     )
@@ -290,10 +340,20 @@ def run_server(
 # ---------------------------------------------------------------------------
 
 def cmd_server(args: argparse.Namespace) -> int:
+    subcommand = getattr(args, "server_subcommand", None)
+
+    if subcommand == "keygen":
+        sys.stdout.write(generate_api_key() + "\n")
+        return 0
+
     port = getattr(args, "port", 4317)
     storage = getattr(args, "storage", None) or os.environ.get(
         "AGENT_STRACE_STORAGE", DEFAULT_TRACE_DIR
     )
     host = getattr(args, "host", "0.0.0.0")
-    run_server(port=port, storage_dir=storage, host=host)
+    auth_key = (
+        getattr(args, "auth_key", None)
+        or os.environ.get("AGENT_STRACE_AUTH_KEY", "")
+    ).strip()
+    run_server(port=port, storage_dir=storage, host=host, auth_key=auth_key)
     return 0