From 4aee3030117f05a786ffa7300cc4ce4fe208313c Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Thu, 2 Jul 2026 17:40:58 +0800 Subject: [PATCH] =?UTF-8?q?abridge:=20SessionForward=20+=20Convert?= =?UTF-8?q?=E2=88=98Session=20composition=20(PR=20#122=20stage=20A)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the abridge SessionForward/AnthropicToOpenAI work from origin/abridge/gateway-tito onto master as the first slice of the #122 split: - SessionForward: forward to session-scoped sidecars (create-on-first- request, path rewrite to /sessions/{id}{path}, delete_session reap). - AnthropicToOpenAI: Anthropic Messages agent over any OpenAI-compatible downstream Handler — protocol translation only, transport-blind, with verbatim assistant replay so session recorders byte-match history. - Forward.handler(): the composition seam — a bound Handler a wrapper can sit on without knowing the transport. - proxy: typed DynamicRoutes protocol for construction-time routes; tunnel 400s on non-object JSON bodies instead of coercing to {}. - sio: RemoteSioError carries the error envelope's status_code, so the tunnel replies with the handler's real status (429/400/503) instead of collapsing every remote failure to 502. (The source branch had this producer half; the port initially took only the consumer.) Review fixes on top of the port (from the pre-PR adversarial review): - _ensure_session accepts any 2xx create response, not just 200. - session_id setter now attaches: assigning an id marks the session ready so the first request doesn't create a fresh session and silently overwrite the assigned id. - Forward.handler() returns a closeable handler and AnthropicToOpenAI delegates aclose() to it, so Proxy.stop reaps the forwarder's httpx pool through the converter composition. - Proxy.session() detects a body failure by catching it, not by probing sys.exc_info() — teardown errors are no longer swallowed for callers that open the session inside an unrelated except block. - delete_session() no longer resurrects a closed pool for one DELETE; it uses a one-shot client, so the documented stop → harvest → reap flow leaks nothing. - Assistant replay memory keys on (id, name, canonical args) per tool call, not ids alone — servers that reuse tool-call ids across turns (TGI's literal "0") no longer corrupt the replayed history. Co-authored-by: FatPigeorz Co-Authored-By: Claude Fable 5 --- agentix/sio.py | 9 +- plugins/abridge/README.md | 13 +- plugins/abridge/agentix/bridge/__init__.py | 13 +- .../agentix/bridge/clients/__init__.py | 34 +- .../bridge/clients/anthropic_to_openai.py | 166 ++++++ .../abridge/agentix/bridge/clients/openai.py | 21 +- plugins/abridge/agentix/bridge/forward.py | 206 ++++++- plugins/abridge/agentix/bridge/proxy.py | 120 ++-- plugins/abridge/agentix/bridge/sidecar.py | 9 +- plugins/abridge/tests/test_sidecar_forward.py | 535 +++++++++++++++++- tests/test_sio_reply_error.py | 48 ++ 11 files changed, 1107 insertions(+), 67 deletions(-) create mode 100644 plugins/abridge/agentix/bridge/clients/anthropic_to_openai.py create mode 100644 tests/test_sio_reply_error.py diff --git a/agentix/sio.py b/agentix/sio.py index 9c5638a..d8b9c5b 100644 --- a/agentix/sio.py +++ b/agentix/sio.py @@ -56,10 +56,15 @@ async def fetch_remote(self, payload): class RemoteSioError(RuntimeError): """Raised by `Namespace.request()` when the reply carries an `:error`.""" - def __init__(self, type_: str, message: str) -> None: + def __init__(self, type_: str, message: str, status_code: int | None = None) -> None: super().__init__(f"{type_}: {message}") self.type = type_ self.message = message + # Upstream HTTP status the host handler chose (e.g. abridge's + # AbridgeError 429/400); None when the error envelope carried no + # status. Lets a protocol layer reply with the real status instead + # of collapsing every remote failure to one blanket code. + self.status_code = status_code # ── module-level bridge state ────────────────────────────────────── @@ -241,10 +246,12 @@ async def _on_reply_error(self, payload: Any) -> None: fut = self._pending_requests.get(req_id) if isinstance(req_id, str) else None if fut is not None and not fut.done(): err = payload.get("error") or {"type": "Unknown", "message": ""} + status_code = err.get("status_code") fut.set_exception( RemoteSioError( err.get("type", "Unknown"), err.get("message", ""), + status_code if isinstance(status_code, int) else None, ) ) diff --git a/plugins/abridge/README.md b/plugins/abridge/README.md index fe09cb1..34bc8d8 100644 --- a/plugins/abridge/README.md +++ b/plugins/abridge/README.md @@ -49,7 +49,7 @@ from agentix.bridge.clients import AnthropicFromOpenAIClient client = AnthropicFromOpenAIClient( base_url="https://api.openai.com/v1", # OpenAI / OpenRouter / vLLM / your gateway api_key="sk-...", - upstream_model="gpt-4o", # the agent keeps sending claude-* model ids + model="gpt-4o", # the agent keeps sending claude-* model ids ) proxy = Proxy(client) @@ -75,12 +75,14 @@ client = OpenAIClient(base_url=..., api_key=..., model="gpt-4o") proxy = Proxy(client) async with proxy.session(sandbox) as handle: - await sandbox.remote(agent, base_url=f"{handle.url}/v1", api_key=PLACEHOLDER_API_KEY) + await sandbox.remote(agent, env=client.environ(handle)) ``` -`OpenAIClient` doesn't ship an `environ(handle)` helper — most OpenAI -SDK callers construct the client with explicit `base_url=`/`api_key=` -arguments rather than reading env vars. +`client.environ(handle)` returns +`{"OPENAI_BASE_URL": handle.url + "/v1", "OPENAI_API_KEY": ""}` — +the `/v1` suffix the OpenAI SDK expects is baked in so a caller can't +drop it. Agents that construct their SDK client explicitly can still +pass `base_url=f"{handle.url}/v1", api_key=PLACEHOLDER_API_KEY` instead. ### Anthropic agent → native Anthropic upstream @@ -252,6 +254,7 @@ agentix/bridge/ ├── openai.py # OpenAIClient (openai SDK) + PLACEHOLDER_API_KEY ├── anthropic.py # AnthropicClient (anthropic SDK) + environ() + PLACEHOLDER_API_KEY ├── anthropic_from_openai.py # AnthropicFromOpenAIClient (openai SDK + translation) + environ() + ├── anthropic_to_openai.py # AnthropicToOpenAI (SDK-free translation over any Handler) + environ() ├── _genai_span.py # populate_openai_span / populate_anthropic_span └── _anthropic_transforms.py # pure Anthropic↔OpenAI converters ``` diff --git a/plugins/abridge/agentix/bridge/__init__.py b/plugins/abridge/agentix/bridge/__init__.py index 053617b..229fc89 100644 --- a/plugins/abridge/agentix/bridge/__init__.py +++ b/plugins/abridge/agentix/bridge/__init__.py @@ -25,20 +25,21 @@ via Python multiple inheritance (mixins) or by passing several to `Proxy(*handlers)`. Two handlers must not register the same path. -`agentix.bridge.clients` ships the three standard handlers -(`OpenAIClient`, `AnthropicClient`, `AnthropicFromOpenAIClient`), each -built on the corresponding provider SDK. Pull from there as building -blocks; abridge core stays shape-blind. +`agentix.bridge.clients` ships the standard handlers (`OpenAIClient`, +`AnthropicClient`, `AnthropicFromOpenAIClient`, and the transport-blind +`AnthropicToOpenAI`). Pull from there as building blocks; abridge core +stays shape-blind. """ from __future__ import annotations -from .forward import Forward +from .forward import Forward, SessionForward from .proxy import ( NAMESPACE, AbridgeError, Client, ClientResponse, + DynamicRoutes, Handler, Proxy, Request, @@ -54,11 +55,13 @@ "Client", "ClientResponse", "Command", + "DynamicRoutes", "Forward", "Handler", "NAMESPACE", "Proxy", "Request", + "SessionForward", "Sidecar", "SidecarError", "TunnelHandle", diff --git a/plugins/abridge/agentix/bridge/clients/__init__.py b/plugins/abridge/agentix/bridge/clients/__init__.py index 3869ba7..aa9c723 100644 --- a/plugins/abridge/agentix/bridge/clients/__init__.py +++ b/plugins/abridge/agentix/bridge/clients/__init__.py @@ -1,24 +1,32 @@ """Bundled handler clients for abridge. -Three out-of-the-box implementations, all built around the official -provider SDKs (`openai`, `anthropic`). Each is a plain class with -`@on(path)`-decorated methods — pass an instance to `Proxy(...)` or -mixin-compose multiple in one user-defined class. +Four out-of-the-box implementations. Each is a plain class with handler +methods — pass an instance to `Proxy(...)` or mixin-compose multiple in +one user-defined class. The first three own their upstream transport via +the official provider SDKs (`openai`, `anthropic`); the fourth is +SDK-free and transport-blind. * `OpenAIClient` — agent speaks OpenAI Chat Completions, upstream is OpenAI-compatible. One `@on("/v1/chat/completions")`. * `AnthropicClient` — agent speaks Anthropic Messages, upstream is native Anthropic. `@on("/v1/messages")` + `@on("/v1/messages/count_tokens")`. * `AnthropicFromOpenAIClient` — agent speaks Anthropic, upstream is - OpenAI-compatible (translation lives here). Same path set as - `AnthropicClient`. + OpenAI-compatible (translation lives here, transport owned by the + `openai` SDK). Same path set as `AnthropicClient`. + * `AnthropicToOpenAI` — the same translation direction, but over ANY + abridge `Handler` downstream (`Forward(...).handler()`, + `SessionForward(...).handler()`, …): the downstream owns HTTP, + sessions, and recording. Pick this one to compose with a + session-scoped recorder; pick `AnthropicFromOpenAIClient` when a + plain SDK-managed upstream is all you need. -The two Anthropic-side classes also expose `environ(handle)` (instance -method) — the env-var bundle (`ANTHROPIC_BASE_URL` + placeholder -`ANTHROPIC_API_KEY`) an in-sandbox Anthropic SDK needs to route through -the tunnel. The OpenAI client doesn't ship an `environ`; agents that -use the OpenAI SDK typically construct the client with -`base_url=handle.url + "/v1"` directly. +All four classes expose `environ(handle)` (instance method) — the env-var +bundle an in-sandbox SDK needs to route through the tunnel, so the wiring step +is `env=client.environ(handle)` uniformly. `OpenAIClient` returns +`{OPENAI_BASE_URL: handle.url + "/v1", OPENAI_API_KEY: placeholder}` — the `/v1` +suffix the OpenAI SDK expects is baked in so a caller can't drop it; the two +Anthropic-side classes return `{ANTHROPIC_BASE_URL: handle.url, ANTHROPIC_API_KEY: +placeholder}` (no `/v1` — the Anthropic SDK appends it itself). The two `populate_*_span` helpers are exposed at this level so user- written clients can stamp the same OTel GenAI attrs the bundled @@ -31,6 +39,7 @@ from .anthropic import PLACEHOLDER_API_KEY as ANTHROPIC_PLACEHOLDER_API_KEY from .anthropic import AnthropicClient from .anthropic_from_openai import AnthropicFromOpenAIClient +from .anthropic_to_openai import AnthropicToOpenAI from .openai import PLACEHOLDER_API_KEY as OPENAI_PLACEHOLDER_API_KEY from .openai import OpenAIClient @@ -38,6 +47,7 @@ "ANTHROPIC_PLACEHOLDER_API_KEY", "AnthropicClient", "AnthropicFromOpenAIClient", + "AnthropicToOpenAI", "OPENAI_PLACEHOLDER_API_KEY", "OpenAIClient", "populate_anthropic_span", diff --git a/plugins/abridge/agentix/bridge/clients/anthropic_to_openai.py b/plugins/abridge/agentix/bridge/clients/anthropic_to_openai.py new file mode 100644 index 0000000..5d5a42f --- /dev/null +++ b/plugins/abridge/agentix/bridge/clients/anthropic_to_openai.py @@ -0,0 +1,166 @@ +"""`AnthropicToOpenAI` — the Convert capability: an Anthropic-Messages agent over +ANY OpenAI-compatible downstream. + +This is one orthogonal capability — protocol translation — and nothing else. It +translates the agent's Anthropic `/v1/messages` to an OpenAI chat-completions body, +hands that body to a `downstream` (any abridge `Handler`), then translates the +OpenAI completion back to Anthropic. It is deliberately transport-blind: the +downstream owns the HTTP, the session, and any recording. + +Compose it with a transport/routing capability: + + # plain OpenAI gateway (no session) + Proxy(AnthropicToOpenAI(Forward(base_url, paths=["/v1/chat/completions"]).handler())) + + # session-scoped recorder (the TITO gateway) — session stays transparent + tito = SessionForward(gateway_url, paths=["/v1/chat/completions"]) + proxy = Proxy(AnthropicToOpenAI(tito.handler(), model="qwen3-4b")) + ... # the agent only ever speaks Anthropic + harvest(tito.session_id) # the session lives in the SessionForward, not here + +`AnthropicToOpenAI` knows nothing about sessions; `SessionForward` knows nothing +about Anthropic. They compose because the seam between them is just abridge's +`Handler` (an OpenAI chat body in, a `ClientResponse` out). No OpenAI SDK +dependency — the upstream hop is whatever `Handler` you pass. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +from agentix.utils import trace + +from ..proxy import ClientResponse, Handler, Request, TunnelHandle, _AsyncCloseable, on +from ._anthropic_transforms import ( + anthropic_messages_to_openai, + anthropic_sse, + count_anthropic_tokens, + openai_to_anthropic_messages, +) +from ._genai_span import populate_anthropic_span +from .anthropic import PLACEHOLDER_API_KEY + +logger = logging.getLogger(__name__) + + +class AnthropicToOpenAI: + """Anthropic Messages agent → any OpenAI-compatible `downstream` Handler. + + `downstream` is the transport/routing capability this converter sits on top of + — a `Forward(...).handler()` for a direct OpenAI gateway, or a + `SessionForward(...).handler()` for a session-scoped recorder like the TITO + gateway. `model`, when set, overrides the agent's model id in the OpenAI body. + `count_tokens` is answered locally (character estimate, no downstream call). + + `aclose()` delegates to a closeable downstream (a `Forward.handler()` carries + one), so `Proxy.stop` reaps the transport's HTTP pool through the composition. + Everything else about the downstream stays yours: hold the forwarder to read + `.session_id` / call `delete_session()`. + """ + + def __init__(self, downstream: Handler, *, model: str | None = None) -> None: + self._downstream = downstream + self._model = model + # Memory of the EXACT assistant message the downstream returned, one + # entry per tool-calling turn. The Anthropic round-trip is lossy (it + # drops reasoning_content and the per-tool-call `index`), so a + # reconstructed assistant won't match what a session-recording backend + # (TITO) stored byte-for-byte. On the next turn we replay the remembered + # original verbatim, keyed by the (id, name, canonical-args) triple of + # each tool call — all of which survive the round-trip. Ids alone are + # NOT a safe key: some servers reuse ids across turns (TGI emits "0"), + # and an id-only key would replay the latest turn into every matching + # history slot. Entries are never evicted — scope one converter + # instance per rollout/conversation. + self._assistant_by_calls: dict[ + tuple[tuple[str, str, str], ...], dict[str, Any] + ] = {} + + @on("/v1/messages") + async def messages(self, request: Request) -> ClientResponse: + openai_body = anthropic_messages_to_openai(request.body, upstream_model=self._model) + # The downstream produces a non-streaming OpenAI completion (a TITO + # recorder needs output_token_logprobs); we re-render SSE locally below if + # the agent asked for streaming. + openai_body["stream"] = False + self._replay_remembered_assistants(openai_body) + with trace.span(f"anthropic messages {request.body.get('model') or ''}"): + resp = await self._downstream(Request(path="/v1/messages", body=openai_body)) + if resp.status_code != 200: + # Pass the downstream's error (4xx/5xx) straight through with its + # status; the agent sees a non-200, not a malformed body. + return resp + openai_resp = json.loads(resp.body) + self._remember_assistant(openai_resp) + anthropic_resp = openai_to_anthropic_messages( + openai_resp, response_model=str(request.body.get("model") or "") + ) + populate_anthropic_span(request=request.body, response=anthropic_resp) + if request.body.get("stream"): + return ClientResponse.sse(anthropic_sse(anthropic_resp)) + return ClientResponse.json(anthropic_resp) + + @staticmethod + def _call_key(message: dict[str, Any]) -> tuple[tuple[str, str, str], ...]: + """One (id, name, canonical-args) triple per tool call. Arguments are + re-serialized with sorted keys so the downstream's original JSON string + and the round-trip's `json.dumps` rendering compare equal regardless of + spacing or key order.""" + key: list[tuple[str, str, str]] = [] + for tc in message.get("tool_calls") or []: + if not isinstance(tc, dict) or not tc.get("id"): + continue + function = tc.get("function") or {} + raw_args = function.get("arguments") or "{}" + try: + args = json.dumps(json.loads(raw_args), sort_keys=True) + except ValueError: + args = raw_args + key.append((str(tc["id"]), str(function.get("name") or ""), args)) + return tuple(key) + + def _remember_assistant(self, openai_resp: dict[str, Any]) -> None: + choice = (openai_resp.get("choices") or [{}])[0] + message = choice.get("message") or {} + calls = self._call_key(message) + if calls: + self._assistant_by_calls[calls] = message + + def _replay_remembered_assistants(self, openai_body: dict[str, Any]) -> None: + messages = openai_body.get("messages") + if not isinstance(messages, list): + return + for i, msg in enumerate(messages): + if not isinstance(msg, dict) or msg.get("role") != "assistant": + continue + remembered = self._assistant_by_calls.get(self._call_key(msg)) + if remembered is not None: + messages[i] = remembered + + @on("/v1/messages/count_tokens") + async def count_tokens(self, request: Request) -> ClientResponse: + return ClientResponse.json( + {"input_tokens": count_anthropic_tokens(request.body).input_tokens} + ) + + async def aclose(self) -> None: + """Close the downstream if it is closeable; a bare function is a no-op. + + Lets `Proxy.stop` reach the transport's HTTP pool through this + converter — otherwise `Proxy(AnthropicToOpenAI(fwd.handler()))` leaks + the forwarder's pool past the proxy lifecycle.""" + if isinstance(self._downstream, _AsyncCloseable): + await self._downstream.aclose() + + def environ(self, handle: TunnelHandle) -> dict[str, str]: + """Anthropic env-var bundle — from the agent's POV the wire is Anthropic, + regardless of the OpenAI downstream.""" + return { + "ANTHROPIC_BASE_URL": handle.url, + "ANTHROPIC_API_KEY": PLACEHOLDER_API_KEY, + } + + +__all__ = ["AnthropicToOpenAI"] diff --git a/plugins/abridge/agentix/bridge/clients/openai.py b/plugins/abridge/agentix/bridge/clients/openai.py index 537c3c2..b906882 100644 --- a/plugins/abridge/agentix/bridge/clients/openai.py +++ b/plugins/abridge/agentix/bridge/clients/openai.py @@ -7,9 +7,10 @@ at our tunnel. The SDK accepts the OpenAI Chat Completions request shape via typed -kwargs. Agents that send non-standard fields the SDK doesn't accept -will see an `UpstreamError`; for arbitrary-shape forwarding write your -own `@on("/v1/chat/completions")` handler with raw httpx. +kwargs. Upstream failures surface as `AbridgeError` carrying the +upstream HTTP status (so the agent sees the real 429/400/... instead of +a blanket 502); for arbitrary-shape forwarding write your own +`@on("/v1/chat/completions")` handler with raw httpx. """ from __future__ import annotations @@ -20,7 +21,7 @@ from agentix.utils import trace -from ..proxy import AbridgeError, ClientResponse, Request, on +from ..proxy import AbridgeError, ClientResponse, Request, TunnelHandle, on from ._genai_span import populate_openai_span if TYPE_CHECKING: @@ -113,5 +114,17 @@ async def chat(self, request: Request) -> ClientResponse: populate_openai_span(request=request.body, response=response_dict) return ClientResponse.json(response_dict) + def environ(self, handle: TunnelHandle) -> dict[str, str]: + """Env-var bundle an in-sandbox OpenAI SDK needs to route through `handle`, + mirroring `AnthropicClient.environ`. `OPENAI_BASE_URL` carries the `/v1` + suffix the OpenAI SDK expects (it appends only `/chat/completions`), so the + load-bearing suffix is baked in here instead of left to the caller. + `OPENAI_API_KEY` is a non-secret placeholder whose shape passes the SDK's + local format check — the real upstream key lives on the host (this client).""" + return { + "OPENAI_BASE_URL": handle.url + "/v1", + "OPENAI_API_KEY": PLACEHOLDER_API_KEY, + } + __all__ = ["PLACEHOLDER_API_KEY", "OpenAIClient"] diff --git a/plugins/abridge/agentix/bridge/forward.py b/plugins/abridge/agentix/bridge/forward.py index eada27e..ff3299d 100644 --- a/plugins/abridge/agentix/bridge/forward.py +++ b/plugins/abridge/agentix/bridge/forward.py @@ -27,6 +27,8 @@ from __future__ import annotations +import asyncio +import contextlib import logging import uuid from collections.abc import Mapping @@ -38,6 +40,28 @@ logger = logging.getLogger(__name__) +class _OwnedHandler: + """A `Handler` that keeps a lifecycle link to the forwarder that owns it. + + `Forward.handler()` hands a bare callable to a wrapper (e.g. + `AnthropicToOpenAI`), which hides the forwarder from the `Proxy` that + closes clients on stop. Exposing `aclose()` here lets the wrapper delegate + close through the seam, so the owner's HTTP pool doesn't outlive the proxy. + """ + + __slots__ = ("_handler", "_owner") + + def __init__(self, handler: Handler, owner: Forward) -> None: + self._handler = handler + self._owner = owner + + async def __call__(self, request: Request) -> ClientResponse: + return await self._handler(request) + + async def aclose(self) -> None: + await self._owner.aclose() + + class Forward: """Forward JSON POSTs on `paths` to `target_url` over httpx. @@ -51,7 +75,10 @@ class Forward: Responses, including 4xx and 5xx responses, remain normal `ClientResponse` values so their status and body survive the tunnel. - `AbridgeError(502)` is reserved for failures to obtain an HTTP response. + `AbridgeError(503)` signals a failure to obtain any HTTP response from the + sidecar (connection refused, DNS, timeout) — a distinct code from a real + upstream 502 the sidecar relays, so the agent can tell "sidecar down" from + "sidecar returned bad gateway". """ def __init__( @@ -81,6 +108,25 @@ async def handler(request: Request) -> ClientResponse: return handler + def handler(self, path: str | None = None) -> Handler: + """The bound forwarder for `path` (or the sole path if there's exactly + one) as a plain `Handler`. This is the composition seam: it lets a + converter/wrapper use this Forward as a transparent downstream — + `AnthropicToOpenAI(SessionForward(gw).handler())` — without knowing it's a + Forward, a SessionForward, or anything else. The returned handler also + carries `aclose()` (delegating to this forwarder), so a wrapper can + propagate `Proxy.stop`'s client cleanup through the seam.""" + routes = self.abridge_routes() + if path is None: + if len(routes) != 1: + raise ValueError( + f"handler() needs an explicit path; forwarder has {sorted(routes)}" + ) + (path,) = routes + if path not in routes: + raise ValueError(f"no route for {path!r}; have {sorted(routes)}") + return _OwnedHandler(routes[path], self) + async def _forward(self, path: str, request: Request) -> ClientResponse: record_id = uuid.uuid4().hex headers = { @@ -89,12 +135,12 @@ async def _forward(self, path: str, request: Request) -> ClientResponse: "x-request-id": record_id, "content-type": "application/json", } - url = self._target + path + url = self._url_for(path) try: resp = await self._get_client().post(url, json=request.body, headers=headers) except httpx.HTTPError as exc: logger.warning("abridge forward %s: %s", url, exc) - raise AbridgeError(f"forward to {url}: {exc}", status_code=502) from exc + raise AbridgeError(f"forward to {url}: {exc}", status_code=503) from exc media_type = resp.headers.get("content-type", "application/json").split(";")[0].strip() return ClientResponse( @@ -103,6 +149,11 @@ async def _forward(self, path: str, request: Request) -> ClientResponse: status_code=resp.status_code, ) + def _url_for(self, path: str) -> str: + """Upstream URL for an inbound `path`. Override to remap — e.g. a + session-scoped sidecar prefixes `/sessions/{id}`.""" + return self._target + path + def _get_client(self) -> httpx.AsyncClient: client = self._client if client is None or client.is_closed: @@ -118,4 +169,151 @@ async def aclose(self) -> None: await client.aclose() -__all__ = ["Forward"] +class SessionForward(Forward): + """Forward to a *session-scoped* sidecar that keys a trajectory by URL path. + + Some sidecars don't accept a bare `/v1/chat/completions` — they require a + session created up front and then addressed by path: `POST {create_path}` + returns `{"session_id": ...}`, and every later call goes to + `{create_path}/{session_id}{path}`. The TITO gateway is the motivating case: + its recording route is `/sessions/{id}/v1/chat/completions`, so a plain + `Forward` (which posts straight to `{target}{path}`) can't reach it, and the + id isn't known until the gateway assigns it. + + `SessionForward` creates the session lazily on the first forwarded request + (or eagerly via `open()`), remembers the assigned id, and rewrites every + inbound `path` to the session-scoped URL. So the in-sandbox agent keeps + calling an unmodified `/v1/chat/completions` and the whole rollout still + lands in one session. + + fwd = SessionForward(gateway_url, paths=["/v1/chat/completions"]) + async with Proxy(fwd).session(sandbox) as handle: + # the agent just POSTs /v1/chat/completions at handle.url; it never + # sees a session id — SessionForward creates + scopes it host-side. + await sandbox.remote(agent, base_url=handle.url) + sid = fwd.session_id # assigned after the run (or `await fwd.open()` up front) + trajectory = (await httpx.AsyncClient().get( + f"{gateway_url}/sessions/{sid}")).json() + await fwd.delete_session() # optional: reap the server-side session + + `.session_id` is only valid once the session exists — reading it before the + first request (or `open()`) raises, so a premature harvest fails loudly rather + than hitting a fabricated id. Assigning `fwd.session_id = existing_id` + *attaches* to an already-created gateway session: requests scope to it and + no create call runs. One instance == one gateway session: all of an + instance's calls accumulate into the same session, so use a fresh + `SessionForward` per rollout (or call `delete_session()` to reap and reset). + The session is intentionally *not* deleted on `aclose()` — the trajectory is + the point and must survive the proxy teardown for harvesting. + """ + + def __init__( + self, + target_url: str, + *, + paths: list[str], + create_path: str = "/sessions", + session_id_field: str = "session_id", + timeout: float = 600.0, + headers: Mapping[str, str] | None = None, + ) -> None: + super().__init__(target_url, paths=paths, timeout=timeout, headers=headers) + # Forward.__init__ stamped a throwaway uuid via the setter below; discard + # it — the gateway assigns the real id on open()/first call. Until then a + # read of `.session_id` raises instead of returning a meaningless value. + self._session_id: str | None = None + self._create_path = "/" + create_path.strip("/") + self._session_field = session_id_field + self._session_ready = False + self._session_lock = asyncio.Lock() + + @property + def session_id(self) -> str: + if self._session_id is None: + raise RuntimeError( + "SessionForward.session_id is unavailable until the gateway session " + "is created — call `await fwd.open()` (or make one forwarded request) " + "before harvesting." + ) + return self._session_id + + @session_id.setter + def session_id(self, value: str) -> None: + # Attach semantics: assigning an id means "this session already exists + # on the gateway" — later requests scope to it and no create call runs. + # Without the ready flag the first request would create a fresh session + # and silently overwrite the assigned id. + self._session_id = value + self._session_ready = True + + async def open(self) -> str: + """Create the sidecar session now (idempotent) and return its id. + + Lazy creation also happens on the first forwarded request, so `open()` + is only needed when the host wants the id before the agent runs. + """ + await self._ensure_session() + return self.session_id + + async def delete_session(self) -> None: + """Reap the server-side session (e.g. after the trajectory is harvested). + + No-op if no session was created. Kept separate from `aclose()` / `Proxy.stop` + so the default flow preserves the trajectory; after this the next request + opens a fresh session. Transport errors are suppressed — best-effort reap. + """ + async with self._session_lock: + sid = self._session_id + if sid is None: + return + url = f"{self._target}{self._create_path}/{sid}" + with contextlib.suppress(httpx.HTTPError): + client = self._client + if client is not None and not client.is_closed: + await client.delete(url, headers=dict(self._headers)) + else: + # Don't resurrect the pool for one reap: the documented + # harvest flow calls this after `Proxy.stop` closed it, and + # a lazily recreated pool here would have no owner left to + # close it. + async with httpx.AsyncClient(timeout=self._timeout) as one_shot: + await one_shot.delete(url, headers=dict(self._headers)) + self._session_ready = False + self._session_id = None + + async def _ensure_session(self) -> None: + if self._session_ready: + return + async with self._session_lock: + if self._session_ready: + return + url = self._target + self._create_path + headers = {**self._headers, "content-type": "application/json"} + try: + resp = await self._get_client().post(url, json={}, headers=headers) + except httpx.HTTPError as exc: + logger.warning("abridge session create %s: %s", url, exc) + raise AbridgeError(f"create session at {url}: {exc}", status_code=503) from exc + if not resp.is_success: + raise AbridgeError( + f"create session at {url}: HTTP {resp.status_code}", status_code=502 + ) + try: + session_id = resp.json()[self._session_field] + except (ValueError, KeyError, TypeError) as exc: + raise AbridgeError( + f"create session at {url}: response missing {self._session_field!r}", + status_code=502, + ) from exc + self.session_id = str(session_id) + logger.info("abridge session created at %s: %s", url, self.session_id) + + def _url_for(self, path: str) -> str: + return f"{self._target}{self._create_path}/{self.session_id}{path}" + + async def _forward(self, path: str, request: Request) -> ClientResponse: + await self._ensure_session() + return await super()._forward(path, request) + + +__all__ = ["Forward", "SessionForward"] diff --git a/plugins/abridge/agentix/bridge/proxy.py b/plugins/abridge/agentix/bridge/proxy.py index f057746..03f8ada 100644 --- a/plugins/abridge/agentix/bridge/proxy.py +++ b/plugins/abridge/agentix/bridge/proxy.py @@ -138,21 +138,36 @@ def sse(cls, body: bytes, *, status_code: int = 200) -> ClientResponse: Handler = Callable[[Request], Awaitable[ClientResponse]] -@runtime_checkable class Client(Protocol): - """Marker protocol for any class with at least one `@on(path)`-decorated - method. - - There's nothing for the protocol to require structurally — `@on` is a - method-level attribute, not a class-level signature, so `isinstance` - against `Client` doesn't validate handler presence (that's - `Proxy.__init__`'s job at construction time). The name exists so - `Proxy(*clients: Client)` reads as "pass handler classes here" rather - than `*clients: object`. A client may additionally implement async + """Marker protocol for a handler object passed to `Proxy(...)`. + + Two kinds qualify, and they compose: a class with at least one + `@on(path)`-decorated method, and/or a class implementing + `DynamicRoutes.abridge_routes()` (paths chosen at construction, e.g. + `Forward`). The protocol requires nothing structurally — `@on` is a + method-level attribute, not a class-level signature — so it is deliberately + NOT `runtime_checkable`: `isinstance(x, Client)` would be true for anything + and is meaningless. Handler presence is validated by `Proxy.__init__` at + construction time. The name just makes `Proxy(*clients: Client)` read as + "pass handler objects here". A client may additionally implement async `aclose()`; `Proxy.stop()` closes such clients once per lifecycle. """ +@runtime_checkable +class DynamicRoutes(Protocol): + """A client that contributes routes chosen at construction time — paths the + class-level `@on` tag can't express, e.g. `Forward(target, paths=[...])`. + + `abridge_routes()` returns `{path: handler}`; `Proxy` merges it alongside any + `@on` handlers under the same duplicate-path rule. This is the typed, blessed + second registration seam (vs. `@on`): a handler returning dynamic routes has a + checkable contract instead of an undocumented duck-typed method. + """ + + def abridge_routes(self) -> dict[str, Handler]: ... + + @runtime_checkable class _AsyncCloseable(Protocol): def aclose(self) -> Awaitable[None]: ... @@ -229,13 +244,12 @@ def _collect_handlers(client: Client) -> dict[str, Handler]: ) handlers[path] = getattr(client, name) - # Dynamic routes: a client may expose `abridge_routes() -> dict[str, - # Handler]` for paths chosen at construction time (e.g. `Forward(target, - # paths=[...])`), which the class-level `@on` tag can't express. They - # compose with `@on` handlers under the same duplicate-path rule. - dynamic = getattr(client, "abridge_routes", None) - if callable(dynamic): - routes = dynamic() + # Dynamic routes: a client implementing `DynamicRoutes` contributes paths + # chosen at construction time (e.g. `Forward(target, paths=[...])`), which + # the class-level `@on` tag can't express. They compose with `@on` handlers + # under the same duplicate-path rule. + if isinstance(client, DynamicRoutes): + routes: object = client.abridge_routes() if not isinstance(routes, dict): raise TypeError( f"{type(client).__name__}.abridge_routes() must return a dict[str, handler]" @@ -370,15 +384,25 @@ def _make_forwarder( async def forward(request: FastAPIRequest) -> Response: body = await _read_json(request) + if body is None: + # Body was present but not a JSON object (unparseable or a non-object + # like an array/string). Fail at the boundary with a precise error + # instead of silently coercing to {} and confusing the upstream. + return JSONResponse( + {"error": {"message": "request body must be a JSON object"}}, + status_code=400, + ) try: # SIO event name IS the path; the host's `Proxy` has a # matching handler registered under the same name. The wire # payload is just the decoded object — no wrapping envelope # beyond request correlation and no HTTP metadata. - result = await asyncio.wait_for( - ns.request(path, body), timeout=request_timeout - ) + # Single timeout source: thread the configured value into + # `ns.request` itself. A redundant outer `wait_for` here was a no-op + # above `ns.request`'s own (smaller) default, silently capping the + # configurable `request_timeout` at that default. + result = await ns.request(path, body, timeout=request_timeout) except TimeoutError: message = "tunnel timed out waiting for host" logger.warning("abridge tunnel %s: %s", path, message) @@ -395,16 +419,20 @@ async def forward(request: FastAPIRequest) -> Response: return forward -async def _read_json(request: FastAPIRequest) -> dict[str, Any]: +async def _read_json(request: FastAPIRequest) -> dict[str, Any] | None: + """Decode the body as a JSON object. `{}` for a genuinely empty body; `None` + when the body is present but not a JSON object (unparseable, or a valid + non-object like an array/string). The caller turns `None` into a 400 so the + failure surfaces at the boundary instead of as a silently-coerced `{}`.""" raw = await request.body() if not raw: return {} try: parsed = json.loads(raw) except ValueError: - return {} + return None if not isinstance(parsed, dict): - return {} + return None return parsed @@ -420,14 +448,12 @@ def _to_http_response(result: object) -> Response: def _status_from_remote_error(exc: RemoteSioError) -> int: - """`RemoteSioError(type, message)` carries no status code. We map - well-known exception type names to HTTP statuses; everything else - becomes 502.""" - if exc.type == "UpstreamError": - # The client raised UpstreamError. The message format may include - # the status code, but it's not structured. Default 502. - return 502 - return 502 + """Use the upstream HTTP status the host handler chose. `AbridgeError` + carries `status_code` (429/400/404/...), which the host threads through the + wire error envelope and the sandbox preserves on `RemoteSioError`. Fall back + to 502 only when the remote error carried no status.""" + status = getattr(exc, "status_code", None) + return status if isinstance(status, int) else 502 # ── host-side: Proxy ───────────────────────────────────────────────────── @@ -583,13 +609,29 @@ async def start(self, sandbox: Sandbox) -> TunnelHandle: Calling `start` again while this proxy is active returns the current handle instead of leaking another tunnel. If startup fails, clients with `aclose()` are still closed. + + ORDERING: open the proxy before any other `sandbox.remote()` / health + call. The `/abridge` namespace must be registered before the runtime + client connects, so `proxy.start` / `proxy.session` has to run first; a + prior remote call that already connected the client makes this raise. """ if self._handle is not None: return self._handle self._clients_closed = False try: - sandbox.register_namespace(self) + try: + sandbox.register_namespace(self) + except RuntimeError as exc: + # register_namespace only raises RuntimeError once the runtime + # client has connected — i.e. a remote()/health() ran first. + # Re-raise in this surface's vocabulary instead of the low-level + # "before entering the async context" message. + raise RuntimeError( + "open the abridge proxy (proxy.session/start) before any other " + "sandbox.remote()/health() call — its /abridge namespace must be " + "registered before the runtime client connects" + ) from exc handle = await sandbox.remote(_start_tunnel, paths=list(self.paths)) except BaseException: try: @@ -651,7 +693,18 @@ async def session(self, sandbox: Sandbox) -> AsyncIterator[TunnelHandle]: handle = await self.start(sandbox) try: yield handle - finally: + except BaseException: + # The body is already raising: a teardown failure must not mask it + # (it would demote the real error to __context__). Log-and-swallow + # stop() errors here. Catching the body exception explicitly (vs + # probing sys.exc_info()) keeps this path off callers that open the + # session inside an unrelated `except` block. + try: + await self.stop(sandbox) + except Exception: + logger.exception("abridge: session teardown failed after body error (suppressed)") + raise + else: await self.stop(sandbox) # ── handy property ──────────────────────────────────────────────── @@ -671,6 +724,7 @@ def url(self) -> str: "AbridgeError", "Client", "ClientResponse", + "DynamicRoutes", "Handler", "NAMESPACE", "Proxy", diff --git a/plugins/abridge/agentix/bridge/sidecar.py b/plugins/abridge/agentix/bridge/sidecar.py index 78f98ba..0005136 100644 --- a/plugins/abridge/agentix/bridge/sidecar.py +++ b/plugins/abridge/agentix/bridge/sidecar.py @@ -210,7 +210,11 @@ async def __aenter__(self) -> str: finally: self._starting = False - raise SidecarError(f"sidecar exhausted port retries on {self._host}") + # Unreachable: the bounded retry loop always returns (success) or + # re-raises (the last attempt can't `continue` — that guard requires + # `attempt < attempts`). Kept as an assert so it documents the invariant + # and satisfies the `-> str` contract without masking a real error. + raise AssertionError("unreachable: sidecar start loop must return or raise") async def __aexit__(self, *exc: object) -> None: await self._terminate() @@ -330,6 +334,9 @@ async def _terminate(self) -> None: raise finally: await self._finish_drainers() + # Clear the handle so a re-entered/inspected sidecar doesn't read + # the dead process as "still running". + self._proc = None __all__ = ["Command", "Sidecar", "SidecarError"] diff --git a/plugins/abridge/tests/test_sidecar_forward.py b/plugins/abridge/tests/test_sidecar_forward.py index e882928..35ab309 100644 --- a/plugins/abridge/tests/test_sidecar_forward.py +++ b/plugins/abridge/tests/test_sidecar_forward.py @@ -19,6 +19,7 @@ Forward, Proxy, Request, + SessionForward, Sidecar, SidecarError, TunnelHandle, @@ -121,7 +122,8 @@ async def fake_post(url, *, json, headers): assert resp.body == b'{"error":"down"}' -async def test_forward_network_error_is_502(monkeypatch) -> None: +async def test_forward_network_error_is_503(monkeypatch) -> None: + """Failure to reach the sidecar is 503 — distinct from a relayed upstream 502.""" fwd = Forward("http://side.car", paths=["/v1/messages"]) async def fake_post(url, *, json, headers): @@ -130,7 +132,7 @@ async def fake_post(url, *, json, headers): monkeypatch.setattr(fwd._client, "post", fake_post) with pytest.raises(AbridgeError) as ei: await fwd.abridge_routes()["/v1/messages"](_req("/v1/messages", {})) - assert ei.value.status_code == 502 + assert ei.value.status_code == 503 async def test_forward_http_status_survives_tunnel_and_sio(monkeypatch) -> None: @@ -179,6 +181,49 @@ async def host_emit(event, data=None, **kwargs): await fwd.aclose() +async def test_abridge_error_status_survives_tunnel_and_sio(monkeypatch) -> None: + """An `AbridgeError(429)` raised by a host handler reaches the agent's HTTP + response as a 429 through the real `:error` reply path — the sandbox must + not collapse it to a blanket 502.""" + import agentix.bridge.proxy as proxy_mod + + import agentix as agentix_mod + + monkeypatch.setattr(agentix_mod, "register_namespace", lambda ns: None) + monkeypatch.setattr(proxy_mod, "_namespace_singleton", None) + + class _RateLimited: + def abridge_routes(self): + return {"/v1/messages": self.handle} + + async def handle(self, request: Request) -> ClientResponse: + raise AbridgeError("rate limited upstream", status_code=429) + + host = Proxy(_RateLimited()) + handle = await proxy_mod._start_tunnel(paths=list(host.paths)) + sandbox_ns = proxy_mod._get_namespace() + + async def sandbox_emit(event, data=None): + await host.trigger_event(event, data) + + async def host_emit(event, data=None, **kwargs): + if event.endswith(":result"): + await sandbox_ns._on_reply_success(data) + elif event.endswith(":error"): + await sandbox_ns._on_reply_error(data) + + monkeypatch.setattr(sandbox_ns, "emit", sandbox_emit) + monkeypatch.setattr(host, "emit", host_emit) + + try: + async with httpx.AsyncClient(base_url=handle.url, timeout=10) as client: + response = await client.post("/v1/messages", json={"model": "claude"}) + assert response.status_code == 429 + assert response.json()["error"]["message"] == "rate limited upstream" + finally: + await proxy_mod._stop_tunnel(handle=handle) + + class _CloseAwareClient: def __init__(self) -> None: self.close_calls = 0 @@ -234,6 +279,30 @@ async def test_proxy_closes_clients_when_start_fails() -> None: assert client.close_calls == 1 +async def test_proxy_session_teardown_error_surfaces_inside_except_block() -> None: + """A teardown failure after a SUCCESSFUL body must raise even when the + caller sits inside an `except` block (a common retry/fallback shape) — + ambient exception state must not be mistaken for a body failure.""" + + class _ExplodingClose: + def abridge_routes(self): + return {"/x": self.handle} + + async def handle(self, request: Request) -> ClientResponse: + return ClientResponse.json({"ok": True}) + + async def aclose(self) -> None: + raise RuntimeError("close boom") + + proxy = Proxy(_ExplodingClose()) + try: + raise ValueError("ambient exception being handled") + except ValueError: + with pytest.raises(ExceptionGroup): + async with proxy.session(_FakeSandbox()): + pass + + async def test_forward_pool_can_reopen_after_idempotent_close() -> None: fwd = Forward("http://side.car", paths=["/v1/messages"]) original = fwd._client @@ -284,3 +353,465 @@ async def test_forward_through_live_sidecar(tmp_path) -> None: assert resp.media_type == "application/json" finally: await fwd.aclose() + + +# ── SessionForward (unit, mocked httpx) ─────────────────────────────── + + +async def test_session_forward_creates_session_then_rewrites_path(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + calls: list = [] + + async def fake_post(url, *, json, headers): + calls.append((url, json, headers)) + if url.endswith("/sessions"): + return httpx.Response(200, content=b'{"session_id": "S9"}', headers={"content-type": "application/json"}) + return httpx.Response(200, content=b'{"ok": true}', headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + resp = await fwd.abridge_routes()["/v1/chat/completions"]( + _req("/v1/chat/completions", {"model": "qwen3-4b"}) + ) + + assert isinstance(resp, ClientResponse) + assert resp.status_code == 200 and resp.body == b'{"ok": true}' + assert fwd.session_id == "S9" + # First upstream call created the session; second routed into it by path. + assert calls[0][0] == "http://gw/sessions" + assert calls[1][0] == "http://gw/sessions/S9/v1/chat/completions" + assert calls[1][1] == {"model": "qwen3-4b"} + assert calls[1][2]["x-session-id"] == "S9" + + +async def test_session_forward_creates_session_once(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + creates = 0 + + async def fake_post(url, *, json, headers): + nonlocal creates + if url.endswith("/sessions"): + creates += 1 + return httpx.Response(200, content=b'{"session_id": "S"}', headers={"content-type": "application/json"}) + return httpx.Response(200, content=b"{}", headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + handler = fwd.abridge_routes()["/v1/chat/completions"] + await handler(_req("/v1/chat/completions", {})) + await handler(_req("/v1/chat/completions", {})) + assert creates == 1 + + +async def test_session_forward_open_precreates_session(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + + async def fake_post(url, *, json, headers): + return httpx.Response(200, content=b'{"session_id": "PRE"}', headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + assert await fwd.open() == "PRE" + assert fwd.session_id == "PRE" + + +async def test_session_forward_accepts_any_2xx_create(monkeypatch) -> None: + """A gateway that answers the create POST with 201 (or any 2xx) is a + success, not a 502.""" + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + + async def fake_post(url, *, json, headers): + return httpx.Response(201, content=b'{"session_id": "S201"}', headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + assert await fwd.open() == "S201" + + +async def test_session_forward_attaches_to_assigned_session(monkeypatch) -> None: + """Assigning `session_id` attaches to that existing gateway session: + forwarded requests are scoped to it with NO create call — the assigned id + must not be silently overwritten by a fresh create.""" + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + fwd.session_id = "EXT" + calls: list = [] + + async def fake_post(url, *, json, headers): + calls.append((url, headers)) + return httpx.Response(200, content=b'{"ok": true}', headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + resp = await fwd.abridge_routes()["/v1/chat/completions"]( + _req("/v1/chat/completions", {"model": "m"}) + ) + + assert resp.status_code == 200 + assert fwd.session_id == "EXT" + # exactly one upstream call — straight into the attached session, no create. + assert [url for url, _ in calls] == ["http://gw/sessions/EXT/v1/chat/completions"] + assert calls[0][1]["x-session-id"] == "EXT" + + +async def test_session_forward_create_failure_is_502(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + + async def fake_post(url, *, json, headers): + return httpx.Response(500, content=b"boom") + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + with pytest.raises(AbridgeError) as ei: + await fwd.open() + assert ei.value.status_code == 502 + + +async def test_session_forward_missing_id_field_is_502(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + + async def fake_post(url, *, json, headers): + return httpx.Response(200, content=b'{"nope": 1}', headers={"content-type": "application/json"}) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + with pytest.raises(AbridgeError) as ei: + await fwd.open() + assert ei.value.status_code == 502 + + +def test_session_forward_session_id_before_open_raises() -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + with pytest.raises(RuntimeError): + _ = fwd.session_id + + +async def test_session_forward_delete_session_reaps_and_resets(monkeypatch) -> None: + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + deleted: list = [] + + async def fake_post(url, *, json, headers): + return httpx.Response(200, content=b'{"session_id": "S"}', headers={"content-type": "application/json"}) + + async def fake_delete(url, *, headers): + deleted.append(url) + return httpx.Response(204) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + monkeypatch.setattr(fwd._client, "delete", fake_delete) + + assert await fwd.open() == "S" + await fwd.delete_session() + assert deleted == ["http://gw/sessions/S"] + # after reaping, the id is gone again — a premature read raises. + with pytest.raises(RuntimeError): + _ = fwd.session_id + + +async def test_delete_session_after_close_does_not_resurrect_pool(monkeypatch) -> None: + """The documented harvest flow — Proxy.stop (pool closed), harvest, then + `delete_session()` — must not leave behind a lazily recreated pool that no + remaining lifecycle will ever close.""" + fwd = SessionForward("http://gw", paths=["/v1/chat/completions"]) + deleted: list = [] + + async def fake_post(url, *, json, headers): + return httpx.Response(200, content=b'{"session_id": "S"}', headers={"content-type": "application/json"}) + + async def fake_delete(self, url, *, headers=None): + deleted.append(url) + return httpx.Response(204) + + assert fwd._client is not None + monkeypatch.setattr(fwd._client, "post", fake_post) + monkeypatch.setattr(httpx.AsyncClient, "delete", fake_delete) + + assert await fwd.open() == "S" + await fwd.aclose() + await fwd.delete_session() + + assert deleted == ["http://gw/sessions/S"] + assert fwd._client is None + + +# ── SessionForward (integration, real sidecar) ──────────────────────── + +SESSION_SERVER = """ +import sys, json +from http.server import BaseHTTPRequestHandler, HTTPServer + +SID = "sess-LIVE" + +class H(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200); self.end_headers(); self.wfile.write(b"ok") + + def do_POST(self): + n = int(self.headers.get("content-length", 0)) + body = self.rfile.read(n) + self.send_response(200) + self.send_header("content-type", "application/json") + self.end_headers() + if self.path == "/sessions": + self.wfile.write(json.dumps({"session_id": SID}).encode()) + else: + self.wfile.write(json.dumps({"path": self.path, "got": json.loads(body or b"{}")}).encode()) + + def log_message(self, *a): + pass + +HTTPServer(("127.0.0.1", int(sys.argv[1])), H).serve_forever() +""" + + +async def test_session_forward_through_live_sidecar(tmp_path) -> None: + script = tmp_path / "sess.py" + script.write_text(SESSION_SERVER) + async with Sidecar(command=[sys.executable, str(script), "{port}"]) as url: + fwd = SessionForward(url, paths=["/v1/chat/completions"]) + try: + resp = await fwd.abridge_routes()["/v1/chat/completions"]( + _req("/v1/chat/completions", {"model": "m"}) + ) + assert resp.status_code == 200 + assert fwd.session_id == "sess-LIVE" + assert b'"path": "/sessions/sess-LIVE/v1/chat/completions"' in resp.body + assert b'"model": "m"' in resp.body + finally: + await fwd.aclose() + + +# ── tunnel boundary + client env helpers ────────────────────────────── + + +async def test_tunnel_rejects_non_object_body(monkeypatch) -> None: + """A present-but-non-object JSON body (an array) is a 400 at the tunnel, + not a silent coercion to {}.""" + import agentix.bridge.proxy as proxy_mod + + import agentix as agentix_mod + + monkeypatch.setattr(agentix_mod, "register_namespace", lambda ns: None) + monkeypatch.setattr(proxy_mod, "_namespace_singleton", None) + + handle = await proxy_mod._start_tunnel(paths=["/v1/messages"]) + try: + async with httpx.AsyncClient(base_url=handle.url, timeout=10) as client: + r = await client.post("/v1/messages", json=["not", "an", "object"]) + assert r.status_code == 400 + finally: + await proxy_mod._stop_tunnel(handle=handle) + + +def test_openai_client_environ_bakes_in_v1() -> None: + """OpenAIClient.environ mirrors the Anthropic ones and bakes in the /v1 suffix.""" + from agentix.bridge.clients import OpenAIClient + + c = OpenAIClient(base_url="https://up.stream/v1", api_key="real-key", model="gpt-4o") + env = c.environ(TunnelHandle(url="http://127.0.0.1:9", port=9)) + assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9/v1" + assert env["OPENAI_API_KEY"].startswith("sk-") + + +# ── composition: Convert (AnthropicToOpenAI) ∘ transport (Forward/SessionForward) ── + +_OPENAI_COMPLETION = ( + b'{"id":"c1","object":"chat.completion","model":"qwen3-4b",' + b'"choices":[{"index":0,"finish_reason":"stop",' + b'"message":{"role":"assistant","content":"hi there"}}],' + b'"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}' +) + + +def test_forward_handler_accessor() -> None: + fwd = Forward("http://x", paths=["/v1/chat/completions"]) + assert callable(fwd.handler()) # sole path, no arg needed + assert callable(fwd.handler("/v1/chat/completions")) + with pytest.raises(ValueError): + Forward("http://x", paths=["/a", "/b"]).handler() # ambiguous → must name a path + + +async def test_anthropic_to_openai_translates_over_any_downstream() -> None: + """The Convert layer is transport-blind: it translates Anthropic→OpenAI, calls + the downstream Handler, and translates the OpenAI completion back to Anthropic.""" + from agentix.bridge.clients import AnthropicToOpenAI + + captured: dict = {} + + async def fake_downstream(request: Request) -> ClientResponse: + captured["body"] = request.body + return ClientResponse(body=_OPENAI_COMPLETION, media_type="application/json") + + conv = AnthropicToOpenAI(fake_downstream, model="qwen3-4b") + resp = await conv.messages( + _req("/v1/messages", { + "model": "claude-3-5-sonnet", + "max_tokens": 64, + "messages": [{"role": "user", "content": "say hi"}], + }) + ) + assert resp.status_code == 200 + assert b"hi there" in resp.body + assert b'"role": "assistant"' in resp.body # Anthropic response shape + # the downstream saw an OpenAI-shaped, non-streaming body with the model override + assert captured["body"]["model"] == "qwen3-4b" + assert captured["body"]["stream"] is False + assert captured["body"]["messages"][0]["role"] == "user" + + +async def test_anthropic_to_openai_composes_with_session_forward(monkeypatch) -> None: + """End-to-end composition: Anthropic agent → AnthropicToOpenAI → SessionForward + creates the session and rewrites the path → OpenAI completion → back to Anthropic. + The converter never touches a session; the SessionForward never sees Anthropic.""" + from agentix.bridge.clients import AnthropicToOpenAI + + tito = SessionForward("http://gw", paths=["/v1/chat/completions"]) + calls: list = [] + + async def fake_post(url, *, json, headers): + calls.append(url) + if url.endswith("/sessions"): + return httpx.Response(200, content=b'{"session_id": "S1"}', headers={"content-type": "application/json"}) + return httpx.Response(200, content=_OPENAI_COMPLETION, headers={"content-type": "application/json"}) + + assert tito._client is not None + monkeypatch.setattr(tito._client, "post", fake_post) + + conv = AnthropicToOpenAI(tito.handler(), model="qwen3-4b") + resp = await conv.messages( + _req("/v1/messages", { + "model": "claude", + "max_tokens": 64, + "messages": [{"role": "user", "content": "say hi"}], + }) + ) + assert resp.status_code == 200 + assert b"hi there" in resp.body + assert tito.session_id == "S1" + assert calls[0] == "http://gw/sessions" + assert calls[1] == "http://gw/sessions/S1/v1/chat/completions" + await tito.aclose() + + +async def test_proxy_stop_closes_forwarder_pool_through_converter() -> None: + """`Proxy(AnthropicToOpenAI(tito.handler()))` — the converter must forward + `aclose()` to its downstream's owner, or the SessionForward httpx pool + outlives `Proxy.stop`. The session itself is untouched (harvest survives).""" + from agentix.bridge.clients import AnthropicToOpenAI + + tito = SessionForward("http://gw", paths=["/v1/chat/completions"]) + pool = tito._client + assert pool is not None + + proxy = Proxy(AnthropicToOpenAI(tito.handler(), model="m")) + async with proxy.session(_FakeSandbox()): + assert not pool.is_closed + + assert pool.is_closed + + +async def test_anthropic_to_openai_replays_remembered_assistant() -> None: + """The lossy Anthropic round-trip drops reasoning_content + tool-call `index`, so + a reconstructed assistant won't byte-match what a session backend stored. The + converter remembers the exact downstream assistant and replays it verbatim on the + next turn, keyed by the surviving tool-call ids.""" + from agentix.bridge.clients import AnthropicToOpenAI + + stored_assistant = { + "role": "assistant", + "content": "", + "reasoning_content": "\n\n", + "tool_calls": [{ + "id": "call_X", "index": 0, "type": "function", + "function": {"name": "python", "arguments": "{\"expression\": \"1+1\"}"}, + }], + } + sent: list = [] + + async def fake_downstream(request: Request) -> ClientResponse: + sent.append(request.body.get("messages")) + return ClientResponse.json({ + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": stored_assistant}], + "model": "m", "usage": {}, + }) + + conv = AnthropicToOpenAI(fake_downstream, model="m") + await conv.messages(_req("/v1/messages", { + "model": "c", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}], + })) + # turn 1: the agent resends the round-tripped assistant (no index, no reasoning_content) + await conv.messages(_req("/v1/messages", { + "model": "c", "max_tokens": 64, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "call_X", "name": "python", "input": {"expression": "1+1"}}]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "call_X", "content": "2"}]}, + ], + })) + + asst = [m for m in sent[1] if m.get("role") == "assistant"] + assert len(asst) == 1 + # value-equal to the stored assistant (incl. reasoning_content + index), not the + # lossy reconstruction — so the backend's byte-match would pass. + assert asst[0] == stored_assistant + assert asst[0]["reasoning_content"] == "\n\n" + assert asst[0]["tool_calls"][0]["index"] == 0 + + +async def test_anthropic_to_openai_replay_survives_reused_tool_call_ids() -> None: + """Some OpenAI-compatible servers reuse tool-call ids across turns (TGI + famously emits id "0"). The replay memory must not key on ids alone, or a + later turn overwrites an earlier one and every matching assistant in the + resent history is replaced by the LAST remembered message — silently + corrupting the recorded conversation.""" + from agentix.bridge.clients import AnthropicToOpenAI + + def _assistant(args: str, reasoning: str) -> dict: + return { + "role": "assistant", + "content": "", + "reasoning_content": reasoning, + "tool_calls": [{ + "id": "0", "index": 0, "type": "function", + "function": {"name": "python", "arguments": args}, + }], + } + + # Distinct calls, same id "0"; arguments spacing differs from the + # round-trip's json.dumps rendering to prove key canonicalization. + assistant_1 = _assistant('{"expression":"1+1"}', "r1") + assistant_2 = _assistant('{"expression":"2+2"}', "r2") + replies = [assistant_1, assistant_2, _assistant('{"expression":"3+3"}', "r3")] + sent: list = [] + + async def fake_downstream(request: Request) -> ClientResponse: + sent.append(request.body.get("messages")) + return ClientResponse.json({ + "choices": [{"index": 0, "finish_reason": "tool_calls", + "message": replies[len(sent) - 1]}], + "model": "m", "usage": {}, + }) + + def _tool_use(args: dict) -> dict: + return {"role": "assistant", "content": [ + {"type": "tool_use", "id": "0", "name": "python", "input": args}]} + + def _tool_result(text: str) -> dict: + return {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "0", "content": text}]} + + conv = AnthropicToOpenAI(fake_downstream, model="m") + turn = [{"role": "user", "content": "hi"}] + await conv.messages(_req("/v1/messages", {"model": "c", "max_tokens": 64, "messages": turn})) + turn = [*turn, _tool_use({"expression": "1+1"}), _tool_result("2")] + await conv.messages(_req("/v1/messages", {"model": "c", "max_tokens": 64, "messages": turn})) + turn = [*turn, _tool_use({"expression": "2+2"}), _tool_result("4")] + await conv.messages(_req("/v1/messages", {"model": "c", "max_tokens": 64, "messages": turn})) + + asst = [m for m in sent[2] if m.get("role") == "assistant"] + assert len(asst) == 2 + # Each history slot replays ITS OWN remembered message, not the latest one. + assert asst[0] == assistant_1 + assert asst[1] == assistant_2 diff --git a/tests/test_sio_reply_error.py b/tests/test_sio_reply_error.py new file mode 100644 index 0000000..a9ffcf6 --- /dev/null +++ b/tests/test_sio_reply_error.py @@ -0,0 +1,48 @@ +"""`Namespace._on_reply_error` threads the error envelope's `status_code` +onto `RemoteSioError`, so protocol layers built on `Namespace.request` (e.g. +abridge's sandbox tunnel) can reply with the real upstream status instead of +collapsing every remote failure to one blanket code.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from agentix.sio import Namespace, RemoteSioError + + +class _Ns(Namespace): + namespace = "/reply-error-test" + + +async def test_reply_error_carries_status_code() -> None: + ns = _Ns() + fut: asyncio.Future = asyncio.get_running_loop().create_future() + ns._pending_requests["r1"] = fut + + await ns._on_reply_error({ + "request_id": "r1", + "error": {"type": "AbridgeError", "message": "rate limited", "status_code": 429}, + }) + + with pytest.raises(RemoteSioError) as ei: + fut.result() + assert ei.value.status_code == 429 + assert ei.value.type == "AbridgeError" + assert ei.value.message == "rate limited" + + +async def test_reply_error_non_int_status_code_is_none() -> None: + ns = _Ns() + fut: asyncio.Future = asyncio.get_running_loop().create_future() + ns._pending_requests["r2"] = fut + + await ns._on_reply_error({ + "request_id": "r2", + "error": {"type": "Boom", "message": "no status", "status_code": "429"}, + }) + + with pytest.raises(RemoteSioError) as ei: + fut.result() + assert ei.value.status_code is None