Skip to content

Commit c4313fb

Browse files
committed
fix(adapters): address review feedback on the middleware framework
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>
1 parent 8e072a5 commit c4313fb

10 files changed

Lines changed: 462 additions & 44 deletions

File tree

nemo_gym/adapters/interceptors/endpoint.py

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ def __init__(
5050
if clean.endswith(suffix):
5151
clean = clean[: -len(suffix)]
5252
break
53+
# ``req.path`` already carries the ``/v1/...`` prefix, so strip a trailing
54+
# ``/v1`` from the base to avoid ``.../v1/v1/...`` double-prefixing.
55+
if clean.endswith("/v1"):
56+
clean = clean[: -len("/v1")]
5357
self._upstream_url = clean
5458
self._api_key = api_key
5559
self._extra_body = extra_body or {}
@@ -61,13 +65,6 @@ def __init__(
6165
# limits live on the global aiohttp client.
6266
self._max_concurrent = max_concurrent
6367

64-
@staticmethod
65-
def _normalize_content(body: dict[str, Any]) -> None:
66-
for choice in body.get("choices", []):
67-
msg = choice.get("message") or choice.get("delta") or {}
68-
if "content" in msg and msg["content"] is None:
69-
msg["content"] = ""
70-
7168
async def intercept_request(
7269
self,
7370
req: AdapterRequest,
@@ -105,7 +102,11 @@ async def intercept_request(
105102

106103
if status in self._retry_on_status and attempt < self._max_retries:
107104
retry_after = resp.headers.get("Retry-After")
108-
delay = float(retry_after) if retry_after else min(2**attempt, 60)
105+
try:
106+
# RFC 7231 allows an HTTP-date here; fall back on parse failure.
107+
delay = float(retry_after) if retry_after else min(2**attempt, 60)
108+
except ValueError:
109+
delay = min(2**attempt, 60)
109110
logger.warning(
110111
"endpoint: %s returned %d, retry %d/%d in %.1fs",
111112
url,
@@ -123,9 +124,6 @@ async def intercept_request(
123124
except (json.JSONDecodeError, UnicodeDecodeError):
124125
parsed = raw
125126

126-
if isinstance(parsed, dict):
127-
self._normalize_content(parsed)
128-
129127
return AdapterResponse(
130128
status_code=status,
131129
headers=resp_headers,

nemo_gym/adapters/middleware.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,10 @@ async def _adapter_middleware(request: Request, call_next): # noqa: ANN202
198198
async def _upstream(req: AdapterRequest) -> AdapterResponse:
199199
new_body = json.dumps(req.body).encode("utf-8")
200200
_override_request_body(request, new_body)
201+
# Re-point the ASGI scope at the prefix-stripped path so the router
202+
# can match it (the /s/<hex>/ session prefix is middleware-only).
203+
request.scope["path"] = req.path
204+
request.scope["raw_path"] = req.path.encode("latin-1")
201205
starlette_resp = await call_next(request)
202206
return await _starlette_response_to_adapter(starlette_resp, req.ctx)
203207

nemo_gym/adapters/pipeline.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from nemo_gym.adapters.types import (
2424
AdapterRequest,
2525
AdapterResponse,
26-
InterceptorContext,
26+
GracefulError,
2727
RequestInterceptor,
2828
RequestToResponseInterceptor,
2929
ResponseInterceptor,
@@ -100,8 +100,9 @@ async def process(
100100
is invoked with the (possibly mutated) request. Response interceptors
101101
then run in reverse order.
102102
"""
103-
ctx = InterceptorContext(request_id=request.ctx.request_id)
104-
set_context(ctx)
103+
# Reuse the request's own context so interceptors reading the global
104+
# get_context() observe the same ctx.extra (e.g. session_id) as req.ctx.
105+
set_context(request.ctx)
105106

106107
current: AdapterRequest | AdapterResponse = request
107108

@@ -113,6 +114,9 @@ async def process(
113114
try:
114115
result = await interceptor.intercept_request(current) # type: ignore[arg-type]
115116
current = result
117+
except GracefulError:
118+
# Control-flow signal (-> 429); never swallowed by best_effort.
119+
raise
116120
except Exception:
117121
if getattr(interceptor, "best_effort", False):
118122
logger.warning(
@@ -133,6 +137,8 @@ async def process(
133137
for interceptor in response_interceptors:
134138
try:
135139
response = await interceptor.intercept_response(response)
140+
except GracefulError:
141+
raise
136142
except Exception:
137143
if getattr(interceptor, "best_effort", False):
138144
logger.warning(

nemo_gym/adapters/proxy.py

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
from starlette.responses import JSONResponse, Response
4242

4343
from nemo_gym.adapters.pipeline import AdapterPipeline
44-
from nemo_gym.adapters.types import AdapterRequest, AdapterResponse, InterceptorContext
44+
from nemo_gym.adapters.types import AdapterRequest, AdapterResponse, GracefulError, InterceptorContext
4545

4646

4747
logger = logging.getLogger(__name__)
@@ -212,6 +212,8 @@ async def _run_pipeline(request: Request, path: str) -> Response:
212212

213213
async def _upstream(req: AdapterRequest) -> AdapterResponse:
214214
target = f"{upstream_url}{req.path}"
215+
if request.url.query:
216+
target = f"{target}?{request.url.query}"
215217
fwd_headers = {k: v for k, v in req.headers.items() if k.lower() not in ("host", "content-length")}
216218
import json as _json
217219

@@ -228,12 +230,6 @@ async def _upstream(req: AdapterRequest) -> AdapterResponse:
228230
except (ValueError, UnicodeDecodeError):
229231
parsed = raw
230232

231-
if isinstance(parsed, dict):
232-
for choice in parsed.get("choices", []):
233-
msg = choice.get("message") or choice.get("delta") or {}
234-
if isinstance(msg, dict) and msg.get("content") is None and "content" in msg:
235-
msg["content"] = ""
236-
237233
return AdapterResponse(
238234
status_code=resp.status,
239235
headers=resp_headers,
@@ -244,6 +240,18 @@ async def _upstream(req: AdapterRequest) -> AdapterResponse:
244240

245241
try:
246242
resp = await pipeline.process(adapter_req, upstream_call=_upstream)
243+
except GracefulError as exc:
244+
# Mirror middleware mode: budget/control-flow termination -> 429 (not 500).
245+
return JSONResponse(
246+
{
247+
"error": {
248+
"message": str(exc),
249+
"type": "invalid_request_error",
250+
"code": "session_budget_exhausted",
251+
},
252+
},
253+
status_code=429,
254+
)
247255
except Exception:
248256
logger.exception("adapter proxy pipeline error")
249257
return JSONResponse(

responses_api_agents/claude_code_agent/tests/test_app.py

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
import asyncio
1717
import json
18+
import types
1819
from pathlib import Path
1920
from unittest.mock import MagicMock, patch
2021

@@ -50,6 +51,14 @@ def _config(**kwargs) -> ClaudeCodeAgentConfig:
5051
def _make_agent(**kwargs) -> ClaudeCodeAgent:
5152
with patch("responses_api_agents.claude_code_agent.app.ClaudeCodeAgent.model_post_init"):
5253
agent = ClaudeCodeAgent(config=_config(**kwargs), server_client=MagicMock(spec=ServerClient))
54+
# Patching model_post_init also skips Pydantic's private-attr initialization,
55+
# leaving __pydantic_private__ as None; restore declared private slots (e.g.
56+
# _proxy_handle) to their defaults so agent methods can read them.
57+
object.__setattr__(
58+
agent,
59+
"__pydantic_private__",
60+
{name: pa.get_default() for name, pa in type(agent).__private_attributes__.items()},
61+
)
5362
agent.sem = asyncio.Semaphore(agent.config.concurrency)
5463
return agent
5564

@@ -58,6 +67,35 @@ def _event(type_: str, **kwargs) -> str:
5867
return json.dumps({"type": type_, **kwargs})
5968

6069

70+
def _run_claude_code_capturing(agent: ClaudeCodeAgent, tmp_path: Path) -> tuple[str, str, dict]:
71+
"""Run ``_run_claude_code`` with the subprocess stubbed out, capturing the
72+
child process ``env`` and ``argv`` so env-threading / model-name handling
73+
can be asserted. Mirrors the ``TestRunClaudeCode`` fixtures."""
74+
captured: dict = {}
75+
76+
class FakeProc:
77+
returncode = 0
78+
79+
async def communicate(self):
80+
return b'{"type":"result","usage":{"input_tokens":1,"output_tokens":1}}\n', b""
81+
82+
async def fake_exec(*cmd, **kwargs):
83+
captured["cmd"] = list(cmd)
84+
captured["env"] = kwargs["env"]
85+
return FakeProc()
86+
87+
# Clear the ambient environment so absence assertions (e.g. no
88+
# ANTHROPIC_BASE_URL in plain mode) are deterministic regardless of the
89+
# developer's/CI shell. The subprocess is mocked, so env is never spawned.
90+
with (
91+
patch.dict("responses_api_agents.claude_code_agent.app.os.environ", {}, clear=True),
92+
patch("responses_api_agents.claude_code_agent.app.Path.home", return_value=tmp_path),
93+
patch("responses_api_agents.claude_code_agent.app.asyncio.create_subprocess_exec", fake_exec),
94+
):
95+
stdout, model = asyncio.run(agent._run_claude_code("hello"))
96+
return stdout, model, captured
97+
98+
6199
class TestSanity:
62100
def test_config_defaults(self) -> None:
63101
cfg = _config()
@@ -234,6 +272,59 @@ async def fake_wait_for(coro, timeout):
234272
assert killed["called"] is True
235273
assert model == "claude-sonnet-4-6"
236274

275+
def test_proxy_mode_threads_base_url_and_preserves_full_model(self, tmp_path: Path) -> None:
276+
# A running adapter proxy in front of the model server: the SDK is
277+
# pointed at the proxy URL and the full (provider-prefixed) model name
278+
# is preserved for the custom upstream.
279+
agent = _make_agent(
280+
model="anthropic/claude-sonnet-4-6",
281+
anthropic_api_key="sk-test", # pragma: allowlist secret
282+
)
283+
agent._proxy_handle = types.SimpleNamespace(url="http://127.0.0.1:9999")
284+
285+
_stdout, model, captured = _run_claude_code_capturing(agent, tmp_path)
286+
env = captured["env"]
287+
288+
assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:9999"
289+
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-test" # pragma: allowlist secret
290+
# Full model name (provider prefix kept) for a custom upstream.
291+
assert model == "anthropic/claude-sonnet-4-6"
292+
assert env["ANTHROPIC_MODEL"] == "anthropic/claude-sonnet-4-6"
293+
assert captured["cmd"][captured["cmd"].index("--model") + 1] == "anthropic/claude-sonnet-4-6"
294+
295+
def test_explicit_base_url_preserves_full_model(self, tmp_path: Path) -> None:
296+
# A direct non-Anthropic endpoint via anthropic_base_url (no proxy):
297+
# base url is threaded through and the full model name is preserved.
298+
agent = _make_agent(
299+
model="anthropic/claude-sonnet-4-6",
300+
anthropic_base_url="https://my-endpoint.example/v1",
301+
anthropic_api_key="sk-direct", # pragma: allowlist secret
302+
)
303+
assert agent._proxy_handle is None
304+
305+
_stdout, model, captured = _run_claude_code_capturing(agent, tmp_path)
306+
env = captured["env"]
307+
308+
assert env["ANTHROPIC_BASE_URL"] == "https://my-endpoint.example/v1"
309+
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-direct" # pragma: allowlist secret
310+
assert model == "anthropic/claude-sonnet-4-6"
311+
assert env["ANTHROPIC_MODEL"] == "anthropic/claude-sonnet-4-6"
312+
313+
def test_plain_anthropic_mode_strips_provider_prefix(self, tmp_path: Path) -> None:
314+
# Plain Anthropic API mode (no proxy, no base url): the provider prefix
315+
# is stripped and no base-url/auth-token env is injected.
316+
agent = _make_agent(model="anthropic/claude-sonnet-4-6")
317+
assert agent._proxy_handle is None
318+
319+
_stdout, model, captured = _run_claude_code_capturing(agent, tmp_path)
320+
env = captured["env"]
321+
322+
assert "ANTHROPIC_BASE_URL" not in env
323+
assert "ANTHROPIC_AUTH_TOKEN" not in env
324+
assert model == "claude-sonnet-4-6"
325+
assert env["ANTHROPIC_MODEL"] == "claude-sonnet-4-6"
326+
assert captured["cmd"][captured["cmd"].index("--model") + 1] == "claude-sonnet-4-6"
327+
237328

238329
class TestExtractInstruction:
239330
def test_user_only(self) -> None:

tests/unit_tests/test_adapter_endpoint.py

Lines changed: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -63,27 +63,15 @@ def _req(body: dict | None = None, path: str = "/v1/chat/completions") -> Adapte
6363

6464

6565
def test_upstream_url_known_suffixes_stripped():
66+
# API suffixes AND a trailing /v1 are stripped so the /v1/... prefix already
67+
# carried on req.path does not double up.
6668
for suffix in ("/chat/completions", "/completions", "/embeddings"):
6769
ic = Interceptor(upstream_url=f"https://api.example.com/v1{suffix}/")
68-
assert ic._upstream_url == "https://api.example.com/v1"
69-
# A non-API suffix is left intact.
70-
assert Interceptor(upstream_url="http://up/v1")._upstream_url == "http://up/v1"
71-
72-
73-
def test_normalize_content_replaces_none_only():
74-
body = {
75-
"choices": [
76-
{"message": {"content": None}},
77-
{"delta": {"content": None}},
78-
{"message": {"content": "keep"}},
79-
{"message": {}}, # no content key -> untouched
80-
]
81-
}
82-
Interceptor._normalize_content(body)
83-
assert body["choices"][0]["message"]["content"] == ""
84-
assert body["choices"][1]["delta"]["content"] == ""
85-
assert body["choices"][2]["message"]["content"] == "keep"
86-
assert "content" not in body["choices"][3]["message"]
70+
assert ic._upstream_url == "https://api.example.com"
71+
# A bare /v1 base is stripped to the host root.
72+
assert Interceptor(upstream_url="http://up/v1")._upstream_url == "http://up"
73+
# A host with no /v1 is left intact.
74+
assert Interceptor(upstream_url="http://up")._upstream_url == "http://up"
8775

8876

8977
async def test_intercept_request_success_merges_body_and_headers():
@@ -107,9 +95,9 @@ async def test_intercept_request_success_merges_body_and_headers():
10795
assert kwargs["headers"]["Content-Type"] == "application/json"
10896
assert "Host" not in kwargs["headers"] and "Content-Length" not in kwargs["headers"]
10997
assert kwargs["headers"]["X-Keep"] == "1"
110-
assert kwargs["url"] == "http://up/v1/v1/chat/completions"
111-
# None content normalized to ""
112-
assert resp.body["choices"][0]["message"]["content"] == ""
98+
assert kwargs["url"] == "http://up/v1/chat/completions"
99+
# content:null passes through unchanged (no proxy-side normalization)
100+
assert resp.body["choices"][0]["message"]["content"] is None
113101
# response headers preserved (original case) as latin-1 byte tuples
114102
assert (b"Content-Type", b"application/json") in resp.headers
115103

@@ -127,6 +115,26 @@ async def test_intercept_request_retries_on_status_then_succeeds():
127115
sleep.assert_awaited() # Retry-After honored
128116

129117

118+
async def test_intercept_request_http_date_retry_after_falls_back_to_backoff():
119+
"""An HTTP-date ``Retry-After`` (RFC 7231 allows it) is not a float, so the
120+
parse must fall back to exponential backoff instead of crashing the retry
121+
loop. The request must still retry and eventually succeed."""
122+
ic = Interceptor(upstream_url="http://up/v1", max_retries=1, retry_on_status=[503])
123+
seq = [
124+
_FakeResp(503, {"Retry-After": "Wed, 21 Oct 2099 07:28:00 GMT"}),
125+
_FakeResp(200, {}, {"ok": True}),
126+
]
127+
with (
128+
patch(f"{_ENDPOINT}.global_request", new=AsyncMock(side_effect=seq)),
129+
patch(f"{_ENDPOINT}.asyncio.sleep", new=AsyncMock()) as sleep,
130+
):
131+
resp = await ic.intercept_request(_req())
132+
assert resp.status_code == 200
133+
assert resp.body == {"ok": True}
134+
# Backoff path was taken (date header could not be parsed as a delay).
135+
sleep.assert_awaited_once_with(1)
136+
137+
130138
async def test_intercept_request_non_json_response_passed_through_as_bytes():
131139
ic = Interceptor(upstream_url="http://up/v1")
132140
with patch(

tests/unit_tests/test_adapter_internals.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,18 @@ async def test_starlette_to_adapter_non_json_content_type_kept_as_bytes():
108108
assert out.body == b"raw"
109109

110110

111+
async def test_starlette_to_adapter_str_body_without_iterator_is_encoded():
112+
# No ``body_iterator`` and a ``str`` body exercises the non-streaming
113+
# str->bytes encode branch.
114+
class _StrBodyResp:
115+
status_code = 200
116+
raw_headers: list = []
117+
body = "plain-str-body"
118+
119+
out = await mw._starlette_response_to_adapter(_StrBodyResp(), InterceptorContext())
120+
assert out.body == b"plain-str-body"
121+
122+
111123
def test_adapter_response_to_starlette_bytes_body():
112124
out = mw._adapter_response_to_starlette(AdapterResponse(status_code=200, headers={"X-A": "1"}, body=b"raw-bytes"))
113125
assert isinstance(out, Response)

0 commit comments

Comments
 (0)