diff --git a/fern/versions/latest/pages/model-server/adapters-rewrites.mdx b/fern/versions/latest/pages/model-server/adapters-rewrites.mdx new file mode 100644 index 0000000000..3a4cf1482c --- /dev/null +++ b/fern/versions/latest/pages/model-server/adapters-rewrites.mdx @@ -0,0 +1,74 @@ +--- +title: "Request-Rewriting Interceptors" +description: "Interceptors that mutate the outbound request body before it reaches the upstream." +position: 6 +--- + +Interceptors that mutate the outbound request body before it reaches the upstream. Run in the REQUEST stage. + +| Name | Stage | Purpose | +|------|-------|---------| +| `drop_params` | request | Remove named parameters from the outbound body. | +| `payload_modifier` | request | Add, remove, and rename body fields. | +| `system_message` | request | Inject a system message (prepend / append / replace). | +| `consolidate_system` | request | Merge displaced system messages into one at position 0. | +| `modify_tools` | request | Strip or add properties on `tools[].function.parameters`. | +| `turn_counter` | request | Per-session turn budget; raises `GracefulError` (→ 429) on exhaustion. | + +### `drop_params` + +```yaml +- name: drop_params + config: + params: ["top_k", "frequency_penalty"] +``` + +### `payload_modifier` + +```yaml +- name: payload_modifier + config: + params_to_remove: ["secret_field"] + params_to_add: { "max_completion_tokens": 4096 } + params_to_rename: { "old_name": "new_name" } +``` + +### `system_message` + +```yaml +- name: system_message + config: + system_message: "You are a careful assistant." + strategy: prepend # one of: prepend | append | replace +``` + +### `consolidate_system` + +```yaml +- name: consolidate_system + config: + separator: "\n\n" +``` + +### `modify_tools` + +```yaml +- name: modify_tools + config: + strip_properties: ["internal_flag"] + add_properties: + reasoning: + type: string + description: "model's chain-of-thought" +``` + +### `turn_counter` + +```yaml +- name: turn_counter + config: + every: 1 # log on every Nth turn + max_turns: 20 # null disables the budget +``` + +The session key is taken from `ctx.extra["session_id"]` (set by the middleware when the path matches `/s//...`); otherwise a body-hash fallback is used. diff --git a/fern/versions/latest/pages/model-server/adapters.mdx b/fern/versions/latest/pages/model-server/adapters.mdx new file mode 100644 index 0000000000..dfb521bb28 --- /dev/null +++ b/fern/versions/latest/pages/model-server/adapters.mdx @@ -0,0 +1,68 @@ +--- +title: "Adapter Middleware" +description: "Interceptor-based middleware framework for Model, Agent, and Resources servers." +position: 3 +--- + +Adapter middleware adds an interceptor chain to a server's request/response path. Each interceptor can observe or mutate the payload without the host server changing. Use it to inject system prompts, drop unsupported params, cache responses, count turns, normalize reasoning fields, log tokens, and so on. + +`adapters` is opt-in. It is declared on `BaseResponsesAPIModelConfig`, `BaseResponsesAPIAgentConfig`, and `BaseResourcesServerConfig`, so every in-tree server inheriting from `SimpleResponsesAPIModel` / `SimpleResponsesAPIAgent` / `SimpleResourcesServer` accepts an `adapters` block automatically. Omitting it leaves behavior identical to the base server. + +## Quickstart + +Add an `adapters` list to any server config: + +```yaml +policy_model: + responses_api_models: + openai_model: + openai_base_url: https://api.openai.com/v1 + openai_api_key: ... + openai_model: gpt-4.1 + adapters: # ← toggle: omit or null = OFF + - {name: logging, config: {}} +``` + +This PR ships the **framework** (`AdapterPipeline`, `InterceptorRegistry`, `install_middleware`, `start_adapter_proxy`) plus two built-in interceptors: + +| Name | Stage | Purpose | +|------|-------|---------| +| `logging` | request + response | Log request body keys and response status/latency. Canonical "did the chain fire?" probe. | +| `endpoint` | request → response | Drive the upstream HTTP call directly. **Only used by `start_adapter_proxy`** (standalone host mode); forbidden inside `install_middleware`. | + +Additional interceptor families (observability, caching, request rewriting) ship in follow-on PRs. + +## Host modes + +The same `AdapterPipeline` runs in either of two hosting modes: + +**Middleware mode** — `install_middleware(app, adapters)` attaches the pipeline to an existing FastAPI app via `app.middleware("http")`. The host server's own routing performs the upstream call via `call_next`. Used by every Model/Agent/Resources server when `adapters` is set on its config. + +**Proxy mode** — `start_adapter_proxy(upstream_url, adapters)` launches a localhost uvicorn that hosts the pipeline with its own forwarding logic. Used by agents that bring their own SDK client (e.g. `claude_code_agent` with `anthropic_base_url`); the agent points its SDK's `*_BASE_URL` at the proxy URL via the `adapter_proxy` field on `BaseResponsesAPIAgentConfig`. + +## Path-Based Session Scoping + +Posts to `/s//` have the prefix stripped before forwarding, and `` is recorded as `ctx.extra["session_id"]`. Follow-on interceptors that key per-session (e.g. `turn_counter`, `caching`) use this id. + +## Custom Interceptors + +Register a class at runtime via `InterceptorRegistry.register`: + +```python +from nemo_gym.adapters import InterceptorRegistry + +InterceptorRegistry.register("my_interceptor", "myproject.adapters.my_interceptor") +``` + +The target module must expose a class named `Interceptor` that subclasses one of `RequestInterceptor`, `RequestToResponseInterceptor`, or `ResponseInterceptor` from `nemo_gym.adapters.types`. The class is instantiated with the YAML `config` dict as kwargs. + +## Configuration Reference + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `adapters` | `list[dict] \| null` | `null` | Ordered interceptor specs on each server config. Each entry is `{name: , config: }`. `null` or `[]` disables the middleware. | +| `adapter_proxy` | `AdapterProxyConfig \| null` | `null` | On `BaseResponsesAPIAgentConfig` only. Configures a localhost proxy in front of an external inference upstream. Fields: `upstream_url`, `adapters`, `host` (default `127.0.0.1`), `port` (default `0` → kernel-assigned), `request_timeout`, `unsafe_allow_remote`. | + + +`start_adapter_proxy` refuses any `host` other than `127.0.0.1`/`localhost` unless you pass `unsafe_allow_remote=True`. The proxy forwards the client's `Authorization` header verbatim to the upstream, so binding to `0.0.0.0` would leak the upstream API key to any caller on the network. + diff --git a/fern/versions/latest/pages/model-server/index.mdx b/fern/versions/latest/pages/model-server/index.mdx index b1c54832d0..547735e905 100644 --- a/fern/versions/latest/pages/model-server/index.mdx +++ b/fern/versions/latest/pages/model-server/index.mdx @@ -32,3 +32,9 @@ Self-hosted inference with vLLM for maximum control. [Model Server Fields](/reference/configuration#model-server-fields) for server configuration syntax and fields. + +## Middleware + +Model servers can attach an interceptor chain that observes or mutates request and response payloads — useful for logging, caching, system-prompt injection, turn budgeting, reasoning-field normalization, and similar cross-cutting hooks. + +See [Adapter Middleware](/model-server/adapters) for the framework, built-in interceptors, and configuration syntax. diff --git a/nemo_gym/adapters/__init__.py b/nemo_gym/adapters/__init__.py new file mode 100644 index 0000000000..807a706deb --- /dev/null +++ b/nemo_gym/adapters/__init__.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Gym adapter framework — interceptor-based middleware for responses_api_models. + +Public API: + install_middleware(app, interceptor_specs) — attach pipeline to a FastAPI app + start_adapter_proxy(upstream_url, adapters) — host pipeline as a localhost uvicorn + AdapterPipeline — the async interceptor chain + InterceptorRegistry — name → class resolution +""" + +from nemo_gym.adapters.middleware import install_middleware +from nemo_gym.adapters.pipeline import AdapterPipeline +from nemo_gym.adapters.proxy import ProxyHandle, start_adapter_proxy +from nemo_gym.adapters.registry import InterceptorRegistry +from nemo_gym.adapters.types import ( + AdapterProxyConfig, + AdapterRequest, + AdapterResponse, + GracefulError, + InterceptorContext, + InterceptorSpec, + RequestInterceptor, + RequestToResponseInterceptor, + ResponseInterceptor, + Stage, +) + + +__all__ = [ + "AdapterPipeline", + "AdapterProxyConfig", + "AdapterRequest", + "AdapterResponse", + "GracefulError", + "InterceptorContext", + "InterceptorRegistry", + "InterceptorSpec", + "ProxyHandle", + "RequestInterceptor", + "RequestToResponseInterceptor", + "ResponseInterceptor", + "Stage", + "install_middleware", + "start_adapter_proxy", +] diff --git a/nemo_gym/adapters/interceptors/__init__.py b/nemo_gym/adapters/interceptors/__init__.py new file mode 100644 index 0000000000..cb1039b6f9 --- /dev/null +++ b/nemo_gym/adapters/interceptors/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Built-in interceptors for the adapter pipeline.""" diff --git a/nemo_gym/adapters/interceptors/consolidate_system.py b/nemo_gym/adapters/interceptors/consolidate_system.py new file mode 100644 index 0000000000..c522dc369e --- /dev/null +++ b/nemo_gym/adapters/interceptors/consolidate_system.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Merge all system messages into one at position 0. + +Models like Qwen3 reject system messages that appear mid-conversation +(``System message must be at the beginning``). +""" + +from __future__ import annotations + +import logging + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + + +def _content_to_str(content) -> str: + """Extract plain text from either a string or OpenAI list-format content.""" + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for item in content: + if isinstance(item, dict): + parts.append(item.get("text", "")) + elif isinstance(item, str): + parts.append(item) + return "\n".join(parts) + return str(content) if content else "" + + +class Interceptor(RequestInterceptor): + def __init__(self, *, separator: str = "\n\n") -> None: + self._sep = separator + self._fix_count = 0 + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + messages: list[dict] = req.body.get("messages", []) + if not messages: + return req + + system_indices: list[int] = [] + system_parts: list[str] = [] + non_system: list[dict] = [] + for i, msg in enumerate(messages): + if msg.get("role") == "system": + system_indices.append(i) + text = _content_to_str(msg.get("content", "")) + if text: + system_parts.append(text) + else: + non_system.append(msg) + + if len(system_indices) <= 1 and (not system_indices or system_indices[0] == 0): + return req + + self._fix_count += 1 + session = req.ctx.extra.get("session_id", "?") + + role_seq = [m.get("role", "?") for m in messages] + displaced = [i for i in system_indices if i > 0] + diag_parts: list[str] = [] + for idx in displaced: + lo = max(0, idx - 2) + hi = min(len(role_seq), idx + 3) + window = " ".join(f"[{j}]={role_seq[j]}" for j in range(lo, hi)) + content_preview = _content_to_str(messages[idx].get("content", ""))[:300] + diag_parts.append(f"system@{idx} window=({window}) content_preview={content_preview!r}") + + logger.warning( + "consolidate_system: fix #%d session=%s n_msgs=%d system_at=%s n_system=%d roles_head=%s | %s", + self._fix_count, + session, + len(messages), + system_indices, + len(system_parts), + role_seq[:8], + " | ".join(diag_parts) if diag_parts else "system missing from pos 0", + ) + + merged: list[dict] = [] + if system_parts: + merged.append({"role": "system", "content": self._sep.join(system_parts)}) + merged.extend(non_system) + req.body["messages"] = merged + return req diff --git a/nemo_gym/adapters/interceptors/drop_params.py b/nemo_gym/adapters/interceptors/drop_params.py new file mode 100644 index 0000000000..52ee734149 --- /dev/null +++ b/nemo_gym/adapters/interceptors/drop_params.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + + +class Interceptor(RequestInterceptor): + def __init__(self, *, params: list[str]) -> None: + self._params = set(params) + self._logged_once = False + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + for p in self._params: + req.body.pop(p, None) + + if not self._logged_once: + logger.info("drop_params: removing %s", sorted(self._params)) + self._logged_once = True + + return req diff --git a/nemo_gym/adapters/interceptors/endpoint.py b/nemo_gym/adapters/interceptors/endpoint.py new file mode 100644 index 0000000000..03f3b90480 --- /dev/null +++ b/nemo_gym/adapters/interceptors/endpoint.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +import json +import logging +import time +from typing import Any + +import aiohttp # type-name imports only (ClientTimeout, ClientError); HTTP calls go through nemo_gym.server_utils.request + +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + RequestToResponseInterceptor, +) +from nemo_gym.server_utils import request as global_request + + +logger = logging.getLogger(__name__) + + +class Interceptor(RequestToResponseInterceptor): + def __init__( + self, + *, + upstream_url: str, + api_key: str | None = None, + extra_body: dict[str, Any] | None = None, + request_timeout: float = 120, + max_retries: int = 0, + retry_on_status: list[int] | None = None, + max_concurrent: int = 64, + ) -> None: + clean = upstream_url.rstrip("/") + for suffix in ("/chat/completions", "/completions", "/embeddings"): + if clean.endswith(suffix): + clean = clean[: -len(suffix)] + break + self._upstream_url = clean + self._api_key = api_key + self._extra_body = extra_body or {} + self._request_timeout = float(request_timeout) + self._timeout = aiohttp.ClientTimeout(total=request_timeout) + self._max_retries = max_retries + self._retry_on_status = set(retry_on_status or [429, 502, 503, 504]) + # ``max_concurrent`` is config-shape compatibility only; connector + # limits live on the global aiohttp client. + self._max_concurrent = max_concurrent + + @staticmethod + def _normalize_content(body: dict[str, Any]) -> None: + for choice in body.get("choices", []): + msg = choice.get("message") or choice.get("delta") or {} + if "content" in msg and msg["content"] is None: + msg["content"] = "" + + async def intercept_request( + self, + req: AdapterRequest, + ) -> AdapterRequest | AdapterResponse: + url = f"{self._upstream_url}{req.path}" + + body = {**req.body, **self._extra_body} + headers = { + k: v for k, v in req.headers.items() if k.lower() not in ("host", "content-length", "transfer-encoding") + } + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + headers.setdefault("Content-Type", "application/json") + + attempt = 0 + while True: + t0 = time.perf_counter() + try: + resp = await global_request( + method="POST", + url=url, + data=json.dumps(body), + headers=headers, + timeout=self._timeout, + ) + async with resp: + raw = await resp.read() + latency = (time.perf_counter() - t0) * 1000 + # Iterate ``resp.headers.items()`` to preserve multi-valued + # keys (e.g. Set-Cookie) that a plain ``dict()`` collapses. + resp_headers: list[tuple[bytes, bytes]] = [ + (k.encode("latin-1"), v.encode("latin-1")) for k, v in resp.headers.items() + ] + status = resp.status + + if status in self._retry_on_status and attempt < self._max_retries: + retry_after = resp.headers.get("Retry-After") + delay = float(retry_after) if retry_after else min(2**attempt, 60) + logger.warning( + "endpoint: %s returned %d, retry %d/%d in %.1fs", + url, + status, + attempt + 1, + self._max_retries, + delay, + ) + attempt += 1 + await asyncio.sleep(delay) + continue + + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, UnicodeDecodeError): + parsed = raw + + if isinstance(parsed, dict): + self._normalize_content(parsed) + + return AdapterResponse( + status_code=status, + headers=resp_headers, + body=parsed, + latency_ms=latency, + ctx=req.ctx, + ) + + except asyncio.TimeoutError: + latency = (time.perf_counter() - t0) * 1000 + if attempt < self._max_retries: + delay = min(2**attempt, 60) + logger.warning( + "endpoint: %s timed out, retry %d/%d in %.1fs", + url, + attempt + 1, + self._max_retries, + delay, + ) + attempt += 1 + await asyncio.sleep(delay) + continue + logger.error("endpoint: %s timed out after %d attempts", url, attempt + 1) + return AdapterResponse( + status_code=504, + headers={}, + body={ + "error": {"message": f"Upstream timed out after {self._request_timeout}s", "type": "timeout"} + }, + latency_ms=latency, + ctx=req.ctx, + ) + + except aiohttp.ClientError as exc: + latency = (time.perf_counter() - t0) * 1000 + if attempt < self._max_retries: + delay = min(2**attempt, 60) + logger.warning( + "endpoint: %s failed (%s), retry %d/%d in %.1fs", + url, + exc, + attempt + 1, + self._max_retries, + delay, + ) + attempt += 1 + await asyncio.sleep(delay) + continue + raise + + async def close(self) -> None: + return None diff --git a/nemo_gym/adapters/interceptors/modify_tools.py b/nemo_gym/adapters/interceptors/modify_tools.py new file mode 100644 index 0000000000..a674fb82f4 --- /dev/null +++ b/nemo_gym/adapters/interceptors/modify_tools.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import copy +import logging +from typing import Any + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + + +def _apply_modifications( + tools: list[dict[str, Any]], + strip: frozenset[str], + add: dict[str, dict[str, Any]], +) -> int: + count = 0 + for tool in tools: + fn = tool.get("function") or tool + params = fn.get("parameters") or {} + props = params.get("properties") + if not isinstance(props, dict): + continue + + req: list[str] | None = params.get("required") + + for field in strip: + if field in props: + del props[field] + count += 1 + if isinstance(req, list) and field in req: + req.remove(field) + + for field, schema in add.items(): + if field not in props: + props[field] = copy.deepcopy(schema) + count += 1 + return count + + +class Interceptor(RequestInterceptor): + def __init__( + self, + *, + strip_properties: list[str] | None = None, + add_properties: dict[str, dict[str, Any]] | None = None, + ) -> None: + self._strip = frozenset(strip_properties or []) + self._add: dict[str, dict[str, Any]] = add_properties or {} + self._logged_once = False + + if not self._strip and not self._add: + logger.warning("modify_tools: no modifications configured — interceptor is a no-op") + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + tools = req.body.get("tools") + if tools: + n = _apply_modifications(tools, self._strip, self._add) + if n and not self._logged_once: + logger.info( + "modify_tools: applied %d change(s) (strip=%s, add=%s)", + n, + sorted(self._strip), + sorted(self._add), + ) + self._logged_once = True + return req diff --git a/nemo_gym/adapters/interceptors/payload_modifier.py b/nemo_gym/adapters/interceptors/payload_modifier.py new file mode 100644 index 0000000000..2c5ded6600 --- /dev/null +++ b/nemo_gym/adapters/interceptors/payload_modifier.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging +from typing import Any + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + + +def _remove_keys(obj: Any, keys: set[str]) -> Any: + if isinstance(obj, dict): + return {k: _remove_keys(v, keys) for k, v in obj.items() if k not in keys} + if isinstance(obj, list): + return [_remove_keys(item, keys) for item in obj] + return obj + + +def _rename_keys(obj: Any, mapping: dict[str, str]) -> Any: + if isinstance(obj, dict): + return {mapping.get(k, k): _rename_keys(v, mapping) for k, v in obj.items()} + if isinstance(obj, list): + return [_rename_keys(item, mapping) for item in obj] + return obj + + +class Interceptor(RequestInterceptor): + def __init__( + self, + *, + params_to_remove: list[str] | None = None, + params_to_add: dict[str, Any] | None = None, + params_to_rename: dict[str, str] | None = None, + ) -> None: + self._remove = set(params_to_remove or []) + self._add: dict[str, Any] = params_to_add or {} + self._rename: dict[str, str] = params_to_rename or {} + self._logged_once = False + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + if self._remove: + req.body = _remove_keys(req.body, self._remove) + + if self._rename: + req.body = _rename_keys(req.body, self._rename) + + if self._add: + req.body.update(self._add) + + if not self._logged_once: + logger.info( + "payload_modifier: remove=%s, add=%s, rename=%s", + sorted(self._remove), + sorted(self._add), + sorted(self._rename), + ) + self._logged_once = True + + return req diff --git a/nemo_gym/adapters/interceptors/request_logging.py b/nemo_gym/adapters/interceptors/request_logging.py new file mode 100644 index 0000000000..0846253e10 --- /dev/null +++ b/nemo_gym/adapters/interceptors/request_logging.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +import logging + +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + RequestInterceptor, + ResponseInterceptor, +) + + +logger = logging.getLogger(__name__) + +_MAX = 512 + + +def _trunc_preview(obj: object) -> str: + if isinstance(obj, bytes): + text = obj.decode(errors="replace") + elif isinstance(obj, dict): + text = json.dumps(obj, default=str, ensure_ascii=False) + else: + text = str(obj) + if len(text) <= _MAX: + return text + return text[:_MAX] + "..." + + +class Interceptor(RequestInterceptor, ResponseInterceptor): + best_effort = True + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + keys = list(req.body.keys()) if isinstance(req.body, dict) else [] + logger.info( + "request %s %s body_keys=%s body_preview=%s", + req.method, + req.path, + keys, + _trunc_preview(req.body), + ) + return req + + async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse: + logger.info( + "response status=%d latency_ms=%.2f body_preview=%s", + resp.status_code, + resp.latency_ms, + _trunc_preview(resp.body), + ) + return resp diff --git a/nemo_gym/adapters/interceptors/system_message.py b/nemo_gym/adapters/interceptors/system_message.py new file mode 100644 index 0000000000..b816784ef7 --- /dev/null +++ b/nemo_gym/adapters/interceptors/system_message.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + +_VALID_STRATEGIES = {"replace", "append", "prepend"} + + +class Interceptor(RequestInterceptor): + def __init__( + self, + *, + system_message: str, + strategy: str = "prepend", + ) -> None: + if strategy not in _VALID_STRATEGIES: + raise ValueError(f"Invalid strategy {strategy!r}, must be one of {_VALID_STRATEGIES}") + self._message = system_message + self._strategy = strategy + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + messages: list = req.body.setdefault("messages", []) + sys_msg = {"role": "system", "content": self._message} + + if self._strategy == "replace": + non_system = [m for m in messages if m.get("role") != "system"] + req.body["messages"] = [sys_msg] + non_system + + elif self._strategy == "append": + messages.append(sys_msg) + + else: # prepend + messages.insert(0, sys_msg) + + return req diff --git a/nemo_gym/adapters/interceptors/turn_counter.py b/nemo_gym/adapters/interceptors/turn_counter.py new file mode 100644 index 0000000000..b02c41aa1b --- /dev/null +++ b/nemo_gym/adapters/interceptors/turn_counter.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +import hashlib +import logging +import time +from dataclasses import dataclass, field +from typing import Any + +from nemo_gym.adapters.types import AdapterRequest, RequestInterceptor + + +logger = logging.getLogger(__name__) + +_WARN_THRESHOLD = 0.80 +_URGENT_THRESHOLD = 0.95 +_STALE_SESSION_SEC = 900.0 +_GC_INTERVAL_SEC = 300.0 + + +def _session_key_from_body(body: dict[str, Any]) -> str: + """Fallback: derive a session key from the first non-system message.""" + messages = body.get("messages") or [] + for msg in messages: + if msg.get("role") == "system": + continue + content = msg.get("content", "") + if isinstance(content, list): + content = "".join(part.get("text", "") for part in content if isinstance(part, dict)) + if content: + return hashlib.sha256(content.encode()).hexdigest()[:8] + return "unknown" + + +@dataclass +class _Session: + count: int = 0 + last_time: float = field(default_factory=time.monotonic) + + +class Interceptor(RequestInterceptor): + def __init__( + self, + *, + every: int = 1, + max_turns: int | None = None, + ) -> None: + self._every = max(every, 1) + self._max = max_turns + self._sessions: dict[str, _Session] = {} + self._lock = asyncio.Lock() + self._last_gc = time.monotonic() + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + key = req.ctx.extra.get("session_id") + if not key: + key = _session_key_from_body(req.body) + logger.warning("no session_id in context — falling back to body-hash key %s", key) + + async with self._lock: + now = time.monotonic() + if now - self._last_gc > _GC_INTERVAL_SEC: + self._gc(now) + self._last_gc = now + + sess = self._sessions.setdefault(key, _Session()) + sess.count += 1 + n = sess.count + dt = now - sess.last_time if sess.count > 1 else 0.0 + sess.last_time = now + active = len(self._sessions) + + if n % self._every == 0 or n == 1: + cap = f"/{self._max}" if self._max else "" + elapsed = f" (+{dt:.1f}s)" if dt > 0 else "" + logger.info( + "task %s turn %d%s%s (%d active)", + key, + n, + cap, + elapsed, + active, + ) + + if self._max is None: + return req + + if n > self._max: + logger.warning( + "task %s REJECTED turn %d (max_turns=%d exceeded)", + key, + n, + self._max, + ) + from nemo_gym.adapters.types import GracefulError + + raise GracefulError( + f"Turn budget exhausted: {n}/{self._max} turns used. " + f"The evaluation framework has terminated this agent session." + ) + + remaining = self._max - n + messages = req.body.get("messages") + if not isinstance(messages, list): + return req + + ratio = n / self._max + if ratio >= _URGENT_THRESHOLD: + messages.append( + { + "role": "system", + "content": ( + f"[SYSTEM] URGENT: Turn {n}/{self._max} — only {remaining} turn(s) left. " + f"You MUST provide your final answer NOW. Do not start new work." + ), + } + ) + elif ratio >= _WARN_THRESHOLD: + messages.append( + { + "role": "system", + "content": ( + f"[SYSTEM] Turn {n}/{self._max} — {remaining} turns remaining. " + f"Begin wrapping up: finish current work and prepare your final answer." + ), + } + ) + + return req + + def _gc(self, now: float) -> None: + """Remove sessions idle longer than ``_STALE_SESSION_SEC``.""" + stale = [k for k, s in self._sessions.items() if now - s.last_time > _STALE_SESSION_SEC] + for k in stale: + del self._sessions[k] + if stale: + logger.debug("turn_counter GC: removed %d stale sessions, %d remaining", len(stale), len(self._sessions)) diff --git a/nemo_gym/adapters/middleware.py b/nemo_gym/adapters/middleware.py new file mode 100644 index 0000000000..c5fbd56785 --- /dev/null +++ b/nemo_gym/adapters/middleware.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapter pipeline middleware for responses_api_models FastAPI apps.""" + +from __future__ import annotations + +import json +import logging +import re +from typing import Any + +from fastapi import FastAPI, Request +from starlette.responses import JSONResponse, Response + +from nemo_gym.adapters.pipeline import AdapterPipeline +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + GracefulError, + InterceptorContext, +) + + +logger = logging.getLogger(__name__) + +# Hop-by-hop headers (RFC 7230 §6.1) plus framing/encoding fields the ASGI +# layer rewrites itself. ``content-encoding`` is included because aiohttp +# auto-decodes upstream bodies — re-emitting the original encoding would +# mislead any gzip-aware client into trying to decompress plain bytes. +_HOP_BY_HOP_HEADERS = frozenset( + { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "content-length", + "content-encoding", + "server", + } +) + +_SESSION_PATH_RE = re.compile(r"^/s/([a-f0-9]+)(/.*)?$") + + +def _override_request_body(request: Request, new_body: bytes) -> None: + """Make downstream handlers see *new_body* instead of the original.""" + request._body = new_body # type: ignore[attr-defined] + if hasattr(request, "_json"): + delattr(request, "_json") + + new_len = str(len(new_body)).encode("ascii") + scope_headers = [(k, v) for k, v in request.scope.get("headers", []) if k.lower() != b"content-length"] + scope_headers.append((b"content-length", new_len)) + request.scope["headers"] = scope_headers + + +async def _starlette_response_to_adapter( + starlette_resp: Response, + ctx: InterceptorContext, +) -> AdapterResponse: + chunks: list[bytes] = [] + body_iter = getattr(starlette_resp, "body_iterator", None) + if body_iter is not None: + async for chunk in body_iter: + if isinstance(chunk, str): + chunk = chunk.encode("utf-8") + chunks.append(chunk) + else: + raw = getattr(starlette_resp, "body", b"") + if isinstance(raw, str): + raw = raw.encode("utf-8") + chunks.append(raw) + raw_body = b"".join(chunks) + + # Preserve duplicate header keys (e.g. multiple Set-Cookie) — dict + # collapse would drop all but the last value. + headers: list[tuple[bytes, bytes]] = list(starlette_resp.raw_headers) + + content_type = "" + for name, value in headers: + if name.lower() == b"content-type": + content_type = value.decode("latin-1").lower() + break + + body: dict[str, Any] | bytes + if "application/json" in content_type or content_type == "": + try: + body = json.loads(raw_body.decode("utf-8")) if raw_body else {} + except (ValueError, UnicodeDecodeError): + body = raw_body + else: + body = raw_body + + return AdapterResponse( + status_code=starlette_resp.status_code, + headers=headers, + body=body, + latency_ms=0.0, + ctx=ctx, + ) + + +def _adapter_response_to_starlette(resp: AdapterResponse) -> Response: + raw = resp.headers or [] + if isinstance(raw, dict): + header_pairs: list[tuple[bytes, bytes]] = [(k.encode("latin-1"), v.encode("latin-1")) for k, v in raw.items()] + else: + header_pairs = list(raw) + + fwd_pairs: list[tuple[bytes, bytes]] = [ + (name, value) for name, value in header_pairs if name.decode("latin-1").lower() not in _HOP_BY_HOP_HEADERS + ] + + if isinstance(resp.body, bytes): + out = Response(content=resp.body, status_code=resp.status_code) + else: + out = JSONResponse(content=resp.body, status_code=resp.status_code) + + # Keep Starlette's framing content-length; drop framing content-type + # only if the adapter chain supplied one. + fwd_keys_lower = {name.lower() for name, _ in fwd_pairs} + framing_kept = [ + (k, v) for k, v in out.raw_headers if k.lower() == b"content-length" or k.lower() not in fwd_keys_lower + ] + out.raw_headers = framing_kept + fwd_pairs + return out + + +def install_middleware( + app: FastAPI, + interceptor_specs: list[dict[str, Any]] | None, +) -> None: + """Install the adapter pipeline as FastAPI middleware on *app*. + + ``interceptor_specs`` is a list of ``{"name": ..., "config": ...}`` dicts. + Empty or ``None`` disables the middleware. + """ + if not interceptor_specs: + return + + # ``endpoint`` performs its own upstream call; inside a host server this + # would double-forward. The proxy host mode (``start_adapter_proxy``) + # appends it itself. + if any(s.get("name") == "endpoint" for s in interceptor_specs): + raise ValueError( + "install_middleware: the 'endpoint' interceptor cannot be used in a " + "middleware-hosted chain — the model server already forwards upstream. " + "Drop it from `adapters` or use start_adapter_proxy() instead." + ) + + pipeline = AdapterPipeline.from_config(interceptor_specs) + + @app.middleware("http") + async def _adapter_middleware(request: Request, call_next): # noqa: ANN202 + if request.method != "POST": + return await call_next(request) + + try: + body = await request.json() + except Exception: + return JSONResponse( + {"error": {"message": "Invalid JSON body", "type": "invalid_request_error"}}, + status_code=400, + ) + + path = request.url.path + ctx = InterceptorContext() + + m = _SESSION_PATH_RE.match(path) + if m: + ctx.extra["session_id"] = m.group(1) + path = m.group(2) or "/" + + adapter_req = AdapterRequest( + method=request.method, + path=path, + headers=dict(request.headers), + body=body, + ctx=ctx, + ) + + async def _upstream(req: AdapterRequest) -> AdapterResponse: + new_body = json.dumps(req.body).encode("utf-8") + _override_request_body(request, new_body) + starlette_resp = await call_next(request) + return await _starlette_response_to_adapter(starlette_resp, req.ctx) + + try: + resp = await pipeline.process(adapter_req, upstream_call=_upstream) + except GracefulError as exc: + logger.warning( + "Adapter middleware: graceful termination for session %s: %s", + ctx.extra.get("session_id", "?"), + exc, + ) + return JSONResponse( + { + "error": { + "message": str(exc), + "type": "invalid_request_error", + "code": "session_budget_exhausted", + }, + }, + status_code=429, + ) + except Exception: + logger.exception("Adapter pipeline error") + return JSONResponse( + {"error": {"message": "Internal adapter middleware error", "type": "server_error"}}, + status_code=500, + ) + + return _adapter_response_to_starlette(resp) diff --git a/nemo_gym/adapters/pipeline.py b/nemo_gym/adapters/pipeline.py new file mode 100644 index 0000000000..6b87dd868d --- /dev/null +++ b/nemo_gym/adapters/pipeline.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Async adapter pipeline — ordered interceptor chain execution.""" + +from __future__ import annotations + +import logging +from typing import Any, Awaitable, Callable + +from nemo_gym.adapters.registry import InterceptorRegistry +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + InterceptorContext, + RequestInterceptor, + RequestToResponseInterceptor, + ResponseInterceptor, + Stage, + set_context, +) + + +UpstreamCall = Callable[[AdapterRequest], Awaitable[AdapterResponse]] + +logger = logging.getLogger(__name__) + + +def _stage_of( + interceptor: RequestInterceptor | RequestToResponseInterceptor | ResponseInterceptor, +) -> Stage: + return getattr(interceptor, "stage", Stage.REQUEST) + + +class AdapterPipeline: + """Async interceptor chain: REQUEST → REQUEST_TO_RESPONSE → RESPONSE.""" + + def __init__( + self, + interceptors: list[RequestInterceptor | RequestToResponseInterceptor | ResponseInterceptor], + ) -> None: + self._chain = list(interceptors) + self._validate_order() + logger.info( + "AdapterPipeline ready (%d interceptors: %s)", + len(self._chain), + [getattr(i, "_registry_name", type(i).__name__) for i in self._chain], + ) + + @classmethod + def from_config( + cls, + interceptor_specs: list[dict[str, Any]], + ) -> AdapterPipeline: + chain: list[RequestInterceptor | RequestToResponseInterceptor | ResponseInterceptor] = [] + for spec in interceptor_specs: + name = spec["name"] + config = spec.get("config") or {} + chain.append(InterceptorRegistry.create(name, config)) + return cls(chain) + + _STAGE_ORDER = [Stage.REQUEST, Stage.REQUEST_TO_RESPONSE, Stage.RESPONSE] + + def _validate_order(self) -> None: + current_idx = 0 + for interceptor in self._chain: + stage = _stage_of(interceptor) + try: + idx = self._STAGE_ORDER.index(stage) + except ValueError: + raise ValueError(f"Unknown stage {stage!r} on {type(interceptor).__name__}") + if idx < current_idx: + raise ValueError( + f"Invalid interceptor order: {type(interceptor).__name__} " + f"(stage={stage.value}) appears after stage " + f"{self._STAGE_ORDER[current_idx].value}. " + f"Required order: request → request_to_response → response" + ) + current_idx = max(current_idx, idx) + + async def process( + self, + request: AdapterRequest, + upstream_call: UpstreamCall | None = None, + ) -> AdapterResponse: + """Run *request* through the chain and return the response. + + If no ``RequestToResponseInterceptor`` short-circuits, ``upstream_call`` + is invoked with the (possibly mutated) request. Response interceptors + then run in reverse order. + """ + ctx = InterceptorContext(request_id=request.ctx.request_id) + set_context(ctx) + + current: AdapterRequest | AdapterResponse = request + + for interceptor in self._chain: + if isinstance(current, AdapterResponse): + break + + if isinstance(interceptor, (RequestInterceptor, RequestToResponseInterceptor)): + try: + result = await interceptor.intercept_request(current) # type: ignore[arg-type] + current = result + except Exception: + if getattr(interceptor, "best_effort", False): + logger.warning( + "Interceptor %s failed (best_effort=True), continuing", + type(interceptor).__name__, + exc_info=True, + ) + continue + raise + + if not isinstance(current, AdapterResponse): + if upstream_call is None: + raise RuntimeError("No interceptor produced a response. Is 'endpoint' in the chain?") + current = await upstream_call(current) + + response = current + response_interceptors = [ic for ic in reversed(self._chain) if isinstance(ic, ResponseInterceptor)] + for interceptor in response_interceptors: + try: + response = await interceptor.intercept_response(response) + except Exception: + if getattr(interceptor, "best_effort", False): + logger.warning( + "Response interceptor %s failed (best_effort=True), continuing", + type(interceptor).__name__, + exc_info=True, + ) + continue + raise + + return response diff --git a/nemo_gym/adapters/proxy.py b/nemo_gym/adapters/proxy.py new file mode 100644 index 0000000000..4efe9837ae --- /dev/null +++ b/nemo_gym/adapters/proxy.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Localhost adapter proxy. + +Hosts the same ``AdapterPipeline`` as ``install_middleware`` in its own +uvicorn server. Used by agents whose SDK clients respect a ``*_BASE_URL`` +env var (Anthropic / OpenAI / Cohere): the agent's outbound traffic is +pointed at the proxy URL, the proxy applies the chain, the proxy's own +upstream-call closure forwards to the real upstream. + +Non-chat paths (``/v1/models``, batches, health checks) pass through to +the upstream verbatim so SDK pre-flight works. +""" + +from __future__ import annotations + +import asyncio +import logging +import socket +import threading +import time +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any + +import aiohttp +import uvicorn +from fastapi import FastAPI, Request +from starlette.responses import JSONResponse, Response + +from nemo_gym.adapters.pipeline import AdapterPipeline +from nemo_gym.adapters.types import AdapterRequest, AdapterResponse, InterceptorContext + + +logger = logging.getLogger(__name__) + +# Routes whose POST traffic runs through the adapter pipeline. Everything +# else (any method, any path) is forwarded upstream verbatim so SDK +# pre-flight (model listing, batches, health checks) still works. +_ADAPTED_ROUTES = frozenset( + { + "/v1/chat/completions", + "/v1/completions", + "/v1/responses", + "/v1/messages", + "/v1/embeddings", + } +) + + +@dataclass +class ProxyHandle: + """Handle for a running adapter proxy.""" + + url: str + port: int + _server: uvicorn.Server + _thread: threading.Thread + + def stop(self, timeout: float = 5.0) -> None: + self._server.should_exit = True + self._thread.join(timeout=timeout) + + def __enter__(self) -> "ProxyHandle": + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.stop() + + +def start_adapter_proxy( + upstream_url: str, + adapters: list[dict[str, Any]], + *, + host: str = "127.0.0.1", + port: int = 0, + request_timeout: float = 120.0, + health_timeout: float = 10.0, + unsafe_allow_remote: bool = False, +) -> ProxyHandle: + """Launch a localhost uvicorn that hosts the adapter pipeline. + + The proxy forwards adapted requests to ``upstream_url`` via its own + ``aiohttp.ClientSession`` (created inside the proxy thread's event loop). + Forwarding happens outside the user's interceptor chain via ``pipeline. + process(req, upstream_call=...)`` — symmetric to ``install_middleware`` + using FastAPI's ``call_next``. + """ + if host not in ("127.0.0.1", "localhost") and not unsafe_allow_remote: + raise ValueError( + f"start_adapter_proxy: refusing host={host!r} — the proxy forwards the " + "client's Authorization header upstream, so binding to a non-localhost " + "interface leaks credentials. Pass unsafe_allow_remote=True to override." + ) + + if any(s.get("name") == "endpoint" for s in adapters): + raise ValueError( + "start_adapter_proxy: the 'endpoint' interceptor cannot be used here — " + "the proxy performs upstream forwarding itself. Drop it from `adapters`." + ) + + pipeline = AdapterPipeline.from_config(list(adapters)) + + app = _build_app(pipeline, upstream_url.rstrip("/"), request_timeout) + actual_port = _bind_port(host, port) + + config = uvicorn.Config( + app, + host=host, + port=actual_port, + log_level="warning", + access_log=False, + ) + server = uvicorn.Server(config) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + + url = f"http://{host}:{actual_port}" + _wait_for_health(url, timeout=health_timeout) + logger.info("adapter proxy ready upstream=%s url=%s", upstream_url, url) + + return ProxyHandle(url=url, port=actual_port, _server=server, _thread=thread) + + +def _bind_port(host: str, preferred: int) -> int: + if preferred: + return preferred + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind((host, 0)) + return s.getsockname()[1] + + +def _wait_for_health(url: str, timeout: float) -> None: + import urllib.request + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(f"{url}/_proxy_health", timeout=1) as r: + if r.status == 200: + return + except Exception: + time.sleep(0.05) + raise RuntimeError(f"adapter proxy at {url} did not become healthy in {timeout:.1f}s") + + +def _build_app(pipeline: AdapterPipeline, upstream_url: str, request_timeout: float) -> FastAPI: + @asynccontextmanager + async def _lifespan(app_: FastAPI): + app_.state.session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=request_timeout)) + try: + yield + finally: + await app_.state.session.close() + + app = FastAPI(lifespan=_lifespan) + app.state.pipeline = pipeline + app.state.upstream_url = upstream_url + app.state.request_timeout = request_timeout + + @app.get("/_proxy_health") + async def _health() -> JSONResponse: + return JSONResponse({"ok": True}) + + async def _dispatch(path: str, request: Request) -> Response: + norm_path = "/" + path.lstrip("/") + if request.method == "POST" and norm_path in _ADAPTED_ROUTES: + return await _run_pipeline(request, norm_path) + return await _passthrough(request, norm_path) + + app.add_api_route( + "/{path:path}", + _dispatch, + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], + ) + + return app + + +async def _run_pipeline(request: Request, path: str) -> Response: + pipeline: AdapterPipeline = request.app.state.pipeline + session: aiohttp.ClientSession = request.app.state.session + upstream_url: str = request.app.state.upstream_url + + try: + body = await request.json() + except Exception: + return JSONResponse( + {"error": {"message": "Invalid JSON body", "type": "invalid_request_error"}}, + status_code=400, + ) + + adapter_req = AdapterRequest( + method=request.method, + path=path, + headers=dict(request.headers), + body=body, + ctx=InterceptorContext(), + ) + + async def _upstream(req: AdapterRequest) -> AdapterResponse: + target = f"{upstream_url}{req.path}" + fwd_headers = {k: v for k, v in req.headers.items() if k.lower() not in ("host", "content-length")} + import json as _json + + t0 = time.perf_counter() + async with session.post(target, data=_json.dumps(req.body), headers=fwd_headers) as resp: + raw = await resp.read() + latency = (time.perf_counter() - t0) * 1000 + # Preserve multi-valued response headers (Set-Cookie etc.) + resp_headers: list[tuple[bytes, bytes]] = [ + (k.encode("latin-1"), v.encode("latin-1")) for k, v in resp.headers.items() + ] + try: + parsed: Any = _json.loads(raw) if raw else {} + except (ValueError, UnicodeDecodeError): + parsed = raw + + if isinstance(parsed, dict): + for choice in parsed.get("choices", []): + msg = choice.get("message") or choice.get("delta") or {} + if isinstance(msg, dict) and msg.get("content") is None and "content" in msg: + msg["content"] = "" + + return AdapterResponse( + status_code=resp.status, + headers=resp_headers, + body=parsed, + latency_ms=latency, + ctx=req.ctx, + ) + + try: + resp = await pipeline.process(adapter_req, upstream_call=_upstream) + except Exception: + logger.exception("adapter proxy pipeline error") + return JSONResponse( + {"error": {"message": "Internal adapter proxy error", "type": "server_error"}}, + status_code=500, + ) + + fwd_headers: list[tuple[bytes, bytes]] = [] + if isinstance(resp.headers, list): + for name, value in resp.headers: + lname = name.decode("latin-1").lower() + if lname in ("content-length", "transfer-encoding", "content-encoding"): + continue + fwd_headers.append((name, value)) + elif isinstance(resp.headers, dict): + for k, v in resp.headers.items(): + if k.lower() in ("content-length", "transfer-encoding", "content-encoding"): + continue + fwd_headers.append((k.encode("latin-1"), v.encode("latin-1"))) + + if isinstance(resp.body, bytes): + out: Response = Response(content=resp.body, status_code=resp.status_code) + else: + out = JSONResponse(content=resp.body, status_code=resp.status_code) + + fwd_keys_lower = {k.lower() for k, _ in fwd_headers} + framing_kept = [ + (k, v) for k, v in out.raw_headers if k.lower() == b"content-length" or k.lower() not in fwd_keys_lower + ] + out.raw_headers = framing_kept + fwd_headers + return out + + +async def _passthrough(request: Request, path: str) -> Response: + """Forward request to upstream verbatim — for SDK pre-flight paths.""" + session: aiohttp.ClientSession = request.app.state.session + upstream_url: str = request.app.state.upstream_url + timeout: float = request.app.state.request_timeout + + body = await request.body() + fwd_headers = {k: v for k, v in request.headers.items() if k.lower() not in ("host", "content-length")} + + target = f"{upstream_url}{path}" + if request.url.query: + target = f"{target}?{request.url.query}" + + try: + async with session.request( + method=request.method, + url=target, + data=body if body else None, + headers=fwd_headers, + ) as resp: + raw = await resp.read() + out_headers: list[tuple[bytes, bytes]] = [] + for k, v in resp.headers.items(): + if k.lower() in ("content-length", "transfer-encoding", "content-encoding"): + continue + out_headers.append((k.encode("latin-1"), v.encode("latin-1"))) + out = Response(content=raw, status_code=resp.status) + out.raw_headers = [(hk, hv) for hk, hv in out.raw_headers if hk.lower() == b"content-length"] + out_headers + return out + except asyncio.TimeoutError: + return JSONResponse( + {"error": {"message": f"Upstream timed out after {timeout}s", "type": "timeout"}}, + status_code=504, + ) + except aiohttp.ClientError as exc: + logger.warning("passthrough %s failed: %s", target, exc) + return JSONResponse( + {"error": {"message": f"Upstream error: {exc}", "type": "upstream_error"}}, + status_code=502, + ) + + +__all__ = ["start_adapter_proxy", "ProxyHandle"] diff --git a/nemo_gym/adapters/registry.py b/nemo_gym/adapters/registry.py new file mode 100644 index 0000000000..d9adb271cd --- /dev/null +++ b/nemo_gym/adapters/registry.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Interceptor registry — maps short names to interceptor classes.""" + +from __future__ import annotations + +import importlib +import logging +from typing import Any, Type + +from nemo_gym.adapters.types import ( + RequestInterceptor, + RequestToResponseInterceptor, + ResponseInterceptor, +) + + +logger = logging.getLogger(__name__) + +InterceptorClass = Type[RequestInterceptor | RequestToResponseInterceptor | ResponseInterceptor] + +# Each module must expose a class named ``Interceptor``. +# +# This dict holds only the framework-level builtins. Interceptor families +# (observability / caching / request-rewriting) add themselves via plain +# ``_BUILTIN[] = `` assignments below the dict literal, +# so each follow-on PR touches its own anchor lines and merges cleanly with +# the others. +_BUILTIN: dict[str, str] = { + # ``endpoint`` drives the upstream HTTP call from inside the pipeline. + # Required for ``start_adapter_proxy`` (standalone host mode); forbidden + # inside ``install_middleware`` because the host server already forwards. + "endpoint": "nemo_gym.adapters.interceptors.endpoint", + # ``logging`` is the canonical "did the chain fire?" probe used by + # framework-level tests. Lightweight; logs request keys + response status. + "logging": "nemo_gym.adapters.interceptors.request_logging", +} + +# Family extensions — follow-on PRs append entries to ``_BUILTIN`` here. +# Each family adds entries under its own family-named comment marker so +# different families don't fight for the same diff context. + +# Request-rewriting family — mutate the outbound request body before +# upstream sees it. +_BUILTIN["drop_params"] = "nemo_gym.adapters.interceptors.drop_params" +_BUILTIN["payload_modifier"] = "nemo_gym.adapters.interceptors.payload_modifier" +_BUILTIN["system_message"] = "nemo_gym.adapters.interceptors.system_message" +_BUILTIN["consolidate_system"] = "nemo_gym.adapters.interceptors.consolidate_system" +_BUILTIN["modify_tools"] = "nemo_gym.adapters.interceptors.modify_tools" +_BUILTIN["turn_counter"] = "nemo_gym.adapters.interceptors.turn_counter" + +# External / plugin registrations at runtime. +_EXTRA: dict[str, str] = {} + + +class InterceptorRegistry: + @staticmethod + def register(name: str, module_path: str) -> None: + _EXTRA[name] = module_path + + @staticmethod + def resolve_class(name: str) -> InterceptorClass: + module_path = _EXTRA.get(name) or _BUILTIN.get(name) + if module_path is None: + available = sorted(set(_BUILTIN) | set(_EXTRA)) + raise ValueError(f"Unknown interceptor {name!r}. Available: {available}") + try: + mod = importlib.import_module(module_path) + except ImportError as exc: + raise ValueError(f"Cannot import interceptor module {module_path!r} for {name!r}: {exc}") from exc + cls = getattr(mod, "Interceptor", None) + if cls is None: + raise ValueError(f"Module {module_path!r} does not expose an 'Interceptor' class") + return cls + + @staticmethod + def create( + name: str, config: dict[str, Any] | None = None + ) -> RequestInterceptor | RequestToResponseInterceptor | ResponseInterceptor: + cls = InterceptorRegistry.resolve_class(name) + instance = cls(**(config or {})) + instance._registry_name = name + return instance + + @staticmethod + def available() -> list[str]: + return sorted(set(_BUILTIN) | set(_EXTRA)) diff --git a/nemo_gym/adapters/types.py b/nemo_gym/adapters/types.py new file mode 100644 index 0000000000..f05e684bd4 --- /dev/null +++ b/nemo_gym/adapters/types.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Core types and ABCs for the adapter interceptor pipeline.""" + +from __future__ import annotations + +import enum +import uuid +from abc import ABC, abstractmethod +from contextvars import ContextVar +from dataclasses import dataclass, field +from typing import Any + +from pydantic import BaseModel, Field + + +@dataclass +class InterceptorContext: + """Per-request state shared across interceptors via ContextVar.""" + + request_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) + extra: dict[str, Any] = field(default_factory=dict) + + +_current_context: ContextVar[InterceptorContext] = ContextVar("adapter_ctx") + + +def get_context() -> InterceptorContext: + try: + return _current_context.get() + except LookupError: + ctx = InterceptorContext() + _current_context.set(ctx) + return ctx + + +def set_context(ctx: InterceptorContext) -> None: + _current_context.set(ctx) + + +@dataclass +class AdapterRequest: + method: str + path: str + headers: dict[str, str] + body: dict[str, Any] + ctx: InterceptorContext = field(default_factory=get_context) + + +@dataclass +class AdapterResponse: + # ``headers`` is a list of byte-tuples (Starlette's ``raw_headers`` shape) + # so multi-valued headers like ``Set-Cookie`` survive. A plain ``dict`` is + # also accepted for convenience. + status_code: int + headers: list[tuple[bytes, bytes]] | dict[str, str] + body: dict[str, Any] | bytes + latency_ms: float = 0.0 + ctx: InterceptorContext = field(default_factory=get_context) + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 400 + + +class Stage(enum.Enum): + REQUEST = "request" + REQUEST_TO_RESPONSE = "request_to_response" + RESPONSE = "response" + + +class RequestInterceptor(ABC): + stage: Stage = Stage.REQUEST + stream_safe: bool = True + best_effort: bool = False + + @abstractmethod + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: ... + + +class RequestToResponseInterceptor(ABC): + """Request-phase interceptor that may short-circuit by returning a response.""" + + stage: Stage = Stage.REQUEST_TO_RESPONSE + stream_safe: bool = True + best_effort: bool = False + + @abstractmethod + async def intercept_request( + self, + req: AdapterRequest, + ) -> AdapterRequest | AdapterResponse: ... + + +class ResponseInterceptor(ABC): + stage: Stage = Stage.RESPONSE + stream_safe: bool = True + best_effort: bool = False + + @abstractmethod + async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse: ... + + +class GracefulError(Exception): + """Terminate the request with a 429 (e.g. session budget exhausted).""" + + +class InterceptorSpec(BaseModel): + """One entry in an ``adapters`` config list.""" + + name: str + config: dict[str, Any] = Field(default_factory=dict) + + +class AdapterProxyConfig(BaseModel): + """Configuration for a localhost adapter proxy in front of an external upstream. + + Used by agents that bring their own inference (e.g. ``claude_code_agent`` + with ``anthropic_base_url``). The agent server starts the proxy alongside + itself; the agent's SDK ``*_BASE_URL`` is rewritten to the proxy's URL so + all outbound model traffic flows through the adapter chain. + """ + + upstream_url: str + adapters: list[InterceptorSpec] = Field(default_factory=list) + host: str = "127.0.0.1" + port: int = 0 + request_timeout: float = 120.0 + unsafe_allow_remote: bool = False diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 83155df73b..e50439b407 100644 --- a/nemo_gym/base_resources_server.py +++ b/nemo_gym/base_resources_server.py @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import abstractmethod +from typing import Any, Optional from fastapi import FastAPI -from pydantic import BaseModel +from pydantic import BaseModel, Field +from nemo_gym.adapters import install_middleware from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest from nemo_gym.openai_utils import ( NeMoGymResponse, @@ -27,7 +29,10 @@ class BaseResourcesServerConfig(BaseRunServerInstanceConfig): - pass + adapters: Optional[list[dict[str, Any]]] = Field( + default=None, + description="Adapter middleware chain: list of {'name': ..., 'config': {...}}. None disables.", + ) class BaseResourcesServer(BaseServer): @@ -66,6 +71,8 @@ def setup_webserver(self) -> FastAPI: app.post("/verify")(self.verify) app.post("/aggregate_metrics")(self.aggregate_metrics) + install_middleware(app, self.config.adapters) + return app async def seed_session(self, body: BaseSeedSessionRequest) -> BaseSeedSessionResponse: diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 42e5f0da65..9e15ee1216 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -12,10 +12,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import atexit from abc import abstractmethod +from typing import Any, Optional from fastapi import Body, FastAPI +from pydantic import Field +from nemo_gym.adapters import AdapterProxyConfig, ProxyHandle, install_middleware, start_adapter_proxy from nemo_gym.base_resources_server import ( AggregateMetrics, AggregateMetricsRequest, @@ -31,7 +35,18 @@ class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig): - pass + adapters: Optional[list[dict[str, Any]]] = Field( + default=None, + description="Adapter middleware chain: list of {'name': ..., 'config': {...}}. None disables.", + ) + adapter_proxy: Optional[AdapterProxyConfig] = Field( + default=None, + description=( + "Optional localhost proxy in front of an external inference upstream. " + "When set, the agent's SDK client should point its *_BASE_URL at " + "``self._proxy_handle.url`` so model traffic flows through the chain." + ), + ) class BaseResponsesAPIAgent(BaseServer): @@ -40,8 +55,21 @@ class BaseResponsesAPIAgent(BaseServer): class SimpleResponsesAPIAgent(BaseResponsesAPIAgent, AggregateMetricsMixin, SimpleServer): config: BaseResponsesAPIAgentConfig + _proxy_handle: Optional[ProxyHandle] = None def setup_webserver(self) -> FastAPI: + if self.config.adapter_proxy is not None: + cfg = self.config.adapter_proxy + self._proxy_handle = start_adapter_proxy( + upstream_url=cfg.upstream_url, + adapters=[spec.model_dump() for spec in cfg.adapters], + host=cfg.host, + port=cfg.port, + request_timeout=cfg.request_timeout, + unsafe_allow_remote=cfg.unsafe_allow_remote, + ) + atexit.register(self._proxy_handle.stop) + app = FastAPI() self.setup_session_middleware(app) @@ -50,6 +78,8 @@ def setup_webserver(self) -> FastAPI: app.post("/run")(self.run) app.post("/aggregate_metrics")(self.aggregate_metrics) + install_middleware(app, self.config.adapters) + return app # TODO: right now there is no validation on the TypedDict NeMoGymResponseCreateParamsNonStreaming diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index e20f14c579..beb6cbb14a 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -13,9 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import abstractmethod +from typing import Any, Optional from fastapi import Body, FastAPI +from pydantic import Field +from nemo_gym.adapters import install_middleware from nemo_gym.openai_utils import ( NeMoGymChatCompletion, NeMoGymChatCompletionCreateParamsNonStreaming, @@ -26,7 +29,10 @@ class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): - pass + adapters: Optional[list[dict[str, Any]]] = Field( + default=None, + description="Adapter middleware chain: list of {'name': ..., 'config': {...}}. None disables.", + ) class BaseResponsesAPIModel(BaseServer): @@ -43,6 +49,8 @@ def setup_webserver(self) -> FastAPI: app.post("/v1/responses")(self.responses) + install_middleware(app, self.config.adapters) + return app @abstractmethod diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index c80bb14b42..295a0cc6d2 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -259,8 +259,10 @@ def _resolve_base_url(self) -> str: async def _run_claude_code(self, instruction: str, system_prompt: Optional[str] = None) -> tuple[str, str]: """Run claude -p --output-format=stream-json and return (stdout, model_name).""" base_url = self._resolve_base_url() - # Keep full model name for local/custom endpoints; strip provider prefix for real Anthropic API. - model = self.config.model if base_url else self.config.model.split("/")[-1] + # Keep full model name for local/custom endpoints (proxy mode counts); + # strip provider prefix only for real Anthropic API. + custom_upstream = bool(base_url) or self._proxy_handle is not None + model = self.config.model if custom_upstream else self.config.model.split("/")[-1] api_key = self.config.anthropic_api_key claude_config_dir = Path.home() / ".claude_code_agent" / uuid4().hex @@ -289,8 +291,17 @@ async def _run_claude_code(self, instruction: str, system_prompt: Optional[str] "IS_SANDBOX": "1", "CLAUDE_CONFIG_DIR": str(claude_config_dir), } - if base_url: - env["ANTHROPIC_BASE_URL"] = base_url + proxy_or_base = self._proxy_handle.url if self._proxy_handle is not None else base_url + if proxy_or_base: + # Custom upstream (proxy in front of an inference endpoint OR a + # direct non-Anthropic endpoint). The claude CLI sends Bearer + # auth via ANTHROPIC_AUTH_TOKEN whenever ANTHROPIC_BASE_URL is + # set; ANTHROPIC_API_KEY (x-api-key) is ignored in that mode. + # Falls back to "local" only when no api_key is configured — + # api.anthropic.com rejects Bearer for sk-ant-... keys, so + # users targeting Anthropic directly must use an OAuth-issued + # token or accept that this auth shape won't work there. + env["ANTHROPIC_BASE_URL"] = proxy_or_base env["ANTHROPIC_AUTH_TOKEN"] = api_key or "local" cmd = [ diff --git a/responses_api_agents/harbor_agent/app.py b/responses_api_agents/harbor_agent/app.py index 72752719c0..061a9701fa 100644 --- a/responses_api_agents/harbor_agent/app.py +++ b/responses_api_agents/harbor_agent/app.py @@ -26,6 +26,7 @@ from fastapi import Body, FastAPI from pydantic import BaseModel, ConfigDict +from nemo_gym.adapters import install_middleware from nemo_gym.base_resources_server import ( BaseRunRequest, BaseVerifyResponse, @@ -200,6 +201,7 @@ def setup_webserver(self) -> FastAPI: app = FastAPI() app.post("/v1/responses")(self.responses) app.post("/run")(self.run) + install_middleware(app, self.config.adapters) return app async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: diff --git a/responses_api_agents/mini_swe_agent/app.py b/responses_api_agents/mini_swe_agent/app.py index 23d64e9389..5e083e1419 100644 --- a/responses_api_agents/mini_swe_agent/app.py +++ b/responses_api_agents/mini_swe_agent/app.py @@ -28,6 +28,7 @@ from minisweagent.run.extra.swegym_runner import _main as run_swegym from pydantic import ConfigDict +from nemo_gym.adapters import install_middleware from nemo_gym.base_resources_server import ( BaseRunRequest, BaseVerifyRequest, @@ -96,6 +97,7 @@ def setup_webserver(self) -> FastAPI: app = FastAPI() app.post("/v1/responses")(self.responses) app.post("/run")(self.run) + install_middleware(app, self.config.adapters) return app async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: diff --git a/tests/unit_tests/adapter_fixtures/consolidate_system_moves_system_to_front_and_merges.json b/tests/unit_tests/adapter_fixtures/consolidate_system_moves_system_to_front_and_merges.json new file mode 100644 index 0000000000..7eb543d87b --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/consolidate_system_moves_system_to_front_and_merges.json @@ -0,0 +1,72 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": {}, + "name": "consolidate_system" + } + ], + "request": { + "body": { + "messages": [ + { + "content": "q1", + "role": "user" + }, + { + "content": "A", + "role": "system" + }, + { + "content": "B", + "role": "system" + } + ], + "model": "m" + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "consolidate_system moves system to front and merges", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/adapter_fixtures/drop_params_drops_temperature_and_top_p.json b/tests/unit_tests/adapter_fixtures/drop_params_drops_temperature_and_top_p.json new file mode 100644 index 0000000000..8b735665ae --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/drop_params_drops_temperature_and_top_p.json @@ -0,0 +1,72 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": { + "params": [ + "temperature", + "top_p" + ] + }, + "name": "drop_params" + } + ], + "request": { + "body": { + "max_tokens": 16, + "messages": [ + { + "content": "hi", + "role": "user" + } + ], + "model": "m", + "temperature": 0.7, + "top_p": 0.95 + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "drop_params drops temperature and top_p", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/adapter_fixtures/modify_tools_strips_a_property___required_entry.json b/tests/unit_tests/adapter_fixtures/modify_tools_strips_a_property___required_entry.json new file mode 100644 index 0000000000..2749bc737e --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/modify_tools_strips_a_property___required_entry.json @@ -0,0 +1,86 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": { + "strip_properties": [ + "debug_flag" + ] + }, + "name": "modify_tools" + } + ], + "request": { + "body": { + "messages": [], + "model": "m", + "tools": [ + { + "function": { + "name": "f", + "parameters": { + "properties": { + "debug_flag": { + "type": "boolean" + }, + "value": { + "type": "string" + } + }, + "required": [ + "debug_flag", + "value" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "modify_tools strips a property + required entry", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/adapter_fixtures/payload_modifier_removes_one_param_and_adds_another.json b/tests/unit_tests/adapter_fixtures/payload_modifier_removes_one_param_and_adds_another.json new file mode 100644 index 0000000000..20c9758a16 --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/payload_modifier_removes_one_param_and_adds_another.json @@ -0,0 +1,67 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": { + "params_to_add": { + "temperature": 0.0 + }, + "params_to_remove": [ + "stream" + ] + }, + "name": "payload_modifier" + } + ], + "request": { + "body": { + "messages": [], + "model": "m", + "stream": true + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "payload_modifier removes one param and adds another", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/adapter_fixtures/system_message_prepends_a_system_message.json b/tests/unit_tests/adapter_fixtures/system_message_prepends_a_system_message.json new file mode 100644 index 0000000000..318e3298b2 --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/system_message_prepends_a_system_message.json @@ -0,0 +1,67 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": { + "strategy": "prepend", + "system_message": "be terse" + }, + "name": "system_message" + } + ], + "request": { + "body": { + "messages": [ + { + "content": "hi", + "role": "user" + } + ], + "model": "m" + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "system_message prepends a system message", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/adapter_fixtures/turn_counter_passes_through_under_the_budget.json b/tests/unit_tests/adapter_fixtures/turn_counter_passes_through_under_the_budget.json new file mode 100644 index 0000000000..1ece10be03 --- /dev/null +++ b/tests/unit_tests/adapter_fixtures/turn_counter_passes_through_under_the_budget.json @@ -0,0 +1,66 @@ +{ + "expected_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + }, + "interceptor_specs": [ + { + "config": { + "max_turns": 5 + }, + "name": "turn_counter" + } + ], + "request": { + "body": { + "messages": [ + { + "content": "hi", + "role": "user" + } + ], + "model": "m" + }, + "headers": { + "content-type": "application/json" + }, + "path": "/v1/chat/completions" + }, + "scenario": "turn_counter passes through under the budget", + "upstream_response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "fixture-canned", + "role": "assistant" + } + } + ], + "id": "chatcmpl-fixture", + "object": "chat.completion" + }, + "headers": { + "content-type": "application/json" + }, + "status_code": 200 + } +} diff --git a/tests/unit_tests/test_adapter_base_class_wiring.py b/tests/unit_tests/test_adapter_base_class_wiring.py new file mode 100644 index 0000000000..4d95047f8e --- /dev/null +++ b/tests/unit_tests/test_adapter_base_class_wiring.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration test for the adapters lift on Agent + Resources bases. + +Confirms that the same ``adapters`` config field + ``install_middleware`` +call shipped on ``BaseResponsesAPIModelConfig`` is now also live on +``BaseResponsesAPIAgentConfig`` and ``BaseResourcesServerConfig``, so any +server inheriting from ``SimpleResponsesAPIAgent`` or ``SimpleResourcesServer`` +picks up the middleware automatically. +""" + +from __future__ import annotations + +import logging +from unittest.mock import MagicMock + +from fastapi import Body +from fastapi.testclient import TestClient + +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseRunRequest, + BaseVerifyRequest, + BaseVerifyResponse, + SimpleResourcesServer, +) +from nemo_gym.base_responses_api_agent import ( + BaseResponsesAPIAgentConfig, + SimpleResponsesAPIAgent, +) +from nemo_gym.openai_utils import ( + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.server_utils import ServerClient + + +class _StubAgent(SimpleResponsesAPIAgent): + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: # type: ignore[override] + raise NotImplementedError + + async def run(self, body: BaseRunRequest = Body()) -> BaseVerifyResponse: # type: ignore[override] + return BaseVerifyResponse( + responses_create_params=body.responses_create_params, + response={"id": "stub", "model": "stub", "usage": {"total_tokens": 7}}, + reward=1.0, + ) + + +class _StubResources(SimpleResourcesServer): + async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse: # type: ignore[override] + return BaseVerifyResponse( + responses_create_params=body.responses_create_params, + response=body.response, + reward=0.5, + ) + + +def _agent_config(adapters: list[dict] | None) -> BaseResponsesAPIAgentConfig: + return BaseResponsesAPIAgentConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="stub_agent", + adapters=adapters, + ) + + +def _resources_config(adapters: list[dict] | None) -> BaseResourcesServerConfig: + return BaseResourcesServerConfig( + host="0.0.0.0", + port=8081, + entrypoint="", + name="stub_resources", + adapters=adapters, + ) + + +def test_agent_base_installs_adapter_chain(caplog) -> None: + agent = _StubAgent( + config=_agent_config(adapters=[{"name": "logging", "config": {}}]), + server_client=MagicMock(spec=ServerClient), + ) + app = agent.setup_webserver() + + # POST to a path the agent server doesn't route — middleware fires + # before routing, so the logging interceptor records the request + # regardless of the 404 that follows. + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + TestClient(app).post("/_probe", json={"x": 1}) + + assert any("request POST /_probe" in rec.message for rec in caplog.records), [ + rec.message for rec in caplog.records + ] + + +def test_agent_base_skips_middleware_when_adapters_none(caplog) -> None: + agent = _StubAgent( + config=_agent_config(adapters=None), + server_client=MagicMock(spec=ServerClient), + ) + app = agent.setup_webserver() + + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + TestClient(app).post("/_probe", json={"x": 1}) + + assert not any("request POST /_probe" in rec.message for rec in caplog.records) + + +def test_resources_base_installs_adapter_chain(caplog) -> None: + resources = _StubResources( + config=_resources_config(adapters=[{"name": "logging", "config": {}}]), + server_client=MagicMock(spec=ServerClient), + ) + app = resources.setup_webserver() + + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + TestClient(app).post("/_probe", json={"x": 1}) + + assert any("request POST /_probe" in rec.message for rec in caplog.records), [ + rec.message for rec in caplog.records + ] + + +class _StubAgentInlineMiddleware(SimpleResponsesAPIAgent): + """Mirrors the harbor_agent / mini_swe_agent override pattern: rebuilds + the FastAPI app from scratch but calls install_middleware inline so the + `adapters` field still applies. + """ + + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: # type: ignore[override] + raise NotImplementedError + + async def run(self, body: BaseRunRequest = Body()) -> BaseVerifyResponse: # type: ignore[override] + raise NotImplementedError + + def setup_webserver(self): + from fastapi import FastAPI + + from nemo_gym.adapters import install_middleware + + app = FastAPI() + app.post("/v1/responses")(self.responses) + app.post("/run")(self.run) + install_middleware(app, self.config.adapters) + return app + + +def test_agent_override_with_inline_install_middleware(caplog) -> None: + """Regression test for the harbor/mini_swe override fix. + + The override doesn't call super().setup_webserver(), so the base class's + install_middleware call doesn't fire. Per the fix, each override adds its + own install_middleware(app, self.config.adapters) call so the YAML + `adapters:` block still takes effect. + """ + agent = _StubAgentInlineMiddleware( + config=_agent_config(adapters=[{"name": "logging", "config": {}}]), + server_client=MagicMock(spec=ServerClient), + ) + app = agent.setup_webserver() + + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + TestClient(app).post("/_probe", json={"x": 1}) + + assert any("request POST /_probe" in rec.message for rec in caplog.records), [ + rec.message for rec in caplog.records + ] + + +def test_resources_base_skips_middleware_when_adapters_none(caplog) -> None: + resources = _StubResources( + config=_resources_config(adapters=None), + server_client=MagicMock(spec=ServerClient), + ) + app = resources.setup_webserver() + + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + TestClient(app).post("/_probe", json={"x": 1}) + + assert not any("request POST /_probe" in rec.message for rec in caplog.records) diff --git a/tests/unit_tests/test_adapter_consolidate_system.py b/tests/unit_tests/test_adapter_consolidate_system.py new file mode 100644 index 0000000000..1f95395cbb --- /dev/null +++ b/tests/unit_tests/test_adapter_consolidate_system.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the consolidate_system interceptor — ported from NEL.""" + +from __future__ import annotations + +from nemo_gym.adapters.interceptors.consolidate_system import Interceptor, _content_to_str +from nemo_gym.adapters.types import AdapterRequest, InterceptorContext + + +def _req(messages: list[dict], session_id: str = "test-session") -> AdapterRequest: + ctx = InterceptorContext() + ctx.extra["session_id"] = session_id + return AdapterRequest( + method="POST", + path="/v1/chat/completions", + headers={"content-type": "application/json"}, + body={"model": "test", "messages": messages}, + ctx=ctx, + ) + + +# --------------------------------------------------------------------------- +# _content_to_str helper +# --------------------------------------------------------------------------- + + +class TestContentToStr: + def test_string_passthrough(self): + assert _content_to_str("hello") == "hello" + + def test_list_of_dicts(self): + content = [{"type": "text", "text": "part1"}, {"type": "text", "text": "part2"}] + assert _content_to_str(content) == "part1\npart2" + + def test_list_of_strings(self): + assert _content_to_str(["a", "b"]) == "a\nb" + + def test_mixed_list(self): + content = [{"type": "text", "text": "dict-part"}, "str-part"] + assert _content_to_str(content) == "dict-part\nstr-part" + + def test_none(self): + assert _content_to_str(None) == "" + + def test_empty_string(self): + assert _content_to_str("") == "" + + def test_empty_list(self): + assert _content_to_str([]) == "" + + def test_dict_without_text_key(self): + assert _content_to_str([{"type": "image_url", "image_url": "x"}]) == "" + + +# --------------------------------------------------------------------------- +# No-op cases: interceptor should return the request unchanged +# --------------------------------------------------------------------------- + + +class TestNoOp: + async def test_empty_messages(self): + ic = Interceptor() + req = _req([]) + result = await ic.intercept_request(req) + assert result.body["messages"] == [] + + async def test_no_messages_key(self): + ic = Interceptor() + ctx = InterceptorContext() + req = AdapterRequest( + method="POST", + path="/v1/chat/completions", + headers={}, + body={"model": "test"}, + ctx=ctx, + ) + result = await ic.intercept_request(req) + assert "messages" not in result.body + + async def test_single_system_at_pos_0(self): + ic = Interceptor() + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"}, + ] + req = _req(messages) + result = await ic.intercept_request(req) + assert result.body["messages"] == messages + assert ic._fix_count == 0 + + async def test_no_system_messages(self): + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + req = _req(messages) + result = await ic.intercept_request(req) + assert result.body["messages"] == messages + assert ic._fix_count == 0 + + +# --------------------------------------------------------------------------- +# Fix cases: interceptor should consolidate system messages +# --------------------------------------------------------------------------- + + +class TestConsolidate: + async def test_system_not_at_pos_0(self): + """Single system message at position > 0 gets moved to front.""" + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "Be helpful."}, + {"role": "assistant", "content": "hello"}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[0] == {"role": "system", "content": "Be helpful."} + assert msgs[1] == {"role": "user", "content": "hi"} + assert msgs[2] == {"role": "assistant", "content": "hello"} + assert ic._fix_count == 1 + + async def test_duplicate_system_messages(self): + """Two system messages get merged into one at position 0.""" + ic = Interceptor() + messages = [ + {"role": "system", "content": "First."}, + {"role": "user", "content": "hi"}, + {"role": "system", "content": "Second."}, + {"role": "assistant", "content": "hello"}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert len(msgs) == 3 + assert msgs[0] == {"role": "system", "content": "First.\n\nSecond."} + assert msgs[1] == {"role": "user", "content": "hi"} + assert msgs[2] == {"role": "assistant", "content": "hello"} + + async def test_empty_system_at_0_real_system_later(self): + """Empty system at pos 0 + real system later triggers fix.""" + ic = Interceptor() + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": "hi"}, + {"role": "system", "content": "Real prompt."}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[0] == {"role": "system", "content": "Real prompt."} + assert msgs[1] == {"role": "user", "content": "hi"} + assert len(msgs) == 2 + + async def test_list_format_content(self): + """System message with OpenAI list-format content is handled.""" + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": [{"type": "text", "text": "Be helpful."}]}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[0] == {"role": "system", "content": "Be helpful."} + assert msgs[1] == {"role": "user", "content": "hi"} + + async def test_three_system_messages(self): + """Three system messages scattered across conversation.""" + ic = Interceptor() + messages = [ + {"role": "system", "content": "A"}, + {"role": "user", "content": "q1"}, + {"role": "system", "content": "B"}, + {"role": "assistant", "content": "a1"}, + {"role": "system", "content": "C"}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[0] == {"role": "system", "content": "A\n\nB\n\nC"} + assert [m["role"] for m in msgs[1:]] == ["user", "assistant"] + + async def test_system_without_content_key(self): + """System message with no content key at all doesn't crash.""" + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system"}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert len(msgs) == 1 + assert msgs[0] == {"role": "user", "content": "hi"} + + +# --------------------------------------------------------------------------- +# Non-system message ordering is preserved +# --------------------------------------------------------------------------- + + +class TestOrderPreservation: + async def test_non_system_order_preserved(self): + ic = Interceptor() + messages = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "system", "content": "sys"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[0]["role"] == "system" + non_system = msgs[1:] + assert [m["content"] for m in non_system] == ["u1", "a1", "u2", "a2"] + + async def test_extra_message_fields_preserved(self): + """Fields like tool_calls, name, etc. are not lost.""" + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "sys"}, + {"role": "assistant", "content": "ok", "tool_calls": [{"id": "1"}]}, + ] + result = await ic.intercept_request(_req(messages)) + msgs = result.body["messages"] + assert msgs[2]["tool_calls"] == [{"id": "1"}] + + +# --------------------------------------------------------------------------- +# Custom separator +# --------------------------------------------------------------------------- + + +class TestCustomSeparator: + async def test_custom_separator(self): + ic = Interceptor(separator=" | ") + messages = [ + {"role": "system", "content": "A"}, + {"role": "user", "content": "hi"}, + {"role": "system", "content": "B"}, + ] + result = await ic.intercept_request(_req(messages)) + assert result.body["messages"][0]["content"] == "A | B" + + +# --------------------------------------------------------------------------- +# Idempotency and fix counter +# --------------------------------------------------------------------------- + + +class TestIdempotency: + async def test_idempotent_after_fix(self): + """Running the interceptor twice on the fixed output is a no-op.""" + ic = Interceptor() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "sys"}, + ] + result1 = await ic.intercept_request(_req(messages)) + assert ic._fix_count == 1 + + result2 = await ic.intercept_request(_req(result1.body["messages"])) + assert ic._fix_count == 1 + assert result2.body["messages"] == result1.body["messages"] + + async def test_fix_counter_increments(self): + ic = Interceptor() + for i in range(3): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": f"sys-{i}"}, + ] + await ic.intercept_request(_req(messages)) + assert ic._fix_count == 3 + + +# --------------------------------------------------------------------------- +# Registry integration +# --------------------------------------------------------------------------- + + +class TestRegistry: + async def test_create_via_registry(self): + from nemo_gym.adapters.registry import InterceptorRegistry + + ic = InterceptorRegistry.create("consolidate_system", {}) + assert isinstance(ic, Interceptor) + + async def test_create_with_separator(self): + from nemo_gym.adapters.registry import InterceptorRegistry + + ic = InterceptorRegistry.create("consolidate_system", {"separator": "---"}) + assert ic._sep == "---" diff --git a/tests/unit_tests/test_adapter_framework.py b/tests/unit_tests/test_adapter_framework.py new file mode 100644 index 0000000000..8ae3b93c15 --- /dev/null +++ b/tests/unit_tests/test_adapter_framework.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Smoke test for the Gym adapter framework. + +Confirms the framework imports, an empty pipeline raises the expected +no-response error, and the registry pre-lists every builtin name. +""" + +from __future__ import annotations + +import pytest + +from nemo_gym.adapters import ( + AdapterPipeline, + AdapterRequest, + InterceptorContext, + InterceptorRegistry, + install_middleware, +) + + +# Names this PR ships in ``nemo_gym.adapters.registry._BUILTIN``. Follow-on +# PRs (observability / caching / request-rewriting) extend the set — +# assertion below is a subset check so adding new builtins doesn't require +# editing this file. +_FRAMEWORK_BUILTINS = { + "endpoint", + "logging", +} + + +def test_framework_imports() -> None: + """All public symbols import cleanly from the package.""" + assert callable(install_middleware) + assert AdapterPipeline is not None + assert InterceptorRegistry is not None + + +@pytest.mark.asyncio +async def test_empty_pipeline_raises_no_response_error() -> None: + """A pipeline with no interceptors cannot produce a response.""" + pipeline = AdapterPipeline([]) + req = AdapterRequest( + method="POST", + path="/v1/chat/completions", + headers={}, + body={"model": "test"}, + ctx=InterceptorContext(), + ) + + with pytest.raises(RuntimeError, match="No interceptor produced a response"): + await pipeline.process(req) + + +def test_registry_pre_lists_framework_builtins() -> None: + """``available()`` includes every framework-level builtin name. + + Subset check (not equality) so this assertion is stable as follow-on + interceptor families extend the builtin set without touching this file. + """ + names = set(InterceptorRegistry.available()) + missing = _FRAMEWORK_BUILTINS - names + assert not missing, f"Framework builtins missing from registry: {sorted(missing)}" diff --git a/tests/unit_tests/test_adapter_interceptors_rewrites.py b/tests/unit_tests/test_adapter_interceptors_rewrites.py new file mode 100644 index 0000000000..4618e02318 --- /dev/null +++ b/tests/unit_tests/test_adapter_interceptors_rewrites.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-interceptor behavior tests — ported from NEL ``tests/test_adapters/test_interceptors.py`` +and ``test_interceptors_extended.py``. + +Mechanical port from ``nemo_evaluator.adapters.*`` to ``nemo_gym.adapters.*``; +``GracefulError`` re-rooted at ``nemo_gym.adapters.types``. +""" + +import pytest + +from nemo_gym.adapters.registry import InterceptorRegistry +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + GracefulError, + InterceptorContext, +) + + +def _req(body=None, **kw): + return AdapterRequest( + method="POST", + path="/v1/chat/completions", + headers={"content-type": "application/json"}, + body=body or {"model": "test", "messages": [{"role": "user", "content": "hi"}]}, + ctx=InterceptorContext(), + ) + + +def _resp(body=None, status_code=200): + return AdapterResponse( + status_code=status_code, + headers={}, + body=body or {}, + ctx=InterceptorContext(), + ) + + + +async def test_drop_params(): + ic = InterceptorRegistry.create("drop_params", {"params": ["temperature", "top_p"]}) + req = _req({"model": "test", "messages": [], "temperature": 0.7, "top_p": 0.9, "max_tokens": 100}) + result = await ic.intercept_request(req) + assert "temperature" not in result.body + assert "top_p" not in result.body + assert result.body["max_tokens"] == 100 + + + +async def test_modify_tools(): + ic = InterceptorRegistry.create("modify_tools", {"strip_properties": ["x"]}) + req = _req( + { + "model": "test", + "messages": [], + "tools": [ + { + "type": "function", + "function": { + "name": "f", + "parameters": { + "type": "object", + "properties": { + "x": {"type": "string"}, + "y": {"type": "integer"}, + }, + "required": ["x", "y"], + }, + }, + } + ], + } + ) + result = await ic.intercept_request(req) + props = result.body["tools"][0]["function"]["parameters"]["properties"] + assert "x" not in props + assert "y" in props + assert "x" not in result.body["tools"][0]["function"]["parameters"]["required"] + + + +async def test_system_message_prepend(): + ic = InterceptorRegistry.create( + "system_message", + { + "system_message": "Be helpful", + "strategy": "prepend", + }, + ) + req = _req({"model": "test", "messages": [{"role": "user", "content": "hi"}]}) + result = await ic.intercept_request(req) + assert result.body["messages"][0] == {"role": "system", "content": "Be helpful"} + assert result.body["messages"][1] == {"role": "user", "content": "hi"} + + + +async def test_system_message_replace(): + ic = InterceptorRegistry.create( + "system_message", + { + "system_message": "New system", + "strategy": "replace", + }, + ) + req = _req( + { + "model": "test", + "messages": [ + {"role": "system", "content": "Old system"}, + {"role": "user", "content": "hi"}, + ], + } + ) + result = await ic.intercept_request(req) + sys_msgs = [m for m in result.body["messages"] if m["role"] == "system"] + assert len(sys_msgs) == 1 + assert sys_msgs[0]["content"] == "New system" + + + +async def test_payload_modifier_remove(): + ic = InterceptorRegistry.create("payload_modifier", {"params_to_remove": ["stream"]}) + req = _req({"model": "test", "messages": [], "stream": True}) + result = await ic.intercept_request(req) + assert "stream" not in result.body + + + +async def test_payload_modifier_add(): + ic = InterceptorRegistry.create("payload_modifier", {"params_to_add": {"temperature": 0.5}}) + req = _req({"model": "test", "messages": []}) + result = await ic.intercept_request(req) + assert result.body["temperature"] == 0.5 + + + +async def test_turn_counter_basic(): + ic = InterceptorRegistry.create("turn_counter", {"max_turns": 5}) + for _ in range(5): + req = _req({"model": "test", "messages": [{"role": "user", "content": "hi"}]}) + await ic.intercept_request(req) + + with pytest.raises(GracefulError, match="Turn budget exhausted"): + await ic.intercept_request(_req({"model": "test", "messages": [{"role": "user", "content": "hi"}]})) + + + +async def test_turn_counter_session_isolation(): + """Repeats of the same problem get independent turn budgets when + the proxy injects distinct session_id values.""" + ic = InterceptorRegistry.create("turn_counter", {"max_turns": 3}) + body = {"model": "test", "messages": [{"role": "user", "content": "same prompt"}]} + + for session_id in ("aaa", "bbb"): + for _ in range(3): + r = _req(body) + r.ctx.extra["session_id"] = session_id + await ic.intercept_request(r) + + r = _req(body) + r.ctx.extra["session_id"] = "aaa" + with pytest.raises(GracefulError, match="Turn budget exhausted"): + await ic.intercept_request(r) + + r = _req(body) + r.ctx.extra["session_id"] = "ccc" + await ic.intercept_request(r) + diff --git a/tests/unit_tests/test_adapter_interceptors_smoke.py b/tests/unit_tests/test_adapter_interceptors_smoke.py new file mode 100644 index 0000000000..eed204262b --- /dev/null +++ b/tests/unit_tests/test_adapter_interceptors_smoke.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Smoke gate: every builtin interceptor resolves and instantiates. + +For every name in ``InterceptorRegistry.available()`` the class is resolved +(import-correctness) and instantiated with ``config={}``. Interceptors that +require non-default kwargs raise ``TypeError`` and are recorded as +``requires_config`` rather than failed. +""" + +from __future__ import annotations + +import logging + +from nemo_gym.adapters import InterceptorRegistry + + +logger = logging.getLogger(__name__) + + +def test_all_builtins_resolve_and_instantiable_or_require_config(caplog) -> None: + caplog.set_level(logging.INFO) + + names = InterceptorRegistry.available() + # Subset check — framework-level minimum, follow-on PRs add more. + assert "endpoint" in names and "logging" in names, f"framework builtins missing from {names}" + assert len(names) >= 2, f"Expected at least 2 builtins, got {len(names)}: {names}" + + instantiates: list[str] = [] + requires_config: list[tuple[str, str]] = [] + + for name in names: + cls = InterceptorRegistry.resolve_class(name) + assert isinstance(cls, type), f"{name!r} did not resolve to a class (got {cls!r})" + + try: + InterceptorRegistry.create(name, config={}) + except TypeError as exc: + requires_config.append((name, str(exc))) + else: + instantiates.append(name) + + summary = ( + f"{len(instantiates)}/{len(names)} instantiate with empty config; " + f"{len(requires_config)}/{len(names)} require config" + ) + print(summary) + print(" empty-config OK:", sorted(instantiates)) + print(" require config:") + for n, msg in sorted(requires_config): + print(f" {n}: {msg}") + + # Sanity: at least one of each side. If every interceptor took empty + # config the contract would be too loose; if none did, the registry + # is probably broken. + assert instantiates, "no interceptor accepted empty config — registry likely broken" + assert requires_config, "every interceptor accepted empty config — likely loss of required-arg validation" diff --git a/tests/unit_tests/test_adapter_middleware_behaviors.py b/tests/unit_tests/test_adapter_middleware_behaviors.py new file mode 100644 index 0000000000..9c85f22a83 --- /dev/null +++ b/tests/unit_tests/test_adapter_middleware_behaviors.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Middleware behavioral tests — replace NEL ``tests/test_adapters/test_proxy.py``. + +NEL's ``test_proxy.py`` exercised a standalone uvicorn proxy spun up by +``start_adapter_proxy(...)``. Gym does not have a standalone proxy — the +adapter pipeline is installed as FastAPI middleware on an existing model +server's app via ``install_middleware``. The architectural shift means the +proxy lifecycle / port-allocation tests do not port. + +What does port — and is exercised here against the middleware via a +FastAPI ``TestClient`` — are the **behavioral invariants** that NEL's +proxy guaranteed and that Gym's Phase-1.5 middleware also implements: + +* ``/s//...`` session-id prefix is parsed off the path and exposed + to interceptors via ``ctx.extra["session_id"]``. +* Hop-by-hop response headers (``transfer-encoding``, etc.) are stripped. +* ``GracefulError`` raised inside an interceptor returns HTTP 429 with a + ``session_budget_exhausted`` error code. +* Invalid JSON bodies return HTTP 400. +* The pipeline runs only on POST; other methods pass through unchanged. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any + +import pytest +from fastapi import Body, FastAPI, Request +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient +from starlette.middleware.sessions import SessionMiddleware + +from nemo_gym.adapters import install_middleware +from nemo_gym.adapters.registry import InterceptorRegistry +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + GracefulError, + RequestInterceptor, + RequestToResponseInterceptor, +) + + +# --------------------------------------------------------------------------- +# Helper interceptors registered for use in these tests +# --------------------------------------------------------------------------- + + +class _HeaderCapture(RequestInterceptor): + """Captures the last request the pipeline saw so tests can introspect ctx.""" + + last_req: AdapterRequest | None = None + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + type(self).last_req = req + return req + + +class _CapturingEcho(RequestToResponseInterceptor): + """Captures the request *and* short-circuits with a canned response. + + Using a short-circuiting interceptor — rather than letting the request + fall through to the FastAPI route — lets us assert on what the pipeline + saw without depending on the route handler matching the session-prefixed + URL. (FastAPI's router doesn't know about session prefixes; that's the + middleware's job, and NEL's standalone proxy used to rewrite the URL + before forwarding. Gym's middleware annotates ctx instead, and trusts + the chain to handle routing semantics.) + """ + + last_req: AdapterRequest | None = None + + async def intercept_request(self, req: AdapterRequest) -> AdapterResponse: + type(self).last_req = req + return AdapterResponse( + status_code=200, + headers={"content-type": "application/json"}, + body={"id": "chatcmpl-canned", "echo_path": req.path}, + ctx=req.ctx, + ) + + +class _GracefulRaiser(RequestInterceptor): + """Always raises ``GracefulError`` to simulate session-budget exhaustion.""" + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + raise GracefulError("turns up") + + +def _install_helper_interceptors() -> None: + """Make the helpers above resolvable via the registry.""" + mod = types.ModuleType("nemo_gym.adapters.interceptors._test_header_capture") + mod.Interceptor = _HeaderCapture + sys.modules[mod.__name__] = mod + InterceptorRegistry.register("_test_header_capture", mod.__name__) + + mod2 = types.ModuleType("nemo_gym.adapters.interceptors._test_graceful_raiser") + mod2.Interceptor = _GracefulRaiser + sys.modules[mod2.__name__] = mod2 + InterceptorRegistry.register("_test_graceful_raiser", mod2.__name__) + + mod3 = types.ModuleType("nemo_gym.adapters.interceptors._test_capturing_echo") + mod3.Interceptor = _CapturingEcho + sys.modules[mod3.__name__] = mod3 + InterceptorRegistry.register("_test_capturing_echo", mod3.__name__) + + +# --------------------------------------------------------------------------- +# App / client factory +# --------------------------------------------------------------------------- + + +def _route_handler_default(body: dict) -> dict: + """Default canned model-server response: echoes path and body.""" + return { + "id": "chatcmpl-route", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "from route"}, + } + ], + "echo": body, + } + + +def _build_test_app( + interceptor_specs: list[dict[str, Any]] | None, + *, + extra_response_headers: dict[str, str] | None = None, +) -> FastAPI: + """Build a FastAPI app whose ``/v1/chat/completions`` POST route returns + a canned response, with the adapter middleware optionally installed on top. + + ``extra_response_headers`` is appended to the model-server's reply so + tests can verify that hop-by-hop headers added by the (mock) upstream + are stripped by the middleware before reaching the client. + """ + app = FastAPI() + + @app.post("/v1/chat/completions") + async def _chat(body: dict): + return JSONResponse( + content=_route_handler_default(body), + headers=extra_response_headers or {}, + ) + + install_middleware(app, interceptor_specs) + return app + + +@pytest.fixture(autouse=True) +def _setup_registry() -> None: + _install_helper_interceptors() + _HeaderCapture.last_req = None + _CapturingEcho.last_req = None + + +# --------------------------------------------------------------------------- +# Session-id parsing +# --------------------------------------------------------------------------- + + +def test_session_id_parsed_into_ctx_extra() -> None: + """``POST /s//v1/...`` strips the prefix and exposes the id on ctx. + + Uses a short-circuiting interceptor (``_test_capturing_echo``) so the + request never falls through to the FastAPI router — which doesn't know + about session-prefixed URLs. The middleware-side behavior being verified + is purely that the pipeline sees ``ctx.extra["session_id"]`` set to the + expected hex string and the stripped path. + """ + app = _build_test_app([{"name": "_test_capturing_echo"}]) + with TestClient(app) as client: + resp = client.post( + "/s/deadbeef1234/v1/chat/completions", + json={"model": "test", "messages": []}, + ) + assert resp.status_code == 200 + assert _CapturingEcho.last_req is not None + assert _CapturingEcho.last_req.ctx.extra.get("session_id") == "deadbeef1234" + # The path the interceptor saw is the clean one with the prefix removed. + assert _CapturingEcho.last_req.path == "/v1/chat/completions" + + +def test_no_session_id_for_plain_path() -> None: + """A plain ``/v1/...`` POST exposes no session_id on the interceptor ctx.""" + app = _build_test_app([{"name": "_test_capturing_echo"}]) + with TestClient(app) as client: + resp = client.post( + "/v1/chat/completions", + json={"model": "test", "messages": []}, + ) + assert resp.status_code == 200 + assert _CapturingEcho.last_req is not None + assert "session_id" not in _CapturingEcho.last_req.ctx.extra + assert _CapturingEcho.last_req.path == "/v1/chat/completions" + + +# --------------------------------------------------------------------------- +# Hop-by-hop header filtering +# --------------------------------------------------------------------------- + + +def test_hop_by_hop_headers_stripped_from_response() -> None: + """Hop-by-hop headers added by the underlying handler are stripped + before being returned to the client.""" + app = _build_test_app( + [{"name": "logging"}], + extra_response_headers={ + "x-upstream-marker": "kept", + "transfer-encoding": "chunked", # must be stripped + "connection": "keep-alive", # must be stripped + }, + ) + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json={"model": "test", "messages": []}) + assert resp.status_code == 200 + keys = {k.lower() for k in resp.headers} + assert "transfer-encoding" not in keys + assert "connection" not in keys + # Non hop-by-hop custom headers must be preserved. + assert resp.headers.get("x-upstream-marker") == "kept" + + +# --------------------------------------------------------------------------- +# GracefulError → 429 +# --------------------------------------------------------------------------- + + +def test_graceful_error_returns_429() -> None: + """When an interceptor raises ``GracefulError``, the middleware emits + HTTP 429 with a ``session_budget_exhausted`` error code.""" + app = _build_test_app([{"name": "_test_graceful_raiser"}]) + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json={"model": "test", "messages": []}) + assert resp.status_code == 429 + body = resp.json() + assert body["error"]["code"] == "session_budget_exhausted" + assert "turns up" in body["error"]["message"] + + +# --------------------------------------------------------------------------- +# Invalid JSON → 400 +# --------------------------------------------------------------------------- + + +def test_invalid_json_body_returns_400() -> None: + """A POST whose body is not valid JSON returns 400 before any + interceptor runs (the pipeline must never see malformed input).""" + app = _build_test_app([{"name": "_test_header_capture"}]) + with TestClient(app) as client: + resp = client.post( + "/v1/chat/completions", + content=b"not-json", + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["type"] == "invalid_request_error" + # The pipeline must not have run. + assert _HeaderCapture.last_req is None + + +# --------------------------------------------------------------------------- +# Non-POST methods pass through (the pipeline only intercepts POST) +# --------------------------------------------------------------------------- + + +def test_get_request_passes_through_middleware() -> None: + """GET requests bypass the pipeline entirely — the model server's own + routes (e.g. /health) keep working unmodified. + + The default app doesn't define a GET route, so we expect 405 (Method + Not Allowed) from the route handler — proving the middleware did not + intercept and short-circuit with its own response.""" + app = _build_test_app([{"name": "_test_header_capture"}]) + with TestClient(app) as client: + resp = client.get("/v1/chat/completions") + assert resp.status_code == 405 # Method Not Allowed from FastAPI + # Pipeline was not invoked for GET. + assert _HeaderCapture.last_req is None + + +# --------------------------------------------------------------------------- +# Multi-valued response headers (Set-Cookie / SessionMiddleware) — the +# ecosystem-review blocker. ``SimpleResponsesAPIModel`` always has Starlette +# ``SessionMiddleware`` attached; before the fix, ``_starlette_response_to_adapter`` +# rebuilt headers as a Python dict and silently collapsed duplicate +# ``Set-Cookie`` headers. This test mounts the adapter middleware on top of +# a FastAPI app that has ``SessionMiddleware`` installed plus an explicit +# route that sets an additional ``Set-Cookie`` — and asserts both cookies +# survive the round trip. +# --------------------------------------------------------------------------- + + +def test_set_cookie_headers_preserved_through_middleware() -> None: + """Multiple ``Set-Cookie`` headers survive the adapter middleware. + + Pre-fix, ``_starlette_response_to_adapter`` collapsed response headers + into a Python ``dict``, so a response carrying two ``Set-Cookie`` + headers (one from ``SessionMiddleware``, one from the route handler) + arrived at the client with only the last value. Post-fix, headers + flow as ``list[tuple[bytes, bytes]]`` end-to-end and duplicates are + preserved. + """ + app = FastAPI() + # Match the order ``SimpleResponsesAPIModel.setup_webserver`` uses: + # session middleware first, then the route handlers, then (after this + # builder returns) the adapter middleware on top. + app.add_middleware(SessionMiddleware, secret_key="test-secret-key") # pragma: allowlist secret + + @app.post("/v1/chat/completions") + async def _chat(request: Request, body: dict = Body(...)): + # Mutate the session so SessionMiddleware actually emits its + # ``Set-Cookie`` on the response path (it skips otherwise). + request.session["test_key"] = "test_value" + resp = JSONResponse(content={"ok": True, "echo": body}) + # Set a second cookie explicitly on top of the one SessionMiddleware + # will inject. Two cookies on the same response is the exact case + # the dict-collapse bug obliterated. + resp.set_cookie("explicit_cookie", "explicit_value") + return resp + + install_middleware(app, [{"name": "logging"}]) + + with TestClient(app) as client: + resp = client.post( + "/v1/chat/completions", + json={"model": "test", "messages": []}, + ) + assert resp.status_code == 200, f"body={resp.text!r}" + # httpx exposes multi-valued headers via ``.get_list``; iterating the + # underlying raw header list is the most portable cross-version check. + raw_set_cookies = [v for k, v in resp.headers.raw if k.lower() == b"set-cookie"] + assert len(raw_set_cookies) >= 2, ( + f"expected ≥2 Set-Cookie headers (SessionMiddleware + explicit), got: {raw_set_cookies!r}" + ) + assert any(b"explicit_cookie=explicit_value" in v for v in raw_set_cookies), ( + f"explicit cookie missing from response: {raw_set_cookies!r}" + ) + # SessionMiddleware's cookie is named "session" by default. + assert any(v.lower().startswith(b"session=") for v in raw_set_cookies), ( + f"SessionMiddleware cookie missing from response: {raw_set_cookies!r}" + ) diff --git a/tests/unit_tests/test_adapter_middleware_integration.py b/tests/unit_tests/test_adapter_middleware_integration.py new file mode 100644 index 0000000000..c0b2864f0b --- /dev/null +++ b/tests/unit_tests/test_adapter_middleware_integration.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Middleware integration tests (Phase 1.5 — wrap, don't replace). + +The middleware now layers on top of an existing FastAPI route handler: +REQUEST interceptors mutate the body, REQUEST_TO_RESPONSE interceptors may +short-circuit, otherwise ``call_next`` invokes the model server's own +handler, and RESPONSE interceptors observe the result. These tests verify: + +* ``log_tokens`` (RESPONSE stage) observes the route handler's output when + the chain has no short-circuiting interceptor. +* REQUEST-stage interceptors mutate the body the route handler sees. +* A short-circuiting REQUEST_TO_RESPONSE interceptor prevents the route + handler from ever running. +* ``adapters=None`` / ``adapters=[]`` keep the middleware uninstalled. +""" + +from __future__ import annotations + +import logging +import sys +import types + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from nemo_gym.adapters import install_middleware +from nemo_gym.adapters.registry import InterceptorRegistry +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + RequestToResponseInterceptor, +) + + +# --------------------------------------------------------------------------- +# Test fixtures: a minimal canned-response interceptor registered as a +# RequestToResponseInterceptor. Used by the short-circuit assertion test. +# --------------------------------------------------------------------------- + + +_CANNED_BODY = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "test-model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello from middleware"}, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8}, +} + + +class _CannedInterceptor(RequestToResponseInterceptor): + """Returns ``_CANNED_BODY`` without touching the network.""" + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest | AdapterResponse: + return AdapterResponse( + status_code=200, + headers={"content-type": "application/json"}, + body=dict(_CANNED_BODY), + latency_ms=1.0, + ctx=req.ctx, + ) + + +def _register_canned_interceptor() -> None: + """Make the canned interceptor resolvable by ``InterceptorRegistry``.""" + module_name = "nemo_gym.adapters.interceptors._test_canned" + if module_name not in sys.modules: + mod = types.ModuleType(module_name) + mod.Interceptor = _CannedInterceptor + sys.modules[module_name] = mod + InterceptorRegistry.register("_test_canned", module_name) + + +# A handler-visibility flag: the test route writes the request body it +# actually received into this list so tests can assert what the model +# server saw (e.g. mutated by REQUEST interceptors). +_route_seen: list[dict] = [] + + +def _build_app() -> FastAPI: + """Minimal FastAPI app with a single ``/v1/chat/completions`` POST route. + + The route stashes the parsed body into ``_route_seen`` and returns a + sentinel body so the test can distinguish whether the request reached + the underlying handler or was short-circuited by middleware. + """ + app = FastAPI() + + @app.post("/v1/chat/completions") + async def _chat_completions(body: dict) -> dict: + _route_seen.append(body) + return { + "id": "chatcmpl-route", + "object": "chat.completion", + "created": 0, + "model": body.get("model", "?"), + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "from route"}, + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, + } + + return app + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _register() -> None: + _register_canned_interceptor() + _route_seen.clear() + + +_SAMPLE_BODY = { + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], +} + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_install_middleware_runs_pipeline_and_logging_emits_log(caplog) -> None: + """A POST flows through the pipeline → route handler → RESPONSE stage. + + The model server's own route handler produces the response; + ``logging`` (a request+response interceptor) emits records on both + sides to prove the pipeline ran. + """ + app = _build_app() + install_middleware( + app, + [ + {"name": "logging", "config": {}}, + ], + ) + + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json=_SAMPLE_BODY) + + assert resp.status_code == 200, resp.text + payload = resp.json() + + # The route handler ran (not a canned short-circuit). + assert payload["choices"][0]["message"]["content"] == "from route" + assert _route_seen == [_SAMPLE_BODY] + + # logging fires on both the request and response halves of the chain. + log_records = [rec for rec in caplog.records if rec.name == "nemo_gym.adapters.interceptors.request_logging"] + assert any("request POST /v1/chat/completions" in r.getMessage() for r in log_records) + assert any("response status=200" in r.getMessage() for r in log_records) + + +def test_short_circuit_interceptor_prevents_route_handler() -> None: + """A short-circuiting REQUEST_TO_RESPONSE interceptor skips ``call_next``. + + The route handler must never run, ``_route_seen`` stays empty, and the + client receives the canned body. + """ + app = _build_app() + install_middleware( + app, + [ + {"name": "logging", "config": {}}, + {"name": "_test_canned", "config": {}}, + ], + ) + + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json=_SAMPLE_BODY) + + assert resp.status_code == 200, resp.text + assert resp.json()["choices"][0]["message"]["content"] == "hello from middleware" + # Crucial: the underlying route handler was bypassed. + assert _route_seen == [], "Route handler ran even though chain short-circuited" + + +def test_install_middleware_defaults_off_when_specs_is_none() -> None: + """``adapters=None`` is a no-op: the route handler runs unchanged.""" + app = _build_app() + install_middleware(app, None) + + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json=_SAMPLE_BODY) + + assert resp.status_code == 200 + # Route handler reply (not the canned interceptor one) — proves no + # middleware was inserted. + assert resp.json()["choices"][0]["message"]["content"] == "from route" + + +def test_install_middleware_defaults_off_when_specs_is_empty() -> None: + """``adapters=[]`` is also a no-op (mirrors the None case).""" + app = _build_app() + install_middleware(app, []) + + with TestClient(app) as client: + resp = client.post("/v1/chat/completions", json=_SAMPLE_BODY) + + assert resp.status_code == 200 + assert resp.json()["choices"][0]["message"]["content"] == "from route" diff --git a/tests/unit_tests/test_adapter_parity_replay_rewrites.py b/tests/unit_tests/test_adapter_parity_replay_rewrites.py new file mode 100644 index 0000000000..a68a04dde3 --- /dev/null +++ b/tests/unit_tests/test_adapter_parity_replay_rewrites.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Parity replay test for the adapter middleware. + +Loads every JSON fixture under ``adapter_fixtures/`` and asserts that the +recorded ``request`` → ``expected_response`` pair still holds when the +middleware is installed with the recorded ``interceptor_specs`` and the +(mocked) upstream returns the recorded ``upstream_response``. ``caching`` +is covered by a dedicated round-trip test below since its parity behavior +is "second hit ≡ first response". +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient + +from nemo_gym.adapters import install_middleware + + +FIXTURE_DIR = Path(__file__).parent / "adapter_fixtures" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _load_fixtures() -> list[tuple[str, dict[str, Any]]]: + """Return ``[(fixture_name, fixture_data), ...]`` sorted by name.""" + if not FIXTURE_DIR.exists(): + return [] + fixtures: list[tuple[str, dict[str, Any]]] = [] + for path in sorted(FIXTURE_DIR.glob("*.json")): + with path.open() as f: + fixtures.append((path.stem, json.load(f))) + return fixtures + + +_FIXTURES = _load_fixtures() + + +def _build_replay_app(interceptor_specs: list[dict[str, Any]], upstream: dict[str, Any]) -> FastAPI: + """FastAPI app whose chat-completions route returns ``upstream`` verbatim, + with the adapter middleware installed on top.""" + app = FastAPI() + + @app.post("/v1/chat/completions") + async def _chat(body: dict): + return JSONResponse( + content=upstream["body"], + status_code=upstream["status_code"], + headers=upstream.get("headers") or {}, + ) + + install_middleware(app, interceptor_specs) + return app + + +_VOLATILE_HEADERS = {"date", "server", "content-length"} + + +def _normalise_headers(headers: dict[str, str]) -> dict[str, str]: + """Drop headers whose values are non-deterministic. + + Matches the normalisation in ``generate_adapter_fixtures.py``. Without this, + replays would fail because ``content-length`` is recomputed per-response + by Starlette and ``date`` / ``server`` vary across runs. + """ + return {k: v for k, v in headers.items() if k.lower() not in _VOLATILE_HEADERS} + + +# --------------------------------------------------------------------------- +# Parametrised replay +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _FIXTURES, reason="no fixtures recorded yet (run generate_adapter_fixtures.py)") +@pytest.mark.parametrize("name,fixture", _FIXTURES, ids=[name for name, _ in _FIXTURES]) +def test_parity_replay(name: str, fixture: dict[str, Any]) -> None: + """Each fixture must replay byte-equal through the middleware chain.""" + request = fixture["request"] + upstream = fixture["upstream_response"] + expected = fixture["expected_response"] + + app = _build_replay_app(fixture["interceptor_specs"], upstream) + with TestClient(app) as client: + resp = client.post(request["path"], json=request["body"], headers=request["headers"]) + + assert resp.status_code == expected["status_code"], ( + f"fixture {name!r}: status mismatch (got {resp.status_code}, expected {expected['status_code']})" + ) + + actual_body = resp.json() + assert actual_body == expected["body"], ( + f"fixture {name!r}: body mismatch\n got: {actual_body!r}\n expected: {expected['body']!r}" + ) + + actual_headers = _normalise_headers(dict(resp.headers)) + expected_headers = _normalise_headers(expected["headers"]) + assert actual_headers == expected_headers, ( + f"fixture {name!r}: headers mismatch\n got: {actual_headers!r}\n expected: {expected_headers!r}" + ) diff --git a/tests/unit_tests/test_adapter_pipeline.py b/tests/unit_tests/test_adapter_pipeline.py new file mode 100644 index 0000000000..bd6a2b87c7 --- /dev/null +++ b/tests/unit_tests/test_adapter_pipeline.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""AdapterPipeline behavior tests — ported from NEL ``tests/test_adapters/test_pipeline.py``. + +Mechanical port from ``nemo_evaluator.adapters.*`` to ``nemo_gym.adapters.*`` +plus one additional Phase-1.5 test for the ``upstream_call`` parameter on +``AdapterPipeline.process()``. +""" + +from __future__ import annotations + +import pytest + +from nemo_gym.adapters.pipeline import AdapterPipeline +from nemo_gym.adapters.types import ( + AdapterRequest, + AdapterResponse, + InterceptorContext, + RequestInterceptor, + RequestToResponseInterceptor, + ResponseInterceptor, +) + + +def _req(**body_overrides): + body = {"model": "test", "messages": [{"role": "user", "content": "hi"}]} + body.update(body_overrides) + return AdapterRequest( + method="POST", + path="/v1/chat/completions", + headers={"content-type": "application/json"}, + body=body, + ctx=InterceptorContext(), + ) + + +class AddHeaderInterceptor(RequestInterceptor): + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + req.headers["x-test"] = "1" + return req + + +class FixedEndpoint(RequestToResponseInterceptor): + def __init__(self): + self.last_request = None + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest | AdapterResponse: + self.last_request = req + return AdapterResponse( + status_code=200, + headers={}, + body={"result": "ok"}, + ctx=req.ctx, + ) + + +class AppendBodyInterceptor(ResponseInterceptor): + async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse: + if isinstance(resp.body, dict): + resp.body["appended"] = True + return resp + + +class TrackingResponseInterceptor(ResponseInterceptor): + def __init__(self): + self.called = False + + async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse: + self.called = True + return resp + + +class CacheHitInterceptor(RequestToResponseInterceptor): + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest | AdapterResponse: + return AdapterResponse( + status_code=200, + headers={}, + body={"cached": True}, + ctx=req.ctx, + ) + + +class FailingRequestInterceptor(RequestInterceptor): + def __init__(self, *, best_effort=False): + self.best_effort = best_effort + + async def intercept_request(self, req: AdapterRequest) -> AdapterRequest: + raise RuntimeError("boom") + + +def test_stage_order_validation(): + with pytest.raises(ValueError, match="Invalid interceptor order"): + AdapterPipeline([AppendBodyInterceptor(), AddHeaderInterceptor()]) + + +async def test_request_then_endpoint_then_response(): + endpoint = FixedEndpoint() + pipeline = AdapterPipeline( + [ + AddHeaderInterceptor(), + endpoint, + ] + ) + resp = await pipeline.process(_req()) + assert resp.status_code == 200 + assert resp.body["result"] == "ok" + assert endpoint.last_request.headers["x-test"] == "1" + + +async def test_short_circuit_skips_endpoint(): + """When a RequestToResponseInterceptor short-circuits, later request-side + interceptors (like the endpoint) are skipped, but response interceptors + still run so they can inspect the short-circuited response.""" + endpoint = FixedEndpoint() + tracker = TrackingResponseInterceptor() + pipeline = AdapterPipeline( + [ + CacheHitInterceptor(), + endpoint, + tracker, + ] + ) + resp = await pipeline.process(_req()) + assert resp.body["cached"] is True + assert endpoint.last_request is None # endpoint was skipped + assert tracker.called is True # response interceptors still fire + + +async def test_best_effort_continues_on_error(): + pipeline = AdapterPipeline( + [ + FailingRequestInterceptor(best_effort=True), + FixedEndpoint(), + ] + ) + resp = await pipeline.process(_req()) + assert resp.status_code == 200 + + +async def test_non_best_effort_raises(): + pipeline = AdapterPipeline( + [ + FailingRequestInterceptor(best_effort=False), + FixedEndpoint(), + ] + ) + with pytest.raises(RuntimeError, match="boom"): + await pipeline.process(_req()) + + +async def test_response_interceptors_run_after_endpoint(): + """Response interceptors placed after the endpoint in the chain must run.""" + appender = AppendBodyInterceptor() + pipeline = AdapterPipeline( + [ + AddHeaderInterceptor(), + FixedEndpoint(), + appender, + ] + ) + resp = await pipeline.process(_req()) + assert resp.body.get("appended") is True + + +async def test_multiple_response_interceptors_run_in_reverse(): + """Multiple response interceptors run in reverse chain order.""" + order: list[str] = [] + + class First(ResponseInterceptor): + async def intercept_response(self, resp): + order.append("first") + return resp + + class Second(ResponseInterceptor): + async def intercept_response(self, resp): + order.append("second") + return resp + + pipeline = AdapterPipeline( + [ + FixedEndpoint(), + First(), + Second(), + ] + ) + await pipeline.process(_req()) + assert order == ["second", "first"] + + +# --------------------------------------------------------------------------- +# Phase 1.5 addition: upstream_call invoked when no interceptor short-circuits +# --------------------------------------------------------------------------- + + +async def test_upstream_call_invoked_when_chain_does_not_short_circuit(): + """When ``upstream_call`` is provided and the chain has no + ``RequestToResponseInterceptor`` that short-circuits, the upstream is + called with the (post-REQUEST-stage) request and its response flows + through ``ResponseInterceptor`` instances in reverse order. + + This is the Phase 1.5 wrap-not-replace contract: an ``endpoint`` + interceptor is no longer required in every chain — the middleware can + supply the upstream via ``upstream_call`` instead. + """ + upstream_seen: list[AdapterRequest] = [] + appender = AppendBodyInterceptor() + + async def _upstream(req: AdapterRequest) -> AdapterResponse: + upstream_seen.append(req) + return AdapterResponse( + status_code=200, + headers={"content-type": "application/json"}, + body={"result": "from-upstream"}, + ctx=req.ctx, + ) + + pipeline = AdapterPipeline( + [ + AddHeaderInterceptor(), # REQUEST stage + appender, # RESPONSE stage + ] + ) + resp = await pipeline.process(_req(), upstream_call=_upstream) + + # Upstream was invoked exactly once with the (mutated) request. + assert len(upstream_seen) == 1 + assert upstream_seen[0].headers["x-test"] == "1" + + # Response flowed back through the RESPONSE stage. + assert resp.status_code == 200 + assert resp.body["result"] == "from-upstream" + assert resp.body["appended"] is True + + +async def test_upstream_call_skipped_when_chain_short_circuits(): + """A short-circuiting ``RequestToResponseInterceptor`` skips + ``upstream_call`` entirely — the upstream must never be invoked.""" + upstream_seen: list[AdapterRequest] = [] + + async def _upstream(req: AdapterRequest) -> AdapterResponse: + upstream_seen.append(req) + return AdapterResponse(status_code=200, headers={}, body={"result": "from-upstream"}, ctx=req.ctx) + + pipeline = AdapterPipeline( + [ + CacheHitInterceptor(), # short-circuits + ] + ) + resp = await pipeline.process(_req(), upstream_call=_upstream) + assert upstream_seen == [] + assert resp.body == {"cached": True} + + +# =========================================================================== +# Library-import path — explicit positive coverage that the NEL-style +# ``AdapterPipeline.process(req)`` call shape (no ``upstream_call``, no +# FastAPI middleware host) still works end-to-end when a +# ``RequestToResponseInterceptor`` short-circuits the chain. This is the +# shape an external library user would import; we pin it here so the +# Phase-1.5 ``upstream_call`` parameter never silently becomes mandatory. +# =========================================================================== + + +async def test_library_import_no_upstream_call_no_middleware_short_circuits(): + """A NEL-style library user instantiates ``AdapterPipeline`` directly and + calls ``process(req)`` with no ``upstream_call`` and no FastAPI host. + A short-circuiting interceptor must produce the response.""" + endpoint = FixedEndpoint() # short-circuits with status 200, body={"result": "ok"} + tracker = TrackingResponseInterceptor() + pipeline = AdapterPipeline([endpoint, tracker]) + + # Call without ``upstream_call`` — the library-import surface. + resp = await pipeline.process(_req()) + + assert isinstance(resp, AdapterResponse) + assert resp.status_code == 200 + assert resp.body == {"result": "ok"} + assert endpoint.last_request is not None # endpoint was the source of the response + assert tracker.called is True # RESPONSE-stage interceptor still ran + + +async def test_library_import_request_stage_mutation_reaches_short_circuit(): + """Request-stage interceptors must mutate the request before the + short-circuiting endpoint sees it, even when no ``upstream_call`` is + provided. Pins the library-import contract end-to-end.""" + endpoint = FixedEndpoint() + pipeline = AdapterPipeline([AddHeaderInterceptor(), endpoint]) + + resp = await pipeline.process(_req()) + + # The mutation from AddHeaderInterceptor must have landed on the request + # that FixedEndpoint observed. + assert endpoint.last_request is not None + assert endpoint.last_request.headers["x-test"] == "1" + assert resp.status_code == 200 + assert resp.body == {"result": "ok"} diff --git a/tests/unit_tests/test_adapter_proxy.py b/tests/unit_tests/test_adapter_proxy.py new file mode 100644 index 0000000000..edfe80decb --- /dev/null +++ b/tests/unit_tests/test_adapter_proxy.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration tests for ``start_adapter_proxy``. + +Drives a real uvicorn proxy thread against a stub upstream (a FastAPI app +hosted by Starlette's ``TestClient``-equivalent in-thread) and asserts: + + - adapted POST routes (``/v1/chat/completions``, ``/v1/messages``) run + the pipeline (logging interceptor fires) + - non-adapted paths (``/v1/models`` GET, ``/health``) pass through + without running the pipeline + - localhost-bind enforcement (``host="0.0.0.0"`` rejected) + - user-supplied ``endpoint`` interceptor rejected + - multi-Set-Cookie response headers survive the proxy +""" + +from __future__ import annotations + +import json +import logging +import threading +import time +import urllib.error +import urllib.request +from typing import Any + +import pytest +import uvicorn +from fastapi import FastAPI, Request +from starlette.responses import JSONResponse, Response + +from nemo_gym.adapters import start_adapter_proxy + + +# --------------------------------------------------------------------------- +# In-thread stub upstream — records requests, lets us inject responses +# --------------------------------------------------------------------------- + + +class _StubUpstream: + def __init__(self) -> None: + self.received: list[dict[str, Any]] = [] + self._lock = threading.Lock() + self.app = FastAPI() + self._build_routes() + self.port: int | None = None + self._server: uvicorn.Server | None = None + self._thread: threading.Thread | None = None + + def _build_routes(self) -> None: + @self.app.api_route( + "/{path:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], + ) + async def echo(path: str, request: Request) -> Response: + try: + body = json.loads(await request.body() or b"{}") + except Exception: + body = None + with self._lock: + self.received.append({"method": request.method, "path": "/" + path, "body": body}) + # Default: OpenAI-compat success + if request.method == "POST": + resp_body = { + "id": "stub-1", + "object": "chat.completion", + "model": body.get("model", "stub") if isinstance(body, dict) else "stub", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "stub-ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + # Inject two Set-Cookie headers to exercise multi-cookie path + out = JSONResponse(resp_body, status_code=200) + out.raw_headers.append((b"set-cookie", b"a=1; Path=/")) + out.raw_headers.append((b"set-cookie", b"b=2; HttpOnly")) + return out + return JSONResponse({"path": "/" + path, "method": request.method}) + + def start(self) -> None: + import socket + + s = socket.socket() + s.bind(("127.0.0.1", 0)) + self.port = s.getsockname()[1] + s.close() + cfg = uvicorn.Config(self.app, host="127.0.0.1", port=self.port, log_level="warning", access_log=False) + self._server = uvicorn.Server(cfg) + self._thread = threading.Thread(target=self._server.run, daemon=True) + self._thread.start() + # Wait until ready + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{self.port}/healthcheck", timeout=1) as r: + if r.status == 200: + return + except Exception: + time.sleep(0.05) + raise RuntimeError("stub upstream did not become healthy") + + def stop(self) -> None: + if self._server: + self._server.should_exit = True + if self._thread: + self._thread.join(timeout=3) + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +@pytest.fixture +def upstream(): + s = _StubUpstream() + s.start() + try: + yield s + finally: + s.stop() + + +def _post_json(url: str, body: dict) -> tuple[int, bytes, list[tuple[str, str]]]: + req = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=5) as r: + return r.status, r.read(), list(r.headers.items()) + except urllib.error.HTTPError as e: + return e.code, e.read(), list(e.headers.items()) + + +def _get(url: str) -> tuple[int, bytes]: + try: + with urllib.request.urlopen(url, timeout=5) as r: + return r.status, r.read() + except urllib.error.HTTPError as e: + return e.code, e.read() + + +# --------------------------------------------------------------------------- +# tests +# --------------------------------------------------------------------------- + + +def test_proxy_runs_pipeline_on_adapted_routes(upstream, caplog) -> None: + proxy = start_adapter_proxy( + upstream_url=upstream.url, + adapters=[{"name": "logging", "config": {}}, {"name": "logging", "config": {}}], + ) + try: + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + status, body, _ = _post_json( + f"{proxy.url}/v1/chat/completions", + {"model": "stub", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert status == 200, body + assert json.loads(body)["choices"][0]["message"]["content"] == "stub-ok" + assert any("request POST /v1/chat/completions" in r.message for r in caplog.records) + finally: + proxy.stop() + + +def test_proxy_passthrough_does_not_run_pipeline(upstream, caplog) -> None: + proxy = start_adapter_proxy( + upstream_url=upstream.url, + adapters=[{"name": "logging", "config": {}}], + ) + try: + with caplog.at_level(logging.INFO, logger="nemo_gym.adapters.interceptors.request_logging"): + status, _ = _get(f"{proxy.url}/v1/models") + # Stub upstream returns 200 for GET /v1/models (via its catch-all) + assert status == 200 + # Crucially, the logging interceptor should NOT have fired for the passthrough + assert not any("request GET /v1/models" in r.message for r in caplog.records) + finally: + proxy.stop() + + +def test_proxy_multi_set_cookie_preserved(upstream) -> None: + proxy = start_adapter_proxy( + upstream_url=upstream.url, + adapters=[{"name": "logging", "config": {}}], + ) + try: + status, _body, headers = _post_json( + f"{proxy.url}/v1/chat/completions", + {"model": "stub", "messages": []}, + ) + assert status == 200 + cookies = [v for k, v in headers if k.lower() == "set-cookie"] + assert len(cookies) == 2, f"expected 2 Set-Cookie headers, got {len(cookies)}: {cookies}" + finally: + proxy.stop() + + +def test_proxy_rejects_remote_host_by_default() -> None: + with pytest.raises(ValueError, match="refusing host"): + start_adapter_proxy(upstream_url="http://x", adapters=[], host="0.0.0.0") + + +def test_proxy_allows_remote_with_explicit_flag(upstream) -> None: + proxy = start_adapter_proxy( + upstream_url=upstream.url, + adapters=[], + host="0.0.0.0", + unsafe_allow_remote=True, + port=0, + ) + try: + assert proxy.url.startswith("http://0.0.0.0:") + finally: + proxy.stop() + + +def test_proxy_rejects_user_supplied_endpoint() -> None: + with pytest.raises(ValueError, match="endpoint.*cannot be used"): + start_adapter_proxy( + upstream_url="http://x", + adapters=[{"name": "endpoint", "config": {"upstream_url": "http://y"}}], + ) + + +def test_proxy_handle_context_manager(upstream) -> None: + with start_adapter_proxy(upstream_url=upstream.url, adapters=[]) as proxy: + status, _ = _get(f"{proxy.url}/_proxy_health") + assert status == 200 + # After context exit, the server should be stopped — port should reject + # connections after a brief moment. We don't strictly assert this since + # uvicorn shutdown is async, but the context manager should have called stop(). diff --git a/tests/unit_tests/test_adapter_registry.py b/tests/unit_tests/test_adapter_registry.py new file mode 100644 index 0000000000..db6837de69 --- /dev/null +++ b/tests/unit_tests/test_adapter_registry.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""InterceptorRegistry tests.""" + +import pytest + +from nemo_gym.adapters.registry import InterceptorRegistry + + +def test_resolve_known_interceptor(): + cls = InterceptorRegistry.resolve_class("logging") + assert cls.__name__ == "Interceptor" + assert cls.__module__ == "nemo_gym.adapters.interceptors.request_logging" + + +def test_resolve_unknown_raises(): + with pytest.raises(ValueError, match="Unknown interceptor 'nonexistent'"): + InterceptorRegistry.resolve_class("nonexistent") + + +def test_create_with_config(): + instance = InterceptorRegistry.create("endpoint", {"upstream_url": "http://x"}) + from nemo_gym.adapters.interceptors.endpoint import Interceptor + + assert isinstance(instance, Interceptor) + + +def test_available_list(): + names = InterceptorRegistry.available() + assert isinstance(names, list) + for expected in ("endpoint", "logging"): + assert expected in names