From 7d147d3d15fd17c21f5fc3f19b23ecb06c471cd3 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Thu, 2 Jul 2026 15:19:44 +0800 Subject: [PATCH 1/2] =?UTF-8?q?abridge:=20direct=20mode=20=E2=80=94=20serv?= =?UTF-8?q?e=20@on=20clients=20as=20a=20standalone=20HTTP=20service?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tunnel exists for sandboxes with no egress: LLM traffic rides the runtime's Socket.IO connection at the cost of two extra hops and the host process fanning in every rollout's calls. When the sandbox can reach the model-serving network (in-cluster vLLM/SGLang), that detour is pure overhead — serve the same @on(path) handlers next to the engine instead and point the agent's SDK straight at it. - build_app(*clients): one POST route per handler, tunnel-identical contract (JSON-object bodies, ClientResponse out, in-band errors) - build_session_app(factory): per-caller sessions — the API key each agent sends (x-api-key / Bearer) hashes to a session id, the factory builds that caller's client, an LRU caps live clients and acloses evictions. Mint a key per rollout and one server groups them all; agent keys are never forwarded upstream. - agentix-bridge-serve: Anthropic-speaking front for an OpenAI-compatible engine (lazy openai extra, loopback bind default) Multi-backend routing and token capture stay out of scope — that is the full gateway's job; this is the tunnel without the tunnel, and the tunnel remains the mode for egress-less sandboxes. Co-Authored-By: Claude Fable 5 --- plugins/abridge/README.md | 29 +++ plugins/abridge/agentix/bridge/serve.py | 250 ++++++++++++++++++++++++ plugins/abridge/pyproject.toml | 6 + plugins/abridge/tests/test_serve.py | 136 +++++++++++++ 4 files changed, 421 insertions(+) create mode 100644 plugins/abridge/agentix/bridge/serve.py create mode 100644 plugins/abridge/tests/test_serve.py diff --git a/plugins/abridge/README.md b/plugins/abridge/README.md index 795aa24..31642a0 100644 --- a/plugins/abridge/README.md +++ b/plugins/abridge/README.md @@ -198,11 +198,40 @@ completion logs (elapsed-ms, status code). Wire-level errors come from `trace.Processor` (e.g. `agentix.plugins.trace-otel`) to export to LangSmith / Langfuse / Datadog / any OTel backend. +## Direct mode — skip the tunnel when the sandbox has network reach + +The tunnel exists for sandboxes with **no egress at all**: LLM traffic +piggybacks on the runtime's Socket.IO connection, at the cost of two +extra hops and the host process fanning in every concurrent rollout's +calls. When the sandbox *can* reach the model-serving network (an +in-cluster vLLM/SGLang, a private gateway), serve the same handlers as +a standalone HTTP service next to the engine and point the agent +straight at it — the host stays out of the data path: + +```bash +OPENAI_API_KEY=EMPTY agentix-bridge-serve \ + --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b +# agent side: ANTHROPIC_BASE_URL=http://:8399 ANTHROPIC_API_KEY= +``` + +Rollout identity travels in the key: mint a fresh placeholder API key +per rollout (the key you already inject into the sandbox), and the +server hashes whatever key each request carries into the +`x-session-id` it stamps upstream — one server groups any number of +concurrent rollouts, agent keys are never forwarded, and the real +upstream key stays server-side. Programmatic surface: +`build_app(*clients)` (shared session) and +`build_session_app(factory)` (one client per caller key, LRU-bounded) +in `agentix.bridge.serve`. Multi-backend routing and token capture +belong to the full gateway (see the roadmap); the tunnel remains the +mode for fully egress-less sandboxes. + ## Module layout ``` agentix/bridge/ ├── proxy.py # Proxy + @on + sandbox tunnel + wire types +├── serve.py # direct mode: @on handlers as a standalone HTTP service ├── forward.py # JSON POST forwarding to a host-side service ├── sidecar.py # local process lifecycle + health supervision └── clients/ # bundled handler implementations diff --git a/plugins/abridge/agentix/bridge/serve.py b/plugins/abridge/agentix/bridge/serve.py new file mode 100644 index 0000000..0251320 --- /dev/null +++ b/plugins/abridge/agentix/bridge/serve.py @@ -0,0 +1,250 @@ +"""Serve abridge clients over plain HTTP — direct mode, no tunnel. + +The `Proxy` tunnels LLM traffic through the host because a sandbox may +have no network egress at all. When the sandbox *can* reach the model +serving network (an in-cluster vLLM, a private gateway), routing every +call through the host adds two hops and one Python process as the +fan-in for every concurrent rollout. This module serves the same +`@on(path)` handler objects as a standalone HTTP service instead: +deploy it next to the engine and point the agent's SDK straight at it — +the host stays out of the data path. + + agent (in sandbox) ──HTTP──▶ agentix-bridge-serve ──▶ engine + translation + identity stamping + +Rollout identity without the host: mint a fresh placeholder API key per +rollout (the key you already inject into the sandbox) and the server +derives `x-session-id` from a hash of whatever key each request +carries — one server groups any number of concurrent rollouts, and the +minting side can compute the same hash to correlate. Agent keys are +treated as identity, never forwarded upstream; the real upstream key +stays on this server. Anything beyond grouping (multi-backend routing, +token capture) is the full gateway's job — this is deliberately just +"the tunnel without the tunnel", and the tunnel remains the mode for +sandboxes with no egress. + + OPENAI_API_KEY=EMPTY agentix-bridge-serve \ + --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import logging +import os +from collections import OrderedDict +from collections.abc import Awaitable, Callable + +import uvicorn +from fastapi import FastAPI +from fastapi import Request as FastAPIRequest +from fastapi.responses import JSONResponse, Response + +from .proxy import AbridgeError, Client, Handler, Request, _AsyncCloseable, _collect_handlers + +logger = logging.getLogger(__name__) + +__all__ = ["build_app", "build_session_app", "main"] + +# `resolve(request) -> Handler`: fixed in `build_app`, per-caller in +# `build_session_app`. +Resolver = Callable[[FastAPIRequest], Awaitable[Handler]] + + +def build_app(*clients: Client) -> FastAPI: + """A FastAPI app with one POST route per `@on(path)` handler. + + Same request/response contract as the sandbox tunnel: JSON-object + bodies in, the handler's `ClientResponse` out, handler errors as + JSON error bodies (`AbridgeError` keeps its status; anything else + is a 502). All requests share the clients' sessions — for + per-caller sessions use `build_session_app`. + """ + handlers: dict[str, Handler] = {} + for client in clients: + for path, method in _collect_handlers(client).items(): + if path in handlers: + raise ValueError(f"two clients register the same @on path {path!r}") + handlers[path] = method + if not handlers: + raise ValueError("no @on-decorated handlers found on any client passed to build_app(...)") + + app = FastAPI() + app.get("/_health")(_health) + for path, method in handlers.items(): + app.post(path)(_make_endpoint(path, _fixed_resolver(method))) + return app + + +def build_session_app( + client_factory: Callable[[str], Client], + *, + max_sessions: int = 256, +) -> FastAPI: + """Like `build_app`, but with one client per caller identity. + + Each request's API key (`x-api-key`, else `Authorization: Bearer`) + hashes to a session id; `client_factory(session_id)` builds that + caller's client on first use. A capped LRU keeps live clients and + closes (via `aclose()`, when implemented) the least recent one past + `max_sessions` — size it above your rollout concurrency so an + in-flight caller is never evicted. Keyless requests share the + `"anonymous"` session. + """ + if max_sessions < 1: + raise ValueError("max_sessions must be >= 1") + + anonymous = client_factory("anonymous") + paths = tuple(_collect_handlers(anonymous)) + if not paths: + raise ValueError("client_factory built a client with no @on-decorated handlers") + + # session id → that caller's handler table (client kept alive by its + # bound methods; kept alongside for aclose on eviction). + sessions: OrderedDict[str, tuple[Client, dict[str, Handler]]] = OrderedDict() + sessions["anonymous"] = (anonymous, _collect_handlers(anonymous)) + lock = asyncio.Lock() + + def _session_resolver(path: str) -> Resolver: + async def resolve(request: FastAPIRequest) -> Handler: + session_id = _session_id(_caller_key(request)) + async with lock: + entry = sessions.get(session_id) + if entry is None: + client = client_factory(session_id) + entry = (client, _collect_handlers(client)) + sessions[session_id] = entry + sessions.move_to_end(session_id) + evicted = [] + while len(sessions) > max_sessions: + _, (old, _handlers) = sessions.popitem(last=False) + evicted.append(old) + for old in evicted: + await _close_client(old) + return entry[1][path] + + return resolve + + app = FastAPI() + app.get("/_health")(_health) + for path in paths: + app.post(path)(_make_endpoint(path, _session_resolver(path))) + return app + + +async def _health() -> dict[str, str]: + return {"status": "ok"} + + +def _fixed_resolver(method: Handler) -> Resolver: + async def resolve(_request: FastAPIRequest) -> Handler: + return method + + return resolve + + +def _caller_key(request: FastAPIRequest) -> str: + key = request.headers.get("x-api-key") + if key: + return key + auth = request.headers.get("authorization", "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() + return "" + + +def _session_id(caller_key: str) -> str: + if not caller_key: + return "anonymous" + # The key is identity, not a secret to echo around: hash it so logs + # and upstream `x-session-id` headers never carry the raw value. + return hashlib.sha256(caller_key.encode()).hexdigest()[:24] + + +async def _close_client(client: Client) -> None: + if isinstance(client, _AsyncCloseable): + try: + await client.aclose() + except Exception: # noqa: BLE001 - eviction must not fail the live request + logger.exception("abridge serve: failed to close evicted client") + + +def _make_endpoint(path: str, resolve: Resolver): + async def endpoint(request: FastAPIRequest) -> Response: + body = await _read_json(request) + try: + handler = await resolve(request) + response = await handler(Request(path=path, body=body)) + except AbridgeError as exc: + logger.warning("abridge serve %s: %s (status=%d)", path, exc.message, exc.status_code) + return JSONResponse({"error": {"message": exc.message}}, status_code=exc.status_code) + except Exception as exc: # noqa: BLE001 - any handler failure becomes a wire error + logger.exception("abridge serve %s: handler raised", path) + return JSONResponse({"error": {"message": f"{type(exc).__name__}: {exc}"}}, status_code=502) + return Response(content=response.body, media_type=response.media_type, status_code=response.status_code) + + return endpoint + + +async def _read_json(request: FastAPIRequest) -> dict: + try: + parsed = await request.json() + except ValueError: + return {} + return parsed if isinstance(parsed, dict) else {} + + +def main(argv: list[str] | None = None) -> None: + """`agentix-bridge-serve` — an Anthropic-speaking front for an + OpenAI-compatible engine, one session per caller key.""" + parser = argparse.ArgumentParser( + prog="agentix-bridge-serve", + description=( + "Serve the Anthropic->OpenAI translation next to an OpenAI-compatible " + "engine (vLLM, SGLang, a gateway). Agents point ANTHROPIC_BASE_URL at " + "this server; each distinct agent API key becomes its own session." + ), + ) + parser.add_argument( + "--upstream-base-url", + default=os.environ.get("OPENAI_BASE_URL"), + help="OpenAI-compatible endpoint, e.g. http://vllm:8000/v1 (env: OPENAI_BASE_URL)", + ) + parser.add_argument( + "--upstream-api-key", + default=os.environ.get("OPENAI_API_KEY", "EMPTY"), + help="real key for the upstream; never the keys agents send (env: OPENAI_API_KEY)", + ) + parser.add_argument( + "--upstream-model", + default=os.environ.get("UPSTREAM_MODEL"), + help="pin every upstream call to this model id (env: UPSTREAM_MODEL)", + ) + parser.add_argument("--host", default="127.0.0.1", help="bind address; expose beyond loopback deliberately") + parser.add_argument("--port", type=int, default=8399) + parser.add_argument("--upstream-timeout", type=float, default=180.0) + parser.add_argument("--max-sessions", type=int, default=256) + args = parser.parse_args(argv) + if not args.upstream_base_url: + parser.error("--upstream-base-url (or OPENAI_BASE_URL) is required") + + # Lazy: the translation client needs the `openai` extra. + from .clients import AnthropicFromOpenAIClient + + def factory(session_id: str) -> AnthropicFromOpenAIClient: + return AnthropicFromOpenAIClient( + base_url=args.upstream_base_url, + api_key=args.upstream_api_key, + model=args.upstream_model, + timeout=args.upstream_timeout, + session_id=session_id, + ) + + app = build_session_app(factory, max_sessions=args.max_sessions) + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() diff --git a/plugins/abridge/pyproject.toml b/plugins/abridge/pyproject.toml index 9a8ae66..3ffd119 100644 --- a/plugins/abridge/pyproject.toml +++ b/plugins/abridge/pyproject.toml @@ -15,6 +15,12 @@ dependencies = [ "uvicorn>=0.30", ] +# Direct mode: serve the Anthropic->OpenAI translation next to an +# OpenAI-compatible engine so sandboxes with network reach skip the +# host tunnel entirely (see agentix/bridge/serve.py). +[project.scripts] +agentix-bridge-serve = "agentix.bridge.serve:main" + # Bundled clients each depend on the provider's SDK; users who only # ship custom `@on` handlers don't need either. Pick what you actually # use: diff --git a/plugins/abridge/tests/test_serve.py b/plugins/abridge/tests/test_serve.py new file mode 100644 index 0000000..6b0222c --- /dev/null +++ b/plugins/abridge/tests/test_serve.py @@ -0,0 +1,136 @@ +"""Tests for `agentix.bridge.serve` — the tunnel-less direct mode. + +The HTTP contract must match the sandbox tunnel (JSON-object bodies, +`ClientResponse` out, in-band errors), and the session-keyed app must +map caller keys to stable, distinct, LRU-managed client sessions. +""" + +from __future__ import annotations + +from typing import Any + +from agentix.bridge import AbridgeError, ClientResponse, Request, on +from agentix.bridge.serve import _session_id, build_app, build_session_app +from fastapi.testclient import TestClient + + +class EchoClient: + def __init__(self, session_id: str = "fixed") -> None: + self.session_id = session_id + self.closed = False + + @on("/v1/echo") + async def echo(self, request: Request) -> ClientResponse: + return ClientResponse.json({"echo": request.body, "session": self.session_id}) + + @on("/v1/teapot") + async def teapot(self, request: Request) -> ClientResponse: + raise AbridgeError("short and stout", status_code=418) + + @on("/v1/boom") + async def boom(self, request: Request) -> ClientResponse: + raise RuntimeError("kaput") + + async def aclose(self) -> None: + self.closed = True + + +def test_build_app_serves_handlers_and_health() -> None: + tc = TestClient(build_app(EchoClient())) + assert tc.get("/_health").json() == {"status": "ok"} + r = tc.post("/v1/echo", json={"x": 1}) + assert r.status_code == 200 + assert r.json()["echo"] == {"x": 1} + assert tc.post("/nope", json={}).status_code == 404 + + +def test_handler_errors_become_wire_errors() -> None: + tc = TestClient(build_app(EchoClient())) + r = tc.post("/v1/teapot", json={}) + assert r.status_code == 418 + assert "short and stout" in r.json()["error"]["message"] + r = tc.post("/v1/boom", json={}) + assert r.status_code == 502 + assert "RuntimeError" in r.json()["error"]["message"] + + +def test_non_object_body_coerced_to_empty_like_the_tunnel() -> None: + tc = TestClient(build_app(EchoClient())) + r = tc.post("/v1/echo", content=b"[1, 2]", headers={"content-type": "application/json"}) + assert r.json()["echo"] == {} + + +def test_session_app_maps_keys_to_stable_distinct_sessions() -> None: + tc = TestClient(build_session_app(EchoClient)) + a1 = tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-a"}).json()["session"] + a2 = tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-a"}).json()["session"] + bearer = tc.post("/v1/echo", json={}, headers={"authorization": "Bearer sk-b"}).json()["session"] + anonymous = tc.post("/v1/echo", json={}).json()["session"] + + assert a1 == a2 == _session_id("sk-a") + assert bearer == _session_id("sk-b") + assert bearer != a1 + assert anonymous == "anonymous" + + +def test_session_app_evicts_and_closes_least_recent() -> None: + built: list[EchoClient] = [] + + def factory(session_id: str) -> EchoClient: + client = EchoClient(session_id) + built.append(client) + return client + + tc = TestClient(build_session_app(factory, max_sessions=1)) + tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-a"}) + tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-b"}) + + still_open = [client.session_id for client in built if not client.closed] + assert still_open == [_session_id("sk-b")] + assert [client.session_id for client in built if client.closed] == ["anonymous", _session_id("sk-a")] + + +def _mock_completion() -> Any: + from openai.types.chat import ChatCompletion + + return ChatCompletion.model_validate( + { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "upstream-model", + "choices": [ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}} + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}, + } + ) + + +def test_agent_keys_become_upstream_session_headers() -> None: + """End to end through the real translation client: the key each + agent sends is hashed into the `x-session-id` the upstream sees, and + the agent's key itself never reaches the upstream call.""" + from agentix.bridge.clients import AnthropicFromOpenAIClient + + seen: list[dict[str, str]] = [] + + def factory(session_id: str) -> AnthropicFromOpenAIClient: + client = AnthropicFromOpenAIClient(api_key="real-upstream-key", model="m", session_id=session_id) + + async def create(**kwargs: Any) -> Any: + seen.append(dict(kwargs["extra_headers"])) + return _mock_completion() + + client._client.chat.completions.create = create # type: ignore[method-assign] + return client + + tc = TestClient(build_session_app(factory)) + body = {"model": "claude", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]} + assert tc.post("/v1/messages", json=body, headers={"x-api-key": "rollout-1"}).status_code == 200 + assert tc.post("/v1/messages", json=body, headers={"x-api-key": "rollout-2"}).status_code == 200 + + assert seen[0]["x-session-id"] == _session_id("rollout-1") + assert seen[1]["x-session-id"] == _session_id("rollout-2") + assert seen[0]["x-session-id"] != seen[1]["x-session-id"] + assert all("rollout-1" not in v and "rollout-2" not in v for headers in seen for v in headers.values()) From c4432d3c655510e84981a6fe741f0ac4d0561e29 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Thu, 2 Jul 2026 15:41:36 +0800 Subject: [PATCH 2/2] abridge serve: in-flight-safe sessions, shutdown cleanup, opt-in auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adversarial review of the first cut confirmed three lifecycle holes and one trust-model gap; all fixed: - eviction is now in-flight-safe: the session table stays LRU-bounded, but an evicted client closes only when its last live request releases (previously a burst of fresh caller keys aclosed clients mid-upstream-call, turning valid LLM calls into 502s — also removes the eviction-thrash DoS lever) - a FastAPI lifespan closes the route-probe client at startup and drains every remaining session client at shutdown, matching the cleanup guarantee the tunnel's Proxy.session() gives - the client factory runs outside the table lock via asyncio.to_thread (SDK construction is ~30ms of synchronous SSL/pool setup that was stalling the whole event loop under the lock) - opt-in caller verification: build_session_app(verify_key=...) / --require-key-prefix gates requests to keys the harness minted (401 otherwise); trust model documented in the module and README, and the README example now shows the non-loopback bind it implies - session_id_for() is public and documented so the minting side can correlate upstream x-session-id values; build_app's docstring now states honestly that AbridgeError statuses survive here while the tunnel wire collapses them to 502 Co-Authored-By: Claude Fable 5 --- plugins/abridge/README.md | 36 +++- plugins/abridge/agentix/bridge/serve.py | 271 ++++++++++++++++++------ plugins/abridge/tests/test_serve.py | 95 +++++++-- 3 files changed, 311 insertions(+), 91 deletions(-) diff --git a/plugins/abridge/README.md b/plugins/abridge/README.md index 31642a0..fe09cb1 100644 --- a/plugins/abridge/README.md +++ b/plugins/abridge/README.md @@ -210,21 +210,35 @@ straight at it — the host stays out of the data path: ```bash OPENAI_API_KEY=EMPTY agentix-bridge-serve \ - --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b -# agent side: ANTHROPIC_BASE_URL=http://:8399 ANTHROPIC_API_KEY= + --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b \ + --host 0.0.0.0 --require-key-prefix rollout-secret- +# agent side: ANTHROPIC_BASE_URL=http://:8399 +# ANTHROPIC_API_KEY=rollout-secret- ``` Rollout identity travels in the key: mint a fresh placeholder API key per rollout (the key you already inject into the sandbox), and the -server hashes whatever key each request carries into the -`x-session-id` it stamps upstream — one server groups any number of -concurrent rollouts, agent keys are never forwarded, and the real -upstream key stays server-side. Programmatic surface: -`build_app(*clients)` (shared session) and -`build_session_app(factory)` (one client per caller key, LRU-bounded) -in `agentix.bridge.serve`. Multi-backend routing and token capture -belong to the full gateway (see the roadmap); the tunnel remains the -mode for fully egress-less sandboxes. +server maps whatever key each request carries to +`session_id_for(key)` — the `x-session-id` it stamps upstream. One +server groups any number of concurrent rollouts; the minting side +calls the same `agentix.bridge.serve.session_id_for` to correlate; +agent keys are never forwarded, and the real upstream key stays +server-side. + +Trust model: the server binds loopback by default and is +unauthenticated unless you gate it — sandboxes run model-generated +code, so when you expose it to them (`--host`), also set +`--require-key-prefix ` (or pass `verify_key=` to +`build_session_app`) so only keys your harness minted are served; +everything else gets a 401. + +Programmatic surface in `agentix.bridge.serve`: `build_app(*clients)` +(shared session) and `build_session_app(factory)` (one client per +caller key; LRU-bounded with in-flight-safe eviction — an evicted +client closes only after its live requests finish — and full cleanup +on shutdown). Multi-backend routing and token capture belong to the +full gateway (see the roadmap); the tunnel remains the mode for fully +egress-less sandboxes. ## Module layout diff --git a/plugins/abridge/agentix/bridge/serve.py b/plugins/abridge/agentix/bridge/serve.py index 0251320..3ffc789 100644 --- a/plugins/abridge/agentix/bridge/serve.py +++ b/plugins/abridge/agentix/bridge/serve.py @@ -14,17 +14,26 @@ Rollout identity without the host: mint a fresh placeholder API key per rollout (the key you already inject into the sandbox) and the server -derives `x-session-id` from a hash of whatever key each request -carries — one server groups any number of concurrent rollouts, and the -minting side can compute the same hash to correlate. Agent keys are -treated as identity, never forwarded upstream; the real upstream key -stays on this server. Anything beyond grouping (multi-backend routing, -token capture) is the full gateway's job — this is deliberately just -"the tunnel without the tunnel", and the tunnel remains the mode for -sandboxes with no egress. +maps whatever key each request carries to `session_id_for(key)` — one +server groups any number of concurrent rollouts, and the minting side +calls the same function to correlate. Agent keys are treated as +identity, never forwarded upstream; the real upstream key stays on +this server. + +Trust model: the server itself is unauthenticated by default and binds +loopback unless told otherwise — expose it only to the network segment +you already trust with the engine. Sandboxes run model-generated code; +when they share a network with the server, pass `verify_key=` (CLI: +`--require-key-prefix`) so only keys your harness minted are served and +everything else gets a 401. + +Anything beyond grouping (multi-backend routing, token capture) is the +full gateway's job — this is deliberately just "the tunnel without the +tunnel", and the tunnel remains the mode for sandboxes with no egress. OPENAI_API_KEY=EMPTY agentix-bridge-serve \ - --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b + --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b \ + --host 0.0.0.0 --require-key-prefix rollout-secret- """ from __future__ import annotations @@ -35,7 +44,9 @@ import logging import os from collections import OrderedDict -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass, field import uvicorn from fastapi import FastAPI @@ -46,21 +57,37 @@ logger = logging.getLogger(__name__) -__all__ = ["build_app", "build_session_app", "main"] +__all__ = ["build_app", "build_session_app", "main", "session_id_for"] -# `resolve(request) -> Handler`: fixed in `build_app`, per-caller in +# `resolve(request)` yields the bound `@on` method for the request — +# fixed in `build_app`, per-caller (with in-flight tracking) in # `build_session_app`. -Resolver = Callable[[FastAPIRequest], Awaitable[Handler]] +Resolver = Callable[[FastAPIRequest], AbstractAsyncContextManager[Handler]] + + +def session_id_for(caller_key: str) -> str: + """The session id the server derives from an agent's API key. + + `sha256(key)` hex truncated to 24 chars; the empty key maps to + `"anonymous"`. Public so the side minting per-rollout keys can + compute the same id and correlate upstream `x-session-id` values + (the raw key is never echoed into logs or upstream headers). + """ + if not caller_key: + return "anonymous" + return hashlib.sha256(caller_key.encode()).hexdigest()[:24] def build_app(*clients: Client) -> FastAPI: """A FastAPI app with one POST route per `@on(path)` handler. - Same request/response contract as the sandbox tunnel: JSON-object - bodies in, the handler's `ClientResponse` out, handler errors as - JSON error bodies (`AbridgeError` keeps its status; anything else - is a 502). All requests share the clients' sessions — for - per-caller sessions use `build_session_app`. + Requests and responses keep the tunnel's shapes: JSON-object bodies + in, the handler's `ClientResponse` out, in-band JSON error bodies. + One deliberate improvement over the tunnel wire: an `AbridgeError`'s + status code reaches the agent here, where the tunnel's SIO leg + collapses handler errors to 502. Other exceptions are a 502. All + requests share the clients' sessions — for per-caller sessions use + `build_session_app`. """ handlers: dict[str, Handler] = {} for client in clients: @@ -78,69 +105,165 @@ def build_app(*clients: Client) -> FastAPI: return app +@dataclass +class _Session: + client: Client + handlers: dict[str, Handler] + refs: int = 0 + evicted: bool = False + + +@dataclass +class _SessionTable: + """Caller-keyed client sessions with in-flight-safe LRU eviction. + + Eviction past `max_sessions` removes the least-recent entry from + the table immediately (the table stays bounded), but a client is + closed only once its in-flight requests drain — a live upstream + call is never killed by another caller's arrival. + """ + + factory: Callable[[str], Client] + max_sessions: int + verify_key: Callable[[str], bool] | None + sessions: OrderedDict[str, _Session] = field(default_factory=OrderedDict) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + @asynccontextmanager + async def handler_for(self, request: FastAPIRequest, path: str) -> AsyncIterator[Handler]: + session = await self._acquire(request) + try: + yield session.handlers[path] + finally: + await self._release(session) + + async def _acquire(self, request: FastAPIRequest) -> _Session: + caller_key = _caller_key(request) + if self.verify_key is not None and not self.verify_key(caller_key): + raise _UnknownKey + session_id = session_id_for(caller_key) + + async with self.lock: + session = self.sessions.get(session_id) + if session is not None: + return self._checkout(session_id, session) + + # Build outside the lock: factory work is synchronous (SSL + # context, SDK pool setup) and must not stall other requests. + client = await asyncio.to_thread(self.factory, session_id) + fresh = _Session(client=client, handlers=_collect_handlers(client)) + + loser: Client | None = None + async with self.lock: + session = self.sessions.get(session_id) + if session is not None: + loser = fresh.client + else: + self.sessions[session_id] = fresh + session = fresh + session = self._checkout(session_id, session) + evicted = self._evict_over_cap() + if loser is not None: + await _close_client(loser) + for old in evicted: + await _close_client(old) + return session + + def _checkout(self, session_id: str, session: _Session) -> _Session: + session.refs += 1 + self.sessions.move_to_end(session_id) + return session + + def _evict_over_cap(self) -> list[Client]: + """Pop least-recent entries past the cap (lock held); return the + clients already idle and safe to close now — busy ones close + when their last in-flight request releases.""" + idle: list[Client] = [] + while len(self.sessions) > self.max_sessions: + _, old = self.sessions.popitem(last=False) + old.evicted = True + if old.refs == 0: + idle.append(old.client) + return idle + + async def _release(self, session: _Session) -> None: + async with self.lock: + session.refs -= 1 + close_now = session.evicted and session.refs == 0 + if close_now: + await _close_client(session.client) + + async def close_all(self) -> None: + async with self.lock: + drained = [session.client for session in self.sessions.values()] + self.sessions.clear() + for client in drained: + await _close_client(client) + + def build_session_app( client_factory: Callable[[str], Client], *, max_sessions: int = 256, + verify_key: Callable[[str], bool] | None = None, ) -> FastAPI: """Like `build_app`, but with one client per caller identity. Each request's API key (`x-api-key`, else `Authorization: Bearer`) - hashes to a session id; `client_factory(session_id)` builds that - caller's client on first use. A capped LRU keeps live clients and - closes (via `aclose()`, when implemented) the least recent one past - `max_sessions` — size it above your rollout concurrency so an - in-flight caller is never evicted. Keyless requests share the - `"anonymous"` session. + maps to `session_id_for(key)`; `client_factory(session_id)` builds + that caller's client on first use. The session table is LRU-bounded + at `max_sessions` with in-flight-safe eviction (an evicted client + closes only after its live requests finish), and every remaining + client closes on app shutdown. + + `verify_key`, when given, gates every request: it receives the raw + caller key (`""` when the request carries none) and a falsy return + is a 401. Without it any reachable caller is served — bind the + server accordingly. """ if max_sessions < 1: raise ValueError("max_sessions must be >= 1") - anonymous = client_factory("anonymous") - paths = tuple(_collect_handlers(anonymous)) + probe = client_factory("probe") + paths = tuple(_collect_handlers(probe)) if not paths: raise ValueError("client_factory built a client with no @on-decorated handlers") - # session id → that caller's handler table (client kept alive by its - # bound methods; kept alongside for aclose on eviction). - sessions: OrderedDict[str, tuple[Client, dict[str, Handler]]] = OrderedDict() - sessions["anonymous"] = (anonymous, _collect_handlers(anonymous)) - lock = asyncio.Lock() - - def _session_resolver(path: str) -> Resolver: - async def resolve(request: FastAPIRequest) -> Handler: - session_id = _session_id(_caller_key(request)) - async with lock: - entry = sessions.get(session_id) - if entry is None: - client = client_factory(session_id) - entry = (client, _collect_handlers(client)) - sessions[session_id] = entry - sessions.move_to_end(session_id) - evicted = [] - while len(sessions) > max_sessions: - _, (old, _handlers) = sessions.popitem(last=False) - evicted.append(old) - for old in evicted: - await _close_client(old) - return entry[1][path] - - return resolve + table = _SessionTable(factory=client_factory, max_sessions=max_sessions, verify_key=verify_key) - app = FastAPI() + @asynccontextmanager + async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + # The probe existed only to enumerate routes at build time. + await _close_client(probe) + yield + await table.close_all() + + app = FastAPI(lifespan=lifespan) app.get("/_health")(_health) for path in paths: - app.post(path)(_make_endpoint(path, _session_resolver(path))) + app.post(path)(_make_endpoint(path, _session_resolver(table, path))) return app +def _session_resolver(table: _SessionTable, path: str) -> Resolver: + def resolve(request: FastAPIRequest) -> AbstractAsyncContextManager[Handler]: + return table.handler_for(request, path) + + return resolve + + async def _health() -> dict[str, str]: return {"status": "ok"} +class _UnknownKey(Exception): + """Raised when `verify_key` rejects the caller's key.""" + + def _fixed_resolver(method: Handler) -> Resolver: - async def resolve(_request: FastAPIRequest) -> Handler: - return method + @asynccontextmanager + async def resolve(_request: FastAPIRequest) -> AsyncIterator[Handler]: + yield method return resolve @@ -155,28 +278,22 @@ def _caller_key(request: FastAPIRequest) -> str: return "" -def _session_id(caller_key: str) -> str: - if not caller_key: - return "anonymous" - # The key is identity, not a secret to echo around: hash it so logs - # and upstream `x-session-id` headers never carry the raw value. - return hashlib.sha256(caller_key.encode()).hexdigest()[:24] - - async def _close_client(client: Client) -> None: if isinstance(client, _AsyncCloseable): try: await client.aclose() - except Exception: # noqa: BLE001 - eviction must not fail the live request - logger.exception("abridge serve: failed to close evicted client") + except Exception: # noqa: BLE001 - cleanup must not fail the live request + logger.exception("abridge serve: failed to close client") def _make_endpoint(path: str, resolve: Resolver): async def endpoint(request: FastAPIRequest) -> Response: body = await _read_json(request) try: - handler = await resolve(request) - response = await handler(Request(path=path, body=body)) + async with resolve(request) as handler: + response = await handler(Request(path=path, body=body)) + except _UnknownKey: + return JSONResponse({"error": {"message": "unknown API key"}}, status_code=401) except AbridgeError as exc: logger.warning("abridge serve %s: %s (status=%d)", path, exc.message, exc.status_code) return JSONResponse({"error": {"message": exc.message}}, status_code=exc.status_code) @@ -222,6 +339,15 @@ def main(argv: list[str] | None = None) -> None: default=os.environ.get("UPSTREAM_MODEL"), help="pin every upstream call to this model id (env: UPSTREAM_MODEL)", ) + parser.add_argument( + "--require-key-prefix", + default=os.environ.get("ABRIDGE_KEY_PREFIX"), + help=( + "serve only requests whose API key starts with this secret prefix; " + "mint rollout keys as . Unset = no auth: any reachable " + "caller is served (env: ABRIDGE_KEY_PREFIX)" + ), + ) parser.add_argument("--host", default="127.0.0.1", help="bind address; expose beyond loopback deliberately") parser.add_argument("--port", type=int, default=8399) parser.add_argument("--upstream-timeout", type=float, default=180.0) @@ -242,7 +368,16 @@ def factory(session_id: str) -> AnthropicFromOpenAIClient: session_id=session_id, ) - app = build_session_app(factory, max_sessions=args.max_sessions) + verify_key: Callable[[str], bool] | None = None + if args.require_key_prefix: + prefix = str(args.require_key_prefix) + + def _has_minted_prefix(key: str) -> bool: + return key.startswith(prefix) + + verify_key = _has_minted_prefix + + app = build_session_app(factory, max_sessions=args.max_sessions, verify_key=verify_key) uvicorn.run(app, host=args.host, port=args.port, log_level="info") diff --git a/plugins/abridge/tests/test_serve.py b/plugins/abridge/tests/test_serve.py index 6b0222c..a5e4249 100644 --- a/plugins/abridge/tests/test_serve.py +++ b/plugins/abridge/tests/test_serve.py @@ -1,16 +1,19 @@ """Tests for `agentix.bridge.serve` — the tunnel-less direct mode. -The HTTP contract must match the sandbox tunnel (JSON-object bodies, -`ClientResponse` out, in-band errors), and the session-keyed app must -map caller keys to stable, distinct, LRU-managed client sessions. +The HTTP contract must match the sandbox tunnel's shapes (JSON-object +bodies, `ClientResponse` out, in-band errors), and the session-keyed +app must map caller keys to stable, distinct sessions whose clients are +never closed under an in-flight request and always closed by shutdown. """ from __future__ import annotations +import asyncio from typing import Any +import httpx from agentix.bridge import AbridgeError, ClientResponse, Request, on -from agentix.bridge.serve import _session_id, build_app, build_session_app +from agentix.bridge.serve import build_app, build_session_app, session_id_for from fastapi.testclient import TestClient @@ -67,13 +70,13 @@ def test_session_app_maps_keys_to_stable_distinct_sessions() -> None: bearer = tc.post("/v1/echo", json={}, headers={"authorization": "Bearer sk-b"}).json()["session"] anonymous = tc.post("/v1/echo", json={}).json()["session"] - assert a1 == a2 == _session_id("sk-a") - assert bearer == _session_id("sk-b") + assert a1 == a2 == session_id_for("sk-a") + assert bearer == session_id_for("sk-b") assert bearer != a1 assert anonymous == "anonymous" -def test_session_app_evicts_and_closes_least_recent() -> None: +def test_session_app_evicts_and_closes_idle_least_recent() -> None: built: list[EchoClient] = [] def factory(session_id: str) -> EchoClient: @@ -85,9 +88,77 @@ def factory(session_id: str) -> EchoClient: tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-a"}) tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-b"}) - still_open = [client.session_id for client in built if not client.closed] - assert still_open == [_session_id("sk-b")] - assert [client.session_id for client in built if client.closed] == ["anonymous", _session_id("sk-a")] + closed = [client.session_id for client in built if client.closed] + assert closed == [session_id_for("sk-a")] + + +def test_shutdown_closes_probe_and_all_session_clients() -> None: + built: list[EchoClient] = [] + + def factory(session_id: str) -> EchoClient: + client = EchoClient(session_id) + built.append(client) + return client + + with TestClient(build_session_app(factory)) as tc: + tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-a"}) + tc.post("/v1/echo", json={}, headers={"x-api-key": "sk-b"}) + + assert [client.session_id for client in built if not client.closed] == [] + + +def test_verify_key_gates_requests() -> None: + tc = TestClient(build_session_app(EchoClient, verify_key=lambda key: key.startswith("ok-"))) + assert tc.post("/v1/echo", json={}, headers={"x-api-key": "bad"}).status_code == 401 + assert tc.post("/v1/echo", json={}).status_code == 401 + r = tc.post("/v1/echo", json={}, headers={"x-api-key": "ok-1"}) + assert r.status_code == 200 + assert r.json()["session"] == session_id_for("ok-1") + + +class BlockingClient(EchoClient): + def __init__(self, session_id: str) -> None: + super().__init__(session_id) + self.entered = asyncio.Event() + self.release = asyncio.Event() + + @on("/v1/block") + async def block(self, request: Request) -> ClientResponse: + self.entered.set() + await self.release.wait() + return ClientResponse.json({"session": self.session_id}) + + +async def test_eviction_defers_close_until_inflight_request_drains() -> None: + built: list[BlockingClient] = [] + + def factory(session_id: str) -> BlockingClient: + client = BlockingClient(session_id) + if session_id != session_id_for("sk-slow"): + client.release.set() + built.append(client) + return client + + app = build_session_app(factory, max_sessions=1) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://serve") as hc: + slow = asyncio.create_task(hc.post("/v1/block", json={}, headers={"x-api-key": "sk-slow"})) + while not any(c.session_id == session_id_for("sk-slow") for c in built): + await asyncio.sleep(0.01) + slow_client = next(c for c in built if c.session_id == session_id_for("sk-slow")) + await asyncio.wait_for(slow_client.entered.wait(), timeout=5) + + # A new caller key evicts the slow session from the table... + r = await hc.post("/v1/block", json={}, headers={"x-api-key": "sk-new"}) + assert r.status_code == 200 + # ...but must not close its client mid-upstream-call. + assert not slow_client.closed + + slow_client.release.set() + response = await asyncio.wait_for(slow, timeout=5) + assert response.status_code == 200 + await asyncio.sleep(0.01) + assert slow_client.closed def _mock_completion() -> Any: @@ -130,7 +201,7 @@ async def create(**kwargs: Any) -> Any: assert tc.post("/v1/messages", json=body, headers={"x-api-key": "rollout-1"}).status_code == 200 assert tc.post("/v1/messages", json=body, headers={"x-api-key": "rollout-2"}).status_code == 200 - assert seen[0]["x-session-id"] == _session_id("rollout-1") - assert seen[1]["x-session-id"] == _session_id("rollout-2") + assert seen[0]["x-session-id"] == session_id_for("rollout-1") + assert seen[1]["x-session-id"] == session_id_for("rollout-2") assert seen[0]["x-session-id"] != seen[1]["x-session-id"] assert all("rollout-1" not in v and "rollout-2" not in v for headers in seen for v in headers.values())