feat(adapters): adapter middleware framework (base)#1646
Conversation
|
🌿 Preview your docs: https://nvidia-preview-feat-adapter-base.docs.buildwithfern.com/nemo/gym Here are the markdown pages you've updated: |
Hey @Glorf, thank you for the docs. Will we also have pointers for the SimpleResponsesAPIAgent, or SimpleResourcesServer somewhere in the docs? Or maybe we should move the adapter framework one level up on the docs organization. |
|
@Glorf, Could you link a small Design Note here as well?
|
|
My current understanding of the Topology is something like this: Is that correct? In this case, should we handle how the proxy take its upstream url from a resolved |
| @@ -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]: | |||
There was a problem hiding this comment.
When _proxy_handle is set, base_url (resolved from the model_server ref via _resolve_base_url()) is computed here but then discarded below — proxy_or_base = self._proxy_handle.url wins, and the proxy's own upstream comes from the static adapter_proxy.upstream_url. So model_server and the proxy are two independent notions of "where the model is," and they don't feed each other.
That's fine for a fixed/external upstream, but it means proxy mode can't target an in-cluster Gym model server resolved from model_server (OS-assigned port) — e.g. one exposing /v1/messages. The dynamic URL just isn't expressible via the static upstream_url.
Is "external upstreams only" the intended scope for this PR? If we later want the capture proxy in front of a Gym model server, the seam is small — either upstream = self._resolve_base_url() or cfg.upstream_url here, or an upstream_model_server: Optional[ModelServerRef] on AdapterProxyConfig resolved at setup_webserver. Non-blocking — just flagging it so the decoupling is a deliberate choice.
ffrujeri
left a comment
ffrujeri
left a comment
There was a problem hiding this comment.
Adapter middleware framework — review summary
Thanks for this — it's a clean, well-structured base framework and the test suite genuinely exercises the hard HTTP behaviors (body replay, multi-Set-Cookie preservation, /s/<hex>/ parsing, GracefulError→429) rather than smoke-testing.
Two correctness items worth a look before merge:
- Session-prefixed paths 404 in middleware mode unless an interceptor short-circuits (
middleware.py:185). get_context()diverges fromreq.ctxbecauseprocess()dropsctx.extra(pipeline.py:103) — session-aware interceptors using the global accessor losesession_id.
The rest are medium/minor: a GracefulError→429-vs-500 asymmetry between middleware and proxy modes, hardcoded content:null normalization in the transport path (vs the "no translation in the proxy" principle), endpoint being unreachable while the docs say it's required by start_adapter_proxy, adapter_proxy being silently ignored on harbor_agent/mini_swe_agent, plus assorted nits.
Notes:
- Coverage of the
adapterspackage sits at exactly the 96%fail_underfloor with zero headroom (proxy.py90%); see the inline coverage note.
| path = request.url.path | ||
| ctx = InterceptorContext() | ||
|
|
||
| m = _SESSION_PATH_RE.match(path) |
There was a problem hiding this comment.
nemo_gym/adapters/middleware.py:185
[BUG] Session-prefixed paths 404 in middleware mode. /s/<hex>/ is stripped into adapter_req.path + ctx.extra["session_id"], but request.scope["path"] is never rewritten, so the default _upstream's call_next(request) routes on the original prefixed path → FastAPI 404. This works only if an interceptor short-circuits; with the shipped logging/endpoint interceptors, POST /s/deadbeef/v1/chat/completions → 404 (reproduced). The adapters.mdx doc presents Path-Based Session Scoping as a working feature. Suggested fix in _upstream before call_next:
| m = _SESSION_PATH_RE.match(path) | |
| async def _upstream(req: AdapterRequest) -> AdapterResponse: | |
| new_body = json.dumps(req.body).encode("utf-8") | |
| _override_request_body(request, new_body) | |
| # Re-point the ASGI scope at the prefix-stripped path so the | |
| # router can match it (session prefix is middleware-only). | |
| request.scope["path"] = req.path | |
| request.scope["raw_path"] = req.path.encode("latin-1") | |
| starlette_resp = await call_next(request) | |
| return await _starlette_response_to_adapter(starlette_resp, req.ctx) |
| is invoked with the (possibly mutated) request. Response interceptors | ||
| then run in reverse order. | ||
| """ | ||
| ctx = InterceptorContext(request_id=request.ctx.request_id) |
There was a problem hiding this comment.
nemo_gym/adapters/pipeline.py:103
[BUG] process() builds a fresh InterceptorContext copying only request_id, not request.ctx.extra, so get_context() returns an empty ctx (no session_id) while req.ctx is populated — req.ctx is get_context() → False (reproduced). Any session-aware interceptor using the documented get_context() accessor silently gets nothing, which undermines the session-scoping feature above. Reusing the request's own ctx fixes it:
| ctx = InterceptorContext(request_id=request.ctx.request_id) | |
| set_context(request.ctx) |
Secondary nit (pipeline.py:104): set_context() is never paired with a ContextVar.reset(). Low impact today (Starlette gives each request its own task/context), but consider capturing the token and resetting in a finally so the var doesn't persist after process() returns.
| retry_on_status: list[int] | None = None, | ||
| max_concurrent: int = 64, | ||
| ) -> None: | ||
| clean = upstream_url.rstrip("/") |
There was a problem hiding this comment.
nemo_gym/adapters/interceptors/endpoint.py:48
[BUG] upstream_url only has /chat/completions|/completions|/embeddings suffixes stripped, then intercept_request builds f"{self._upstream_url}{req.path}" with req.path already starting /v1/.... A common base like http://host:8000/v1 yields http://host:8000/v1/v1/chat/completions (reproduced; the double-prefixed URL also appears unasserted in this PR's endpoint test logs). The same pattern is reachable in proxy mode via adapter_proxy.upstream_url (proxy.py:214). Could you either also strip a trailing /v1 or document the expected upstream_url format (bare host root)?
|
|
||
| try: | ||
| resp = await pipeline.process(adapter_req, upstream_call=_upstream) | ||
| except Exception: |
There was a problem hiding this comment.
nemo_gym/adapters/proxy.py:247
[BUG] GracefulError is mapped to a 429 (session_budget_exhausted) in middleware.py:206, but here the bare except Exception turns the same error into a 500. The same interceptor chain returns different statuses depending on host mode. Consider adding a except GracefulError: arm mirroring the middleware path. (These lines are also currently uncovered — see the coverage note.)
| except (ValueError, UnicodeDecodeError): | ||
| parsed = raw | ||
|
|
||
| if isinstance(parsed, dict): |
There was a problem hiding this comment.
nemo_gym/adapters/proxy.py:231
[GUIDELINE] This content: null → "" normalization is hardcoded unconditionally in the transport-forwarding path (here and again in endpoint.py:65-69,127) and duplicated across both. That's semantic payload mutation in the proxy, which runs against the "no translation in the proxy — translation stays in anthropic_converter" principle. There's also a third asymmetry: middleware.py does not normalize, so middleware-mode responses keep content: null. Consider moving this into an opt-in ResponseInterceptor so the behavior is uniform and explicit.
| _server: uvicorn.Server | ||
| _thread: threading.Thread | ||
|
|
||
| def stop(self, timeout: float = 5.0) -> None: |
There was a problem hiding this comment.
[ASYNC, minor] ProxyHandle.stop has no idempotency guard and no log if join(timeout=5) fails to exit — a silent daemon-thread/port leak for a long-lived process that restarts proxies. Minor (the daemon thread dies with the process), but a guard + a warning on join timeout would help.
| headers.setdefault("Content-Type", "application/json") | ||
|
|
||
| attempt = 0 | ||
| while True: |
There was a problem hiding this comment.
nemo_gym/adapters/interceptors/endpoint.py:86
[ASYNC, minor] This retry loop stacks on top of global_request's built-in 3-try retry, so on exception paths you can get ~3×(max_retries+1) attempts with two backoff schemes. Fine at the default max_retries=0; worth a comment so users who raise max_retries understand the multiplier.
| """ | ||
|
|
||
| upstream_url: str | ||
| adapters: list[InterceptorSpec] = Field(default_factory=list) |
There was a problem hiding this comment.
nemo_gym/adapters/types.py:137
[CONFIG, minor] adapters is typed list[dict] on the three base configs but list[InterceptorSpec] here in AdapterProxyConfig. The raw-dict form gets no Hydra schema validation until the runtime registry name lookup — consider unifying on InterceptorSpec for the top-level adapters too.
| return 200 <= self.status_code < 400 | ||
|
|
||
|
|
||
| class Stage(enum.Enum): |
There was a problem hiding this comment.
[DOCSTRING, nit] A few public API symbols lack docstrings: Stage, RequestInterceptor, ResponseInterceptor, get_context/set_context (here in types.py), plus endpoint.Interceptor (and its no-op close() stub, which has no caller) and request_logging.Interceptor. Short prose docstrings would help since these are the framework's extension surface.
|
|
||
| t0 = time.perf_counter() | ||
| async with session.post(target, data=_json.dumps(req.body), headers=fwd_headers) as resp: | ||
| raw = await resp.read() |
There was a problem hiding this comment.
nemo_gym/adapters/proxy.py:220
[DESIGN] Reinforcing the design-note question on SSE: the proxy fully buffers (raw = await resp.read() here, then json.loads with a raw-bytes fallback at :227) — no incremental streaming. For a long generation the Claude CLI receives nothing until upstream completes, which risks its first-event/idle timeout (apparent hang). The shipped logging interceptor already guards bytes (request_logging.py:34) so it won't crash, but capture degrades to a 512-char preview, and any future dict-assuming ResponseInterceptor would break on an SSE body. Likely intended scope for PR 1 — flagging so the SSE limitation is a deliberate, documented choice.
Glorf
force-pushed
the
feat/ecs-fargate-sandbox
branch
from
a80772d to
11124f9
Compare
Per @ffrujeri's review of #1646: - middleware: re-point the ASGI scope path so /s/<hex>/ session-prefixed requests route to the real handler instead of 404'ing - pipeline: reuse request.ctx so get_context() exposes ctx.extra (session_id); never swallow GracefulError under best_effort (both phases) - proxy: map GracefulError -> 429 (parity with middleware mode); preserve the query string when forwarding to the upstream - endpoint: strip a trailing /v1 from upstream_url (avoid /v1/v1); guard the Retry-After parse against HTTP-date values - drop the proxy/endpoint content:null normalization (translation belongs in the /v1/messages route, not the transport proxy) - tests: lock in each fix + add claude_code env/model-name threading coverage; adapters package coverage 95% -> 100% Signed-off-by: Michal Bien <mbien@nvidia.com>
Introduces an interceptor-based middleware framework that runs at every
in-tree server's request/response boundary. PR 1 of 4 — base framework
only. Follow-on PRs add observability / caching / request-rewriting
interceptor families.
Framework
nemo_gym/adapters/
pipeline.py AdapterPipeline + stage-order validation
(REQUEST → REQUEST_TO_RESPONSE → RESPONSE)
middleware.py install_middleware(app, specs) — FastAPI middleware
that wraps call_next; body replay, multi-Set-Cookie
preservation, /s/<hex>/... session-prefix routing,
GracefulError → 429
proxy.py start_adapter_proxy(upstream_url, adapters) —
localhost uvicorn host mode for external-inference
agents whose SDK respects *_BASE_URL env vars
registry.py short-name → Interceptor class + runtime register()
types.py AdapterRequest/Response, three Interceptor ABCs,
Stage enum, GracefulError, ContextVar-backed
per-request context, InterceptorSpec and
AdapterProxyConfig pydantic models
interceptors/
endpoint.py Drives upstream HTTP call (required by proxy mode,
forbidden inside install_middleware)
request_logging.py
Canonical "did the chain fire?" probe (logs body
keys + response status/latency)
Server wire-up
adapters: list[dict] | None = None on three base configs:
BaseResponsesAPIModelConfig
BaseResponsesAPIAgentConfig
BaseResourcesServerConfig
install_middleware(app, self.config.adapters) at the tail of each
base's setup_webserver. Every in-tree server inheriting from
SimpleResponsesAPI{Model,Agent} or SimpleResourcesServer picks up
the adapters knob automatically.
adapter_proxy: AdapterProxyConfig | None = None on
BaseResponsesAPIAgentConfig. When set, SimpleResponsesAPIAgent.
setup_webserver starts a localhost uvicorn proxy in a daemon thread,
stores the ProxyHandle on self._proxy_handle, registers atexit cleanup.
Per-server override fixes
harbor_agent and mini_swe_agent override setup_webserver without
calling super(). Inline install_middleware(app, self.config.adapters)
calls added to both so the adapters field still applies.
claude_code_agent integration
Reads self._proxy_handle.url when adapter_proxy is set; threads the
proxy URL into the claude CLI subprocess via ANTHROPIC_BASE_URL +
ANTHROPIC_AUTH_TOKEN. Preserves the full model-name prefix in proxy
mode (no .split("/")[-1] stripping).
Safety
- Stage ordering validated at startup
- Unknown interceptor name raises at config-validation time
- install_middleware rejects `endpoint` in chain (host already forwards)
- start_adapter_proxy rejects user-supplied `endpoint` AND refuses
host="0.0.0.0" unless unsafe_allow_remote=True (otherwise leaks
the upstream API key to any caller on the network)
- best_effort=True interceptors swallow exceptions; strict ones
propagate
Tests
42 tests, framework-level coverage. Follow-on PRs add per-interceptor
test suites.
Docs
fern/versions/latest/pages/model-server/adapters.mdx new
fern/versions/latest/pages/model-server/index.mdx link added
Signed-off-by: Michal Bien <mbien@nvidia.com>
Raise core coverage past the 96% gate (93.29% -> 96%) with real unit tests for the previously-untested adapter paths: - endpoint interceptor: upstream-URL suffix stripping, None-content normalization, success path (body/header merge, api-key injection, hop-by-hop stripping), retry-on-status, timeout -> 504, ClientError raise/retry, and close(). - adapter internals: lazy context creation in a fresh contextvars.Context, AdapterResponse.ok, request_logging._trunc_preview branches, middleware request/response conversion (streaming/plain/non-json bodies, bytes body) and the endpoint guard, proxy host/endpoint guards + _bind_port/ _wait_for_health, and registry/pipeline error + best-effort paths. All upstream HTTP/SSH is mocked; tests are offline and deterministic. Signed-off-by: Michal Bien <mbien@nvidia.com>
Per @ffrujeri's review of #1646: - middleware: re-point the ASGI scope path so /s/<hex>/ session-prefixed requests route to the real handler instead of 404'ing - pipeline: reuse request.ctx so get_context() exposes ctx.extra (session_id); never swallow GracefulError under best_effort (both phases) - proxy: map GracefulError -> 429 (parity with middleware mode); preserve the query string when forwarding to the upstream - endpoint: strip a trailing /v1 from upstream_url (avoid /v1/v1); guard the Retry-After parse against HTTP-date values - drop the proxy/endpoint content:null normalization (translation belongs in the /v1/messages route, not the transport proxy) - tests: lock in each fix + add claude_code env/model-name threading coverage; adapters package coverage 95% -> 100% Signed-off-by: Michal Bien <mbien@nvidia.com>
Per @ffrujeri's review of #1646: - middleware: re-point the ASGI scope path so /s/<hex>/ session-prefixed requests route to the real handler instead of 404'ing - pipeline: reuse request.ctx so get_context() exposes ctx.extra (session_id); never swallow GracefulError under best_effort (both phases) - proxy: map GracefulError -> 429 (parity with middleware mode); preserve the query string when forwarding to the upstream - endpoint: strip a trailing /v1 from upstream_url (avoid /v1/v1); guard the Retry-After parse against HTTP-date values - drop the proxy/endpoint content:null normalization (translation belongs in the /v1/messages route, not the transport proxy) - tests: lock in each fix + add claude_code env/model-name threading coverage; adapters package coverage 95% -> 100% Signed-off-by: Michal Bien <mbien@nvidia.com>
ECS Fargate `SandboxProvider` on top of the sandbox API (#1377, now merged to `main`). Auto-mirrors public images to ECR on demand; SSH reverse-tunnel for exec / file-transfer / model egress. Rebased onto `main` now that the sandbox base (#1377) has landed — this PR stands on its own and is ready for review. **Sandbox-bound agents line of work:** **ECS Fargate (this PR)** → adapter middleware base (#1646) → sandbox-bound CLI agents + capture (#1647). Interceptor follow-ups on the adapter base: #1649 (caching), #1650 (observability), #1651 (rewrites). 3 commits, +5974: ECS provider + engine, on-demand ECR mirroring + conda activation fix, and 233 unit tests (engine.py to 99%). `uv.lock` reconciled with main's security upgrades (#1657). --------- Signed-off-by: Michal Bien <mbien@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2 participants
Introduces an interceptor-based middleware framework that runs at every in-tree server's request/response boundary. PR 1 of 4 — base framework only. Follow-on PRs add observability / caching / request-rewriting interceptor families.
Framework
nemo_gym/adapters/
pipeline.py AdapterPipeline + stage-order validation
(REQUEST → REQUEST_TO_RESPONSE → RESPONSE)
middleware.py install_middleware(app, specs) — FastAPI middleware
that wraps call_next; body replay, multi-Set-Cookie
preservation, /s//... session-prefix routing,
GracefulError → 429
proxy.py start_adapter_proxy(upstream_url, adapters) —
localhost uvicorn host mode for external-inference
agents whose SDK respects *_BASE_URL env vars
registry.py short-name → Interceptor class + runtime register()
types.py AdapterRequest/Response, three Interceptor ABCs,
Stage enum, GracefulError, ContextVar-backed
per-request context, InterceptorSpec and
AdapterProxyConfig pydantic models
interceptors/
endpoint.py Drives upstream HTTP call (required by proxy mode,
forbidden inside install_middleware)
request_logging.py
Canonical "did the chain fire?" probe (logs body
keys + response status/latency)
Server wire-up
adapters: list[dict] | None = None on three base configs:
BaseResponsesAPIModelConfig
BaseResponsesAPIAgentConfig
BaseResourcesServerConfig
install_middleware(app, self.config.adapters) at the tail of each
base's setup_webserver. Every in-tree server inheriting from
SimpleResponsesAPI{Model,Agent} or SimpleResourcesServer picks up
the adapters knob automatically.
adapter_proxy: AdapterProxyConfig | None = None on
BaseResponsesAPIAgentConfig. When set, SimpleResponsesAPIAgent.
setup_webserver starts a localhost uvicorn proxy in a daemon thread,
stores the ProxyHandle on self._proxy_handle, registers atexit cleanup.
Per-server override fixes
harbor_agent and mini_swe_agent override setup_webserver without
calling super(). Inline install_middleware(app, self.config.adapters)
calls added to both so the adapters field still applies.
claude_code_agent integration
Reads self._proxy_handle.url when adapter_proxy is set; threads the
proxy URL into the claude CLI subprocess via ANTHROPIC_BASE_URL +
ANTHROPIC_AUTH_TOKEN. Preserves the full model-name prefix in proxy
mode (no .split("/")[-1] stripping).
Safety
endpointin chain (host already forwards)endpointAND refuses host="0.0.0.0" unless unsafe_allow_remote=True (otherwise leaks the upstream API key to any caller on the network)Tests
42 tests, framework-level coverage. Follow-on PRs add per-interceptor
test suites.
Docs
fern/versions/latest/pages/model-server/adapters.mdx new
fern/versions/latest/pages/model-server/index.mdx link added
Recreated with an upstream head so it can join the GitHub stack (supersedes #1518, which had a fork head).