openvidia 2.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,1374 @@
1
+ """
2
+ Responses API → chat/completions shim.
3
+
4
+ Translates /v1/responses (the API used by Codex CLI) into
5
+ /v1/chat/completions, which the proxy already knows how to forward to NVIDIA
6
+ NIM. The translation is bidirectional:
7
+
8
+ - request: input (string | items[]) → messages[]
9
+ - response: chat completion → output items (text, function_call)
10
+ - streaming: SSE chat chunks → SSE Responses events
11
+ - tools: function definitions → chat tools, and back
12
+
13
+ No abstractions — just payload translation. All names, signatures and control
14
+ flow are frozen; only prose (docstrings/comments) is reformatted.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import asyncio
20
+ import json
21
+ import time
22
+ import uuid
23
+ from collections.abc import AsyncGenerator
24
+ from typing import Any
25
+
26
+ import httpx
27
+ from fastapi import Request
28
+ from fastapi.responses import JSONResponse, StreamingResponse
29
+
30
+ from . import config
31
+ from ._upstream_utils import get_upstream_sem, is_resource_exhausted
32
+ from .proxy_state import ProxyState
33
+
34
+ # ── Bounded rotation + saturation fast-fail ─────────────────────────────
35
+ # Codex CLI blocks historically when /v1/responses rotated serially across
36
+ # ALL 25 candidate keys with no per-attempt timeout (client default read=
37
+ # 120s): 25 × 120s = up to 50 minutes before returning 503. We now cap the
38
+ # number of keys we actually send to per rotation phase (the weighted
39
+ # least-loaded ordering already puts the best keys first), give each send a
40
+ # bounded connect+read+write+pool timeout, and fast-fail the whole loop as
41
+ # soon as the live (RPM-eligible) pool is too small to bother.
42
+ # Upstream gateway timeouts: the provider's edge gave up waiting for the
43
+ # model. Every key would hit the same wall, so these must not be charged to
44
+ # the key that happened to carry the request.
45
+ _GATEWAY_TIMEOUTS = {502, 503, 504}
46
+
47
+ _MAX_ROTATE_ATTEMPTS = 5 # hard cap on sends per model phase
48
+ _ROTATE_SEND_TIMEOUT = httpx.Timeout(**config.httpx_timeout_kwargs())
49
+ _MODEL_PROBE_TIMEOUT = httpx.Timeout(connect=5.0, read=90.0, write=10.0, pool=95.0)
50
+ _MIN_LIVE_FRACTION = 0.05 # <5% of valid keys live → 503 fast
51
+
52
+
53
+ # A reasoning model withholds its first byte until it has finished thinking —
54
+ # measured at 117-144s for z-ai/glm-5.2 on the NVIDIA free tier, for prompts
55
+ # of any size. During that window the SSE socket carries nothing, and a client
56
+ # (or any proxy in between) cannot tell a thinking model from a dead
57
+ # connection. Emit an SSE comment periodically: comments are stripped by the
58
+ # SSE parser before events are dispatched, so the client's protocol handling
59
+ # never sees them — the bytes exist purely to keep the stream demonstrably
60
+ # alive.
61
+ _KEEPALIVE_INTERVAL = 10.0
62
+ _KEEPALIVE_BYTES = b": keepalive\n\n"
63
+
64
+
65
+ async def _keepalive_until(task, result_box: list):
66
+ """Yield keepalive comments until ``task`` finishes; append its result."""
67
+ while True:
68
+ try:
69
+ result_box.append(
70
+ await asyncio.wait_for(asyncio.shield(task), timeout=_KEEPALIVE_INTERVAL)
71
+ )
72
+ return
73
+ except TimeoutError:
74
+ yield _KEEPALIVE_BYTES
75
+ except Exception: # noqa: BLE001 — surfaced to the caller via the box
76
+ result_box.append((None, None, None))
77
+ return
78
+
79
+
80
+ # Statuses that describe the REQUEST, not the key. Every key in the pool would
81
+ # answer these identically, so rotating just converts one malformed request
82
+ # into N cooled-down keys. This is the single most expensive thing the rotation
83
+ # loop used to do: a long conversation whose tool-call pairing broke would take
84
+ # five keys out of service per attempt.
85
+ _DETERMINISTIC_STATUSES = frozenset({400, 404, 405, 413, 415, 422})
86
+
87
+ # Bound on how long a phase will sit in the ResourceExhausted back-off, which
88
+ # is refunded rather than counted against ``max_attempts``. Without a separate
89
+ # cap the refund makes that branch unbounded: 3 passes over a 25-key pool at
90
+ # 0.8s each is a minute of the caller waiting under the semaphore.
91
+ _MAX_EXHAUSTED_WAITS = 6
92
+
93
+
94
+ async def _send_once(client, method, upstream, payload, content, hdrs, send_timeout, stream):
95
+ """Build and send one upstream request, retrying once on a GOAWAY.
96
+
97
+ HTTP/2 servers recycle connections by sending GOAWAY, which httpx surfaces
98
+ as RemoteProtocolError(ConnectionTerminated). It says nothing about the key
99
+ — it is the pooled socket reaching end of life. Treating it as a key fault
100
+ is how four keys got marked bad in the same second for a single connection
101
+ handover. The right response is to send it again, on the same key, over a
102
+ fresh connection.
103
+ """
104
+ for goaway_retry in range(2):
105
+ if content is not None:
106
+ req = client.build_request(
107
+ method, upstream, content=content, headers=hdrs, timeout=send_timeout
108
+ )
109
+ else:
110
+ req = client.build_request(
111
+ method, upstream, json=payload, headers=hdrs, timeout=send_timeout
112
+ )
113
+ try:
114
+ # Global concurrency semaphore (shared with proxy_app): never
115
+ # exceed the cap on concurrent upstream sends.
116
+ async with get_upstream_sem():
117
+ return await client.send(req, stream=stream)
118
+ except httpx.RemoteProtocolError:
119
+ if goaway_retry:
120
+ raise
121
+ continue
122
+ raise AssertionError("unreachable") # pragma: no cover
123
+
124
+
125
+ def json_headers(key: str, _idx: int = 0) -> dict[str, str]:
126
+ """The upstream header set for a JSON request on ``key``.
127
+
128
+ One definition instead of five byte-identical closures. ``_idx`` exists
129
+ because _rotation_phase calls its header factory with (key, index); no
130
+ variant ever used it.
131
+ """
132
+ return {
133
+ "Authorization": f"Bearer {key}",
134
+ "Content-Type": "application/json",
135
+ "User-Agent": "openvidia/2.0",
136
+ }
137
+
138
+
139
+ def _detail_from_body(body: str) -> str:
140
+ """The human-readable sentence out of a NIM error body, if there is one."""
141
+ body = (body or "").strip()
142
+ if not body:
143
+ return ""
144
+ try:
145
+ parsed = json.loads(body)
146
+ except (json.JSONDecodeError, TypeError):
147
+ return body[:400]
148
+ if isinstance(parsed, dict):
149
+ detail = parsed.get("detail") or parsed.get("message") or parsed.get("title")
150
+ if isinstance(parsed.get("error"), dict):
151
+ detail = parsed["error"].get("message") or detail
152
+ if isinstance(detail, str) and detail.strip():
153
+ return detail.strip()[:400]
154
+ return body[:400]
155
+
156
+
157
+ def _upstream_error_message(outcome: dict | None, fallback: str = "all keys failed") -> str:
158
+ """Turn a rotation outcome into something a user can act on.
159
+
160
+ "all keys exhausted" was true and useless: it hid the provider's own
161
+ explanation — context overflow, unsupported parameter, unknown model —
162
+ behind a message about key management. When upstream said something,
163
+ say it.
164
+ """
165
+ if not outcome:
166
+ return fallback
167
+ status = outcome.get("status")
168
+ detail = _detail_from_body(outcome.get("body") or "")
169
+
170
+ # A deterministic rejection is not a key problem, so it must never be
171
+ # described as one. "all keys failed (last upstream status 404)" sent
172
+ # people hunting through their key pool for what was a path the provider
173
+ # simply does not serve — no key was even tried twice.
174
+ if outcome.get("deterministic"):
175
+ if detail:
176
+ return f"upstream rejected the request ({status}): {detail}"
177
+ return f"upstream rejected the request with {status}"
178
+
179
+ if not detail:
180
+ return f"{fallback} (last upstream status {status})" if status else fallback
181
+ return f"upstream {status}: {detail}" if status else detail
182
+
183
+
184
+ async def _rotation_phase(
185
+ client,
186
+ upstream,
187
+ payload,
188
+ headers_factory,
189
+ state,
190
+ candidates,
191
+ *,
192
+ max_attempts,
193
+ timeout,
194
+ stream,
195
+ log_tag,
196
+ outcome_box=None,
197
+ method="POST",
198
+ content=None,
199
+ probe_timeout=None,
200
+ ):
201
+ """Single bounded rotation phase. Returns (resp_or_None, used_key, used_idx).
202
+
203
+ Parameters:
204
+ method: HTTP method (default POST, used by shims; catch-all passes request.method).
205
+ content: raw body bytes (used by catch-all; when None, json=payload is used).
206
+ probe_timeout: if given, the first attempt uses this timeout to
207
+ detect dead models; subsequent attempts use ``timeout``.
208
+ outcome_box: optional dict filled with the last upstream status and
209
+ error body. The caller needs it to answer with what the provider
210
+ actually said instead of a generic "all keys exhausted", which is
211
+ the only place a context-overflow message or an "unsupported
212
+ parameter" ever appears.
213
+ """
214
+
215
+ def _record(status=None, body="", deterministic=False, exhausted=False):
216
+ if outcome_box is not None:
217
+ outcome_box["status"] = status
218
+ outcome_box["body"] = body
219
+ outcome_box["deterministic"] = deterministic
220
+ # The caller must not score a full upstream worker pool against the
221
+ # model either — see the circuit-breaker note at the send site.
222
+ outcome_box["exhausted"] = exhausted
223
+
224
+ attempts = 0
225
+ exhausted_waits = 0
226
+ network_errors = 0 # consecutive ConnectTimeout/ConnectError across keys
227
+ _MAX_NETWORK_ERRORS = 3 # if 3+ keys fail to connect, the network is down
228
+ for pass_num in range(3):
229
+ if pass_num > 0:
230
+ await asyncio.sleep(1.0)
231
+ for idx, k in candidates:
232
+ # Return, not break: breaking only left the inner loop, so a phase
233
+ # that had spent its attempt budget still paid two 1s sleeps before
234
+ # giving up.
235
+ if attempts >= max_attempts:
236
+ return None, None, None
237
+ # An account-wide limit is shared by every key in the pool, so the
238
+ # next one cannot beat it. Checked before sending, not after: the
239
+ # point is to stop spending keys, and that includes the key this
240
+ # iteration was about to spend.
241
+ if state.is_pool_throttled():
242
+ _record(status=429, body="account-wide rate limit")
243
+ return None, None, None
244
+ if not state.key_can_send_rpm(k) or state.is_key_on_cooldown(k):
245
+ continue
246
+ attempts += 1
247
+ hdrs = headers_factory(k, idx)
248
+ # Count the request against the sliding window BEFORE sending, and
249
+ # here rather than in each caller. Every upstream send in the proxy
250
+ # goes through this function, so this is the one place that cannot
251
+ # be forgotten — and it was: the catch-all (which serves
252
+ # /v1/chat/completions, i.e. all opencode traffic) never recorded
253
+ # anything, leaving key_can_send_rpm() permanently True and the
254
+ # whole preventive throttle inert on the main path.
255
+ state.record_request(k)
256
+ # Claim the key BEFORE sending. get_candidate_keys() scores a key by
257
+ # in_flight + recent RPM, and neither is set until a request finishes:
258
+ # without this claim, N requests arriving together all score every key
259
+ # at zero, tie-break on index, and pile onto key[0] while the rest of
260
+ # the pool idles. That is how a 26-key pool produces 429s.
261
+ state.begin_in_flight(k)
262
+ released = False
263
+ resp = None
264
+ _model = payload.get("model", "") if isinstance(payload, dict) else ""
265
+ _t0 = time.monotonic()
266
+ # First attempt: probe timeout to detect dead models quickly.
267
+ _send_timeout = timeout
268
+ if probe_timeout is not None and attempts == 1:
269
+ _send_timeout = probe_timeout
270
+ try:
271
+ resp = await _send_once(
272
+ client, method, upstream, payload, content, hdrs, _send_timeout, stream
273
+ )
274
+ except httpx.ReadTimeout:
275
+ # The key connected and the upstream ACCEPTED the request — it is
276
+ # simply still thinking. Two things must NOT happen here:
277
+ # - cooling the key down, which blames 25 healthy keys for one
278
+ # slow model and drains the pool
279
+ # - rotating, since the next key runs the same model and will
280
+ # wait exactly as long
281
+ # Stop the phase instead and let the caller fall back to a
282
+ # different (faster) model, which is the real escalation.
283
+ state.log_cb(
284
+ f" {log_tag}: key[{idx}] no first byte in {_send_timeout.read:.0f}s "
285
+ f"— model too slow, not a key fault"
286
+ )
287
+ state.record_model_result(_model, too_slow=True)
288
+ state.record_key_model_result(k, _model, ok=False)
289
+ state.end_in_flight(k)
290
+ released = True
291
+ _record(status=504, body="model produced no first byte", deterministic=False)
292
+ break
293
+ except httpx.HTTPError as e:
294
+ err_msg = str(e) or type(e).__name__
295
+ # ConnectTimeout / ConnectError / NetworkError are NOT key faults:
296
+ # they mean the network path to NVIDIA is down or slow. Every key
297
+ # would fail identically, so cooling them down just drains the
298
+ # pool one key at a time during a transient blip. Instead, pause
299
+ # briefly and retry the SAME key rather than rotating to the next
300
+ # victim. Only after repeated network errors on the same key do
301
+ # we reluctantly cool it down (maybe its credential is bad).
302
+ is_network = isinstance(e, (httpx.ConnectTimeout, httpx.ConnectError))
303
+ if is_network:
304
+ network_errors += 1
305
+ state.log_cb(f" {log_tag}: key[{idx}] {err_msg} (network — not a key fault)")
306
+ state.end_in_flight(k)
307
+ released = True
308
+ # If 3+ keys fail to connect, the network is down — stop
309
+ # rotating (every key would fail identically) and let the
310
+ # caller return 503 instead of draining the pool.
311
+ if network_errors >= _MAX_NETWORK_ERRORS:
312
+ state.log_cb(
313
+ f" {log_tag}: {network_errors} consecutive network errors "
314
+ f"— network appears down, stopping rotation"
315
+ )
316
+ _record(status=503, body="network unreachable", deterministic=False)
317
+ return None, None, None
318
+ # Don't cool the key. Brief pause so a transient blip
319
+ # clears, then let the next candidate try.
320
+ await asyncio.sleep(0.5)
321
+ _record(status=None, body=err_msg, deterministic=False)
322
+ continue
323
+ state.log_cb(f" {log_tag}: key[{idx}] {err_msg} (rotating)")
324
+ state.end_in_flight(k)
325
+ released = True
326
+ state.mark_key_failed(k)
327
+ _record(status=None, body=err_msg, deterministic=False)
328
+ continue
329
+ finally:
330
+ # A 200 keeps the claim: the key stays busy for as long as the
331
+ # stream runs, and the caller releases it when the body ends.
332
+ if not released and (resp is None or resp.status_code != 200):
333
+ state.end_in_flight(k)
334
+ if resp.status_code == 200:
335
+ state.record_model_result(_model, ok=True, ttft=time.monotonic() - _t0)
336
+ state.record_key_model_result(k, _model, ok=True, ttft=time.monotonic() - _t0)
337
+ return resp, k, idx
338
+ err_status = resp.status_code
339
+ # Read error body for detailed logging before closing
340
+ error_body = ""
341
+ try:
342
+ error_body = await resp.aread()
343
+ error_body_bytes = error_body
344
+ error_body = error_body.decode("utf-8", errors="replace")[:500]
345
+ except Exception:
346
+ error_body_bytes = None
347
+ await resp.aclose()
348
+ state.log_cb(f" {log_tag}: key[{idx}] HTTP {err_status}")
349
+ # A full worker pool is not the model's fault any more than it is
350
+ # the key's. Deciding this BEFORE scoring matters: the exemption
351
+ # below only skipped the key cooldown, while record_model_result had
352
+ # already run — so one request during an upstream peak spent seven
353
+ # attempts, booked seven model failures against a three-failure
354
+ # threshold, and opened the model circuit for two minutes. With
355
+ # fallback=off that circuit answers every later request with an
356
+ # instant 503, which reads as the model being dead when upstream was
357
+ # merely busy.
358
+ _exhausted = err_status in (429, 503) and is_resource_exhausted(error_body_bytes)
359
+ if not _exhausted:
360
+ state.record_model_result(_model, status=err_status)
361
+ state.record_key_model_result(k, _model, ok=False)
362
+ _record(status=err_status, body=error_body, exhausted=_exhausted)
363
+
364
+ # The request is wrong, not the key. Stop immediately and hand the
365
+ # provider's own explanation back to the caller.
366
+ if err_status in _DETERMINISTIC_STATUSES:
367
+ state.log_cb(f" {log_tag}: HTTP {err_status} is a request error — not rotating")
368
+ _record(status=err_status, body=error_body, deterministic=True)
369
+ # A context-overflow rejection states the model's real window.
370
+ # Learning it here is what replaced the 1.3 MB probe: the answer
371
+ # arrives from a request the user was making anyway.
372
+ if _model and error_body:
373
+ from .compaction import note_context_limit
374
+
375
+ learned = note_context_limit(_model, error_body)
376
+ if learned:
377
+ state.log_cb(f"⧉ learned {_model} context window = {learned} tokens")
378
+ return None, None, None
379
+
380
+ # ResourceExhausted = worker concurrency limit, not RPM.
381
+ # Don't burn the key — just skip and let the next key try.
382
+ # NVIDIA sends it under 429 *and* under 503 ("Worker local total
383
+ # request limit reached (48/48)"). Matching on the status alone let
384
+ # the 503 form fall through to the gateway-timeout branch below,
385
+ # which cooled a key down for every request the shared worker pool
386
+ # refused — draining a healthy 31-key pool during an upstream peak.
387
+ if _exhausted:
388
+ exhausted_waits += 1
389
+ if exhausted_waits > _MAX_EXHAUSTED_WAITS:
390
+ state.log_cb(f" {log_tag}: upstream workers full — giving up")
391
+ return None, None, None
392
+ attempts -= 1 # Don't burn attempt budget on worker-level transient peak
393
+ state.log_cb(
394
+ f" {log_tag}: key[{idx}] ResourceExhausted (worker full) — pausing 0.8s"
395
+ )
396
+ state.stats.record_key_usage(k, ok=False, error="ResourceExhausted")
397
+ await asyncio.sleep(0.8)
398
+ continue
399
+ # A gateway timeout is the MODEL being slow, not the key being bad:
400
+ # NVIDIA's edge gave up waiting for it, and every other key would hit
401
+ # the same wall. Cooling keys down for it empties the pool one 504 at
402
+ # a time, which is precisely what a 26-key pool exists to prevent.
403
+ # Use 30s cooldown (was 10s) and don't burn the attempt budget: the
404
+ # model is the bottleneck, not the key, so trying another key just
405
+ # spreads the same 504 across the pool faster.
406
+ if err_status in _GATEWAY_TIMEOUTS:
407
+ state.log_cb(
408
+ f" {log_tag}: HTTP {err_status} gateway timeout — key[{idx}] 30s cooldown"
409
+ )
410
+ state.mark_key_failed(k, status=err_status, retry_after=30)
411
+ attempts -= 1 # Don't burn budget on a model-level timeout
412
+ continue
413
+ state.mark_key_failed(
414
+ k, status=err_status, error_body=error_body if error_body else None
415
+ )
416
+ return None, None, None
417
+
418
+
419
+ def _live_pool_snapshot(state, candidates):
420
+ """Return ``(live, total_pool)`` for saturation gating.
421
+
422
+ ``live`` = candidates that are cooldown-free AND RPM-eligible right now
423
+ (the set a rotation loop could actually succeed on). ``total_pool`` is
424
+ the FULL pool size (including cooldown keys) — weighing by the full pool
425
+ rather than just `len(candidates)` is what makes the saturation gate
426
+ fire correctly: when most of the 25 keys are on cooldown, the live
427
+ fraction genuinely drops below threshold, which is exactly when we want
428
+ to skip rotation and go to model-fallback / 503 instead of serially
429
+ hammering the few surviving keys.
430
+ """
431
+ live = sum(
432
+ 1 for _, k in candidates if state.key_can_send_rpm(k) and not state.is_key_on_cooldown(k)
433
+ )
434
+ total_pool = len(state.keys)
435
+ return live, total_pool
436
+
437
+
438
+ # ── Defensive sanitization of chat/completions messages ────────────────
439
+ def _sanitize_chat_messages(messages: list[dict]) -> list[dict]:
440
+ """
441
+ Guarantee every message is a valid OpenAI chat/completions message so
442
+ upstream never rejects a request with::
443
+
444
+ data did not match any variant of untagged enum
445
+ ChatCompletionRequestToolMessageContent
446
+
447
+ Coercion rules:
448
+ - content is forced to str. Lists are flattened to text; dicts are
449
+ JSON-encoded; None becomes "" (or " " for tool/user fallbacks).
450
+ - tool messages require a non-null string content and a tool_call_id;
451
+ if the caller forgot the id we synthesize one (NVIDIA enforces the
452
+ match against a prior assistant tool_calls entry).
453
+ - assistant messages carrying tool_calls get content coerced to ""
454
+ (never None) — NVIDIA NIM rejects content:null on non-tool messages.
455
+ - Unknown roles are dropped silently.
456
+ """
457
+ out: list[dict] = []
458
+ for m in messages:
459
+ if not isinstance(m, dict):
460
+ continue
461
+ role = m.get("role")
462
+ if role not in ("system", "user", "assistant", "tool"):
463
+ continue
464
+ content = m.get("content")
465
+ tool_calls = m.get("tool_calls")
466
+
467
+ # Flatten array-shaped content to a text string. Sources include
468
+ # tool_call_output arrays, Codex content parts, and multimodal
469
+ # payloads we reduce to plain text.
470
+ if isinstance(content, list):
471
+ parts: list[str] = []
472
+ for part in content:
473
+ if isinstance(part, str):
474
+ parts.append(part)
475
+ elif isinstance(part, dict):
476
+ if part.get("type") in ("text", "input_text", "output_text"):
477
+ parts.append(str(part.get("text", "")))
478
+ else:
479
+ parts.append(json.dumps(part, ensure_ascii=False))
480
+ content = chr(10).join(p for p in parts if p)
481
+ elif isinstance(content, dict):
482
+ content = json.dumps(content, ensure_ascii=False)
483
+ elif content is None:
484
+ content = ""
485
+ elif not isinstance(content, str):
486
+ content = str(content)
487
+
488
+ if role == "tool":
489
+ tcid = m.get("tool_call_id") or ""
490
+ if not tcid:
491
+ # No matching tool_call_id → synthesize one so the upstream
492
+ # schema validator accepts the tool response.
493
+ tcid = f"call_{uuid.uuid4().hex[:24]}"
494
+ if not content:
495
+ content = " "
496
+ out.append({"role": "tool", "tool_call_id": tcid, "content": content})
497
+ continue
498
+
499
+ if role == "assistant" and tool_calls:
500
+ # NVIDIA NIM rejects content: null on assistant messages; use "".
501
+ clean_calls = []
502
+ for tc in tool_calls:
503
+ if not isinstance(tc, dict):
504
+ continue
505
+ fn = tc.get("function") or {}
506
+ args = fn.get("arguments", "")
507
+ if not isinstance(args, str):
508
+ args = json.dumps(args, ensure_ascii=False)
509
+ if not args:
510
+ args = "{}"
511
+ clean_calls.append(
512
+ {
513
+ "id": tc.get("id") or f"call_{uuid.uuid4().hex[:24]}",
514
+ "type": "function",
515
+ "function": {
516
+ "name": fn.get("name", ""),
517
+ "arguments": args,
518
+ },
519
+ }
520
+ )
521
+ if not clean_calls:
522
+ # Malformed tool_calls: degrade to a plain assistant text turn.
523
+ out.append({"role": "assistant", "content": content or " "})
524
+ continue
525
+ out.append(
526
+ {
527
+ "role": "assistant",
528
+ "content": content or "",
529
+ "tool_calls": clean_calls,
530
+ }
531
+ )
532
+ continue
533
+
534
+ out.append({"role": role, "content": content or " "})
535
+ return out
536
+
537
+
538
+ # ── Request: Responses input → chat/completions messages ─────────────
539
+
540
+
541
+ def _input_to_messages(input_data: Any) -> list[dict]:
542
+ """
543
+ Convert the Responses API `input` field into chat/completions messages.
544
+
545
+ `input` is either a bare string (→ single user message) or an array of
546
+ InputItems. Codex CLI sends messages with role="developer" (mapped to
547
+ "system") and content parts with type="input_text" (mapped to plain text).
548
+ """
549
+ if isinstance(input_data, str):
550
+ return [{"role": "user", "content": input_data}]
551
+
552
+ messages: list[dict] = []
553
+ for item in input_data:
554
+ typ = item.get("type", "message")
555
+
556
+ if typ == "message":
557
+ role = item.get("role", "user")
558
+ # Codex uses "developer" → remap to "system" for chat/completions.
559
+ if role == "developer":
560
+ role = "system"
561
+ content = item.get("content", "")
562
+ # content may be a string or an array of content parts.
563
+ if isinstance(content, list):
564
+ # Codex uses type="input_text" on the way in and "output_text"
565
+ # for assistant turns replayed from history. Omitting the
566
+ # latter emptied every previous assistant message: in a
567
+ # multi-turn session (Codex defaults to store:false, so the
568
+ # whole history is resent each time) the model saw its own
569
+ # replies as blank and repeated work it had already done.
570
+ text_parts = [
571
+ p.get("text", "")
572
+ for p in content
573
+ if isinstance(p, dict)
574
+ and p.get("type") in ("text", "input_text", "output_text")
575
+ ]
576
+ content = "\n".join(t for t in text_parts if t)
577
+ messages.append({"role": role, "content": content})
578
+
579
+ elif typ == "function_call_output":
580
+ # Result of a prior tool call — emit as a tool message.
581
+ call_id = item.get("call_id", "")
582
+ output = item.get("output", "")
583
+ if isinstance(output, dict):
584
+ output = json.dumps(output)
585
+ messages.append({"role": "tool", "tool_call_id": call_id, "content": str(output)})
586
+
587
+ elif typ == "function_call":
588
+ # A function call from a previous turn — reconstruct as an
589
+ # assistant message with tool_calls so context is preserved.
590
+ name = item.get("name", "")
591
+ arguments = item.get("arguments", "")
592
+ call_id = item.get("call_id", "")
593
+ messages.append(
594
+ {
595
+ "role": "assistant",
596
+ "tool_calls": [
597
+ {
598
+ "id": call_id,
599
+ "type": "function",
600
+ "function": {"name": name, "arguments": arguments},
601
+ }
602
+ ],
603
+ }
604
+ )
605
+ # Unknown item types are ignored — forward compatibility.
606
+
607
+ return messages
608
+
609
+
610
+ def _tools_to_chat_tools(tools: list[dict]) -> list[dict]:
611
+ """
612
+ Convert Responses tools into chat/completions tools[].
613
+
614
+ Codex CLI sends two function shapes:
615
+ - flat: {type:"function", name:"x", description:"...", parameters:{...}}
616
+ - nested: {type:"function", function:{name:"x", ...}}
617
+
618
+ Non-function types (namespace, web_search, image_generation) are filtered
619
+ out because NVIDIA NIM does not understand them.
620
+ """
621
+ chat_tools = []
622
+ for tool in tools:
623
+ if tool.get("type") != "function":
624
+ continue
625
+ # Flat shape (Codex): name and parameters at the top level.
626
+ if "name" in tool:
627
+ chat_tools.append(
628
+ {
629
+ "type": "function",
630
+ "function": {
631
+ "name": tool.get("name", ""),
632
+ "description": tool.get("description", ""),
633
+ "parameters": tool.get("parameters", {}),
634
+ },
635
+ }
636
+ )
637
+ # Nested shape (OpenAI standard).
638
+ elif "function" in tool:
639
+ fn = tool["function"]
640
+ chat_tools.append(
641
+ {
642
+ "type": "function",
643
+ "function": {
644
+ "name": fn.get("name", ""),
645
+ "description": fn.get("description", ""),
646
+ "parameters": fn.get("parameters", {}),
647
+ },
648
+ }
649
+ )
650
+ return chat_tools
651
+
652
+
653
+ def _build_chat_payload(body: dict, model_override: str | None) -> dict:
654
+ """Build the chat/completions payload from a Responses request body."""
655
+ messages = []
656
+
657
+ # Instructions (system prompt) → leading system message.
658
+ instructions = body.get("instructions")
659
+ if instructions:
660
+ messages.append({"role": "system", "content": instructions})
661
+
662
+ messages.extend(_input_to_messages(body.get("input", "")))
663
+
664
+ # model_override (from state.active_model) takes precedence; otherwise use
665
+ # the NVIDIA default — never forward "openvidia/openvidia" to NVIDIA.
666
+ # No hardcoded model: resolved live from the user's selection.
667
+ from .proxy_app import default_model
668
+
669
+ effective_model = model_override or default_model()
670
+
671
+ payload: dict[str, Any] = {
672
+ "model": effective_model,
673
+ "messages": messages,
674
+ }
675
+
676
+ # Tools translation.
677
+ tools = body.get("tools", [])
678
+ if tools:
679
+ payload["tools"] = _tools_to_chat_tools(tools)
680
+
681
+ # tool_choice translation (Responses → chat/completions).
682
+ # Responses uses: "auto" | "none" | "required" | {type:"function", name}
683
+ tc = body.get("tool_choice")
684
+ if tc is not None:
685
+ if isinstance(tc, str):
686
+ payload["tool_choice"] = tc
687
+ elif isinstance(tc, dict):
688
+ if tc.get("type") == "auto":
689
+ payload["tool_choice"] = "auto"
690
+ elif tc.get("type") == "none":
691
+ payload["tool_choice"] = "none"
692
+ elif tc.get("type") == "required":
693
+ payload["tool_choice"] = "required"
694
+ elif tc.get("type") == "function":
695
+ payload["tool_choice"] = {
696
+ "type": "function",
697
+ "function": {"name": tc.get("name", "")},
698
+ }
699
+
700
+ # parallel_tool_calls passthrough (Codex sends it).
701
+ if "parallel_tool_calls" in body:
702
+ payload["parallel_tool_calls"] = body["parallel_tool_calls"]
703
+
704
+ # metadata passthrough (Codex may attach; NVIDIA ignores unknown fields).
705
+ if body.get("metadata"):
706
+ payload["metadata"] = body["metadata"]
707
+
708
+ # stop sequences passthrough.
709
+ if body.get("stop"):
710
+ payload["stop"] = body["stop"]
711
+
712
+ # stream_options: Codex sends {"include_usage": true} on streaming
713
+ # requests so the final SSE chunk carries token usage. NVIDIA NIM
714
+ # accepts stream_options.include_usage.
715
+ if body.get("stream_options"):
716
+ payload["stream_options"] = body["stream_options"]
717
+
718
+ # Pass-through optional parameters (only those compatible with chat/completions).
719
+ for key in (
720
+ "temperature",
721
+ "top_p",
722
+ "max_tokens",
723
+ "max_completion_tokens",
724
+ "stream",
725
+ "seed",
726
+ "presence_penalty",
727
+ "frequency_penalty",
728
+ "logit_bias",
729
+ "user",
730
+ ):
731
+ val = body.get(key)
732
+ if val is not None:
733
+ payload[key] = val
734
+
735
+ # The Responses API caps output with max_output_tokens; the loop above only
736
+ # looked for the Chat Completions spellings, so a client asking for a short
737
+ # answer got an unbounded one.
738
+ if body.get("max_output_tokens") is not None and "max_tokens" not in payload:
739
+ payload["max_tokens"] = body["max_output_tokens"]
740
+
741
+ # Structured output. Responses nests it as text.format; Chat Completions
742
+ # calls it response_format. Accept either, prefer what the client sent.
743
+ if body.get("response_format"):
744
+ payload["response_format"] = body["response_format"]
745
+ else:
746
+ text_cfg = body.get("text")
747
+ fmt = text_cfg.get("format") if isinstance(text_cfg, dict) else None
748
+ if isinstance(fmt, dict) and fmt.get("type"):
749
+ payload["response_format"] = fmt
750
+
751
+ payload["messages"] = _sanitize_chat_messages(payload["messages"])
752
+ # Translate the client's reasoning effort into NVIDIA NIM thinking flags.
753
+ # _fill_missing means an explicit client choice always beats the dashboard
754
+ # default. The Responses API spells this `reasoning: {effort: "..."}` —
755
+ # reading `reasoning_effort` (a Chat Completions field) meant the branch
756
+ # below never ran for Codex and the dashboard silently won every time.
757
+ reasoning = body.get("reasoning")
758
+ effort = reasoning.get("effort") if isinstance(reasoning, dict) else None
759
+ effort = effort or body.get("reasoning_effort")
760
+ if effort == "low":
761
+ config._fill_missing(payload, {"chat_template_kwargs": {"enable_thinking": False}})
762
+ elif effort in ("medium", "high"):
763
+ config._fill_missing(payload, {"chat_template_kwargs": {"enable_thinking": True}})
764
+ # Dashboard thinking toggle — never overrides what the client asked for.
765
+ config.apply_model_options(payload)
766
+ return payload
767
+
768
+
769
+ # ── Response: chat/completions → Responses output ────────────────────
770
+
771
+
772
+ def _chat_response_to_responses(chat_data: dict, model: str) -> dict:
773
+ """Translate a non-streaming chat/completions response into Responses format."""
774
+ resp_id = f"resp_{uuid.uuid4().hex[:24]}"
775
+ created = chat_data.get("created", int(time.time()))
776
+
777
+ output: list[dict] = []
778
+ choice = (chat_data.get("choices") or [{}])[0]
779
+ msg = choice.get("message", {})
780
+
781
+ # Tool calls → function_call items.
782
+ tool_calls = msg.get("tool_calls", [])
783
+ if tool_calls:
784
+ for tc in tool_calls:
785
+ fn = tc.get("function", {})
786
+ output.append(
787
+ {
788
+ "type": "function_call",
789
+ "id": f"fc_{uuid.uuid4().hex[:24]}",
790
+ "call_id": tc.get("id", ""),
791
+ "name": fn.get("name", ""),
792
+ "arguments": fn.get("arguments", ""),
793
+ }
794
+ )
795
+
796
+ # Text content → message item with a content array (output_text).
797
+ text = msg.get("content")
798
+ if text:
799
+ output.append(
800
+ {
801
+ "type": "message",
802
+ "id": f"msg_{uuid.uuid4().hex[:24]}",
803
+ "role": "assistant",
804
+ "status": "completed",
805
+ "content": [{"type": "output_text", "text": text}],
806
+ }
807
+ )
808
+
809
+ # Status derived from finish_reason. "tool_calls" is a completed response
810
+ # that happens to end in a tool call — marking it "incomplete" told the
811
+ # client the answer had been truncated, and a conforming one may refuse to
812
+ # run the tools. Only a real cut-off is incomplete.
813
+ finish = choice.get("finish_reason", "stop")
814
+ status = "incomplete" if finish in ("length", "content_filter") else "completed"
815
+
816
+ # Usage mapping (prompt/completion → input/output/total). NVIDIA may
817
+ # return completion_tokens_details to pass through if present.
818
+ usage = chat_data.get("usage", {})
819
+ resp_usage = {
820
+ "input_tokens": usage.get("prompt_tokens", 0),
821
+ "output_tokens": usage.get("completion_tokens", 0),
822
+ "total_tokens": usage.get("total_tokens", 0),
823
+ }
824
+ ctd = usage.get("completion_tokens_details")
825
+ if isinstance(ctd, dict) and ctd:
826
+ resp_usage["output_tokens_details"] = ctd
827
+ ptd = usage.get("prompt_tokens_details")
828
+ if isinstance(ptd, dict) and ptd:
829
+ resp_usage["input_tokens_details"] = ptd
830
+
831
+ return {
832
+ "id": resp_id,
833
+ "object": "response",
834
+ "created_at": created,
835
+ "model": model,
836
+ "output": output,
837
+ "status": status,
838
+ "usage": resp_usage,
839
+ }
840
+
841
+
842
+ # ── Streaming: SSE chat chunks → SSE Responses events ─────────────────
843
+
844
+
845
+ def _sse_event(event_type: str, data: dict) -> bytes:
846
+ """Serialize a single SSE event for the Responses streaming protocol."""
847
+ return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
848
+
849
+
850
+ async def _stream_responses(
851
+ state: ProxyState,
852
+ chat_payload: dict,
853
+ model: str,
854
+ client,
855
+ request: Request,
856
+ ) -> AsyncGenerator[bytes, None]:
857
+ """
858
+ Forward chat/completions with stream:true to NVIDIA and re-translate the
859
+ SSE chat chunks into Responses-protocol SSE events.
860
+ """
861
+ resp_id = f"resp_{uuid.uuid4().hex[:24]}"
862
+ created_ts = int(time.time())
863
+
864
+ # Opening events: response.created then response.in_progress (Codex
865
+ # expects both before any output).
866
+ yield _sse_event(
867
+ "response.created",
868
+ {
869
+ "type": "response.created",
870
+ "response": {
871
+ "id": resp_id,
872
+ "object": "response",
873
+ "created_at": created_ts,
874
+ "model": model,
875
+ "output": [],
876
+ "status": "in_progress",
877
+ },
878
+ },
879
+ )
880
+ yield _sse_event(
881
+ "response.in_progress",
882
+ {
883
+ "type": "response.in_progress",
884
+ "response": {
885
+ "id": resp_id,
886
+ "object": "response",
887
+ "created_at": created_ts,
888
+ "model": model,
889
+ "output": [],
890
+ "status": "in_progress",
891
+ },
892
+ },
893
+ )
894
+
895
+ upstream = "https://integrate.api.nvidia.com/v1/chat/completions"
896
+
897
+ # Candidate keys resolved inside the generator, not before — avoids
898
+ # concurrent requests sharing a single precomputed list.
899
+ async with state.lock:
900
+ candidates = state.get_candidate_keys()
901
+
902
+ # Key rotation: bounded attempts + saturation fast-fail (was serial
903
+ # across ALL candidates with the client default 120s read → Codex block).
904
+ resp = None
905
+ used_key = None
906
+ used_idx = None
907
+ _outcome: dict = {}
908
+
909
+ _live, _valid = _live_pool_snapshot(state, candidates)
910
+ if _valid and _live < max(1, int(_valid * _MIN_LIVE_FRACTION)):
911
+ state.log_cb(
912
+ f" responses shim: pool saturated ({_live}/{_valid} live) → skip primary "
913
+ f"(fallback/503)"
914
+ )
915
+ else:
916
+ _box: list = []
917
+ _task = asyncio.ensure_future(
918
+ _rotation_phase(
919
+ client,
920
+ upstream,
921
+ chat_payload,
922
+ json_headers,
923
+ state,
924
+ candidates,
925
+ max_attempts=_MAX_ROTATE_ATTEMPTS,
926
+ timeout=_ROTATE_SEND_TIMEOUT,
927
+ stream=True,
928
+ log_tag="responses shim",
929
+ outcome_box=_outcome,
930
+ )
931
+ )
932
+ # No deadline and no model substitution: the selected model is THE
933
+ # model. We wait for it (keepalives keep the stream alive meanwhile)
934
+ # and never silently answer from a different one.
935
+ async for _ka in _keepalive_until(_task, _box):
936
+ yield _ka
937
+ resp, used_key, used_idx = _box[0]
938
+
939
+ if resp is None or used_key is None:
940
+ # Definitive failure: emit error + response.failed and terminate.
941
+ yield _sse_event(
942
+ "error",
943
+ {
944
+ "type": "error",
945
+ "code": "invalid_request_error"
946
+ if _outcome.get("deterministic")
947
+ else "server_error",
948
+ "message": _upstream_error_message(_outcome),
949
+ "param": None,
950
+ },
951
+ )
952
+ yield _sse_event(
953
+ "response.failed",
954
+ {
955
+ "type": "response.failed",
956
+ "response": {
957
+ "id": resp_id,
958
+ "object": "response",
959
+ "created_at": created_ts,
960
+ "model": model,
961
+ "output": [],
962
+ "status": "failed",
963
+ },
964
+ },
965
+ )
966
+ return
967
+
968
+ state.stats.success += 1
969
+ state.stats.record_key_usage(used_key, ok=True)
970
+ if used_idx is not None:
971
+ state.log_cb(f"✔ key[{used_idx}] OK")
972
+
973
+ # Accumulators for the streamed output. Codex (Responses API) uses item
974
+ # type "message" with a content array of output_text parts — NOT item
975
+ # type "text".
976
+ text_item_id = f"msg_{uuid.uuid4().hex[:24]}"
977
+ text_started = False
978
+ text_full = ""
979
+ text_output_index = 0
980
+ # output_index: 0 = text (if any), then tool_calls in arrival order.
981
+ tool_calls_map: dict[int, dict] = {} # index → tool_call state
982
+ next_output_index = 0
983
+ items_done_emitted = False
984
+
985
+ def _emit_items_done() -> list[bytes]:
986
+ """Emit output_text.done + output_item.done for text and all tool calls.
987
+ Returns the list of SSE events to yield."""
988
+ events: list[bytes] = []
989
+ if text_started:
990
+ events.append(
991
+ _sse_event(
992
+ "response.output_text.done",
993
+ {
994
+ "type": "response.output_text.done",
995
+ "item_id": text_item_id,
996
+ "output_index": text_output_index,
997
+ "text": text_full,
998
+ },
999
+ )
1000
+ )
1001
+ events.append(
1002
+ _sse_event(
1003
+ "response.output_item.done",
1004
+ {
1005
+ "type": "response.output_item.done",
1006
+ "output_index": text_output_index,
1007
+ "item": {
1008
+ "type": "message",
1009
+ "id": text_item_id,
1010
+ "role": "assistant",
1011
+ "status": "completed",
1012
+ "content": [{"type": "output_text", "text": text_full}],
1013
+ },
1014
+ },
1015
+ )
1016
+ )
1017
+ for tcm in tool_calls_map.values():
1018
+ events.append(
1019
+ _sse_event(
1020
+ "response.function_call_arguments.done",
1021
+ {
1022
+ "type": "response.function_call_arguments.done",
1023
+ "item_id": tcm["fc_id"],
1024
+ "output_index": tcm["output_index"],
1025
+ "arguments": tcm["arguments"],
1026
+ },
1027
+ )
1028
+ )
1029
+ events.append(
1030
+ _sse_event(
1031
+ "response.output_item.done",
1032
+ {
1033
+ "type": "response.output_item.done",
1034
+ "output_index": tcm["output_index"],
1035
+ "item": {
1036
+ "type": "function_call",
1037
+ "id": tcm["fc_id"],
1038
+ "call_id": tcm["call_id"],
1039
+ "name": tcm["name"],
1040
+ "arguments": tcm["arguments"],
1041
+ "status": "completed",
1042
+ },
1043
+ },
1044
+ )
1045
+ )
1046
+ return events
1047
+
1048
+ # Usage arrives in its own trailing chunk (choices == []) AFTER the one
1049
+ # carrying finish_reason. Reading it from a single `chunk` variable at the
1050
+ # end of the loop found the finish chunk instead, so usage was always
1051
+ # zeros — and when the loop never iterated (an immediate [DONE], a client
1052
+ # that disconnected on the first line, an empty body) that variable was
1053
+ # never assigned at all, raising UnboundLocalError inside the generator
1054
+ # after the 200 headers had shipped: a truncated stream with no
1055
+ # response.completed. Accumulate it as it arrives instead.
1056
+ final_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
1057
+
1058
+ try:
1059
+ async for line in resp.aiter_lines():
1060
+ if await request.is_disconnected():
1061
+ break
1062
+ if not line or not line.startswith("data: "):
1063
+ continue
1064
+ data_str = line[6:]
1065
+ if data_str.strip() == "[DONE]":
1066
+ break
1067
+
1068
+ try:
1069
+ chunk = json.loads(data_str)
1070
+ except json.JSONDecodeError:
1071
+ continue
1072
+
1073
+ usage = chunk.get("usage")
1074
+ if isinstance(usage, dict) and usage:
1075
+ final_usage = {
1076
+ "input_tokens": usage.get("prompt_tokens", 0),
1077
+ "output_tokens": usage.get("completion_tokens", 0),
1078
+ "total_tokens": usage.get("total_tokens", 0),
1079
+ }
1080
+
1081
+ choice = (chunk.get("choices") or [{}])[0]
1082
+ delta = choice.get("delta", {})
1083
+
1084
+ # Text delta → response.output_text.delta
1085
+ text_content = delta.get("content")
1086
+ if text_content:
1087
+ if not text_started:
1088
+ text_started = True
1089
+ text_output_index = next_output_index
1090
+ next_output_index += 1
1091
+ yield _sse_event(
1092
+ "response.output_item.added",
1093
+ {
1094
+ "type": "response.output_item.added",
1095
+ "output_index": text_output_index,
1096
+ "item": {
1097
+ "type": "message",
1098
+ "id": text_item_id,
1099
+ "role": "assistant",
1100
+ "status": "in_progress",
1101
+ "content": [],
1102
+ },
1103
+ },
1104
+ )
1105
+ yield _sse_event(
1106
+ "response.output_text.delta",
1107
+ {
1108
+ "type": "response.output_text.delta",
1109
+ "item_id": text_item_id,
1110
+ "output_index": text_output_index,
1111
+ "delta": text_content,
1112
+ },
1113
+ )
1114
+ text_full += text_content
1115
+
1116
+ # Tool call deltas → function_call item add + arguments deltas.
1117
+ tc_deltas = delta.get("tool_calls", [])
1118
+ for tc in tc_deltas:
1119
+ idx = tc.get("index", 0)
1120
+ if idx not in tool_calls_map:
1121
+ fc_id = f"fc_{uuid.uuid4().hex[:24]}"
1122
+ call_id = tc.get("id", f"call_{uuid.uuid4().hex[:24]}")
1123
+ tc_output_index = next_output_index
1124
+ next_output_index += 1
1125
+ tool_calls_map[idx] = {
1126
+ "fc_id": fc_id,
1127
+ "call_id": call_id,
1128
+ "name": tc.get("function", {}).get("name", ""),
1129
+ "arguments": "",
1130
+ "output_index": tc_output_index,
1131
+ }
1132
+ yield _sse_event(
1133
+ "response.output_item.added",
1134
+ {
1135
+ "type": "response.output_item.added",
1136
+ "output_index": tc_output_index,
1137
+ "item": {
1138
+ "type": "function_call",
1139
+ "id": fc_id,
1140
+ "call_id": call_id,
1141
+ "name": tc.get("function", {}).get("name", ""),
1142
+ "arguments": "",
1143
+ "status": "in_progress",
1144
+ },
1145
+ },
1146
+ )
1147
+ tcm = tool_calls_map[idx]
1148
+ args_delta = tc.get("function", {}).get("arguments", "")
1149
+ if args_delta:
1150
+ tcm["arguments"] += args_delta
1151
+ yield _sse_event(
1152
+ "response.function_call_arguments.delta",
1153
+ {
1154
+ "type": "response.function_call_arguments.delta",
1155
+ "item_id": tcm["fc_id"],
1156
+ "output_index": tcm["output_index"],
1157
+ "delta": args_delta,
1158
+ },
1159
+ )
1160
+
1161
+ # Finish: close all open items, then emit response.completed.
1162
+ finish = choice.get("finish_reason")
1163
+ if finish and not items_done_emitted:
1164
+ items_done_emitted = True
1165
+ for ev in _emit_items_done():
1166
+ yield ev
1167
+ # Deliberately no break: the items are closed, but the usage
1168
+ # chunk still follows. The loop ends on [DONE] a moment later.
1169
+ continue
1170
+
1171
+ # Safety catch: ensure items are closed even if finish_reason was omitted or [DONE] arrived early
1172
+ if not items_done_emitted:
1173
+ items_done_emitted = True
1174
+ for ev in _emit_items_done():
1175
+ yield ev
1176
+
1177
+ # Final response.completed event. A non-empty output array lets Codex
1178
+ # reconstruct the assistant message instead of relying on delta replay.
1179
+ final_output: list[dict] = []
1180
+ if text_started:
1181
+ final_output.append(
1182
+ {
1183
+ "type": "message",
1184
+ "id": text_item_id,
1185
+ "role": "assistant",
1186
+ "status": "completed",
1187
+ "content": [{"type": "output_text", "text": text_full}],
1188
+ }
1189
+ )
1190
+ for tcm in tool_calls_map.values():
1191
+ final_output.append(
1192
+ {
1193
+ "type": "function_call",
1194
+ "id": tcm["fc_id"],
1195
+ "call_id": tcm["call_id"],
1196
+ "name": tcm["name"],
1197
+ "arguments": tcm["arguments"],
1198
+ "status": "completed",
1199
+ }
1200
+ )
1201
+ yield _sse_event(
1202
+ "response.completed",
1203
+ {
1204
+ "type": "response.completed",
1205
+ "response": {
1206
+ "id": resp_id,
1207
+ "object": "response",
1208
+ "created_at": created_ts,
1209
+ "model": model,
1210
+ "output": final_output,
1211
+ "status": "completed",
1212
+ "usage": final_usage,
1213
+ },
1214
+ },
1215
+ )
1216
+ finally:
1217
+ # Runs on client disconnect (GeneratorExit) and on any httpx error
1218
+ # while reading the body. Previously both skipped the release, and
1219
+ # the key kept its in-flight claim for the life of the process.
1220
+ await resp.aclose()
1221
+ # Release the load-balancer claim taken in _rotation_phase. Without
1222
+ # this the key looks permanently busy and the scheduler stops
1223
+ # choosing it.
1224
+ state.end_in_flight(used_key)
1225
+
1226
+
1227
+ # ── Entry point ─────────────────────────────────────────────────────────
1228
+
1229
+
1230
+ async def handle_responses(
1231
+ request: Request,
1232
+ state: ProxyState,
1233
+ client,
1234
+ ) -> JSONResponse | StreamingResponse:
1235
+ """
1236
+ Handle POST /v1/responses.
1237
+
1238
+ Forwards directly to NVIDIA chat/completions (this path is internal — it
1239
+ does not go through the catch-all proxy_app handler).
1240
+ """
1241
+ raw = await request.body()
1242
+ try:
1243
+ body = json.loads(raw)
1244
+ except json.JSONDecodeError:
1245
+ return JSONResponse({"error": "invalid JSON"}, status_code=400)
1246
+
1247
+ state.log_cb(
1248
+ f" responses: model={body.get('model', '?')} stream={body.get('stream', False)} tools={len(body.get('tools', []))}"
1249
+ )
1250
+
1251
+ model_override = state.active_model
1252
+ chat_payload = _build_chat_payload(body, model_override)
1253
+
1254
+ # ── Circuit breaker: skip dead models, auto-failover to next preset ───
1255
+ _req_model = chat_payload.get("model", "")
1256
+ if _req_model and state.is_model_circuit_open(_req_model):
1257
+ # Respect user's fallback preference: "off" = never failover, return 503
1258
+ _opts = config.model_options()
1259
+ _fallback_mode = _opts.get("fallback", "off")
1260
+ _per = (_opts.get("per_model") or {}).get(_req_model, {})
1261
+ _fallback_mode = _per.get("fallback", _fallback_mode)
1262
+ if _fallback_mode == "off":
1263
+ state.log_cb(f"🔴 {_req_model} circuit OPEN, fallback=off → 503")
1264
+ return JSONResponse(
1265
+ {"error": f"{_req_model} is down (circuit open), fallback disabled"},
1266
+ status_code=503,
1267
+ )
1268
+ _presets = config.load_saved_presets()
1269
+ _fallback = next(
1270
+ (m for m in _presets if m != _req_model and not state.is_model_circuit_open(m)),
1271
+ None,
1272
+ )
1273
+ if _fallback:
1274
+ state.log_cb(f"🔴 {_req_model} circuit OPEN → auto-failover to {_fallback}")
1275
+ chat_payload["model"] = _fallback
1276
+ config.apply_model_options(chat_payload) # re-apply opts for new model
1277
+ else:
1278
+ state.log_cb(f"🔴 {_req_model} circuit OPEN, no healthy fallback")
1279
+ return JSONResponse(
1280
+ {"error": f"{_req_model} is down (circuit open), no fallback available"},
1281
+ status_code=503,
1282
+ )
1283
+
1284
+ # Auto-compaction: summarize oversized history in place (trim fallback).
1285
+ from .compaction import maybe_compact
1286
+
1287
+ chat_payload["messages"] = await maybe_compact(
1288
+ chat_payload["messages"],
1289
+ state=state,
1290
+ client=client,
1291
+ log=state.log_cb,
1292
+ model=chat_payload.get("model", ""),
1293
+ )
1294
+
1295
+ want_stream = body.get("stream", False)
1296
+
1297
+ if want_stream:
1298
+ chat_payload["stream"] = True
1299
+ return StreamingResponse(
1300
+ _stream_responses(
1301
+ state,
1302
+ chat_payload,
1303
+ model_override or chat_payload["model"],
1304
+ client,
1305
+ request,
1306
+ ),
1307
+ media_type="text/event-stream",
1308
+ headers={
1309
+ "Cache-Control": "no-cache",
1310
+ "Connection": "keep-alive",
1311
+ "X-Accel-Buffering": "no",
1312
+ },
1313
+ )
1314
+
1315
+ # Non-streaming: resolve candidates right before the HTTP loop.
1316
+ async with state.lock:
1317
+ candidates = state.get_candidate_keys()
1318
+ if not candidates:
1319
+ return JSONResponse({"error": "no keys available"}, status_code=503)
1320
+
1321
+ upstream = "https://integrate.api.nvidia.com/v1/chat/completions"
1322
+ used_key = None
1323
+ used_idx = None
1324
+ resp = None
1325
+ _outcome: dict = {}
1326
+
1327
+ # Bounded attempts + saturation fast-fail (was serial across all
1328
+ # candidates with the 120s client default → Codex block).
1329
+ _live, _valid = _live_pool_snapshot(state, candidates)
1330
+ if _valid and _live < max(1, int(_valid * _MIN_LIVE_FRACTION)):
1331
+ state.log_cb(f" responses shim: pool saturated ({_live}/{_valid} live) → 503 fast")
1332
+ else:
1333
+ resp, used_key, used_idx = await _rotation_phase(
1334
+ client,
1335
+ upstream,
1336
+ chat_payload,
1337
+ json_headers,
1338
+ state,
1339
+ candidates,
1340
+ max_attempts=_MAX_ROTATE_ATTEMPTS,
1341
+ timeout=_ROTATE_SEND_TIMEOUT,
1342
+ stream=False,
1343
+ log_tag="responses shim",
1344
+ outcome_box=_outcome,
1345
+ )
1346
+ if used_idx is not None and isinstance(resp, httpx.Response):
1347
+ await resp.aread()
1348
+
1349
+ if resp is None or used_key is None:
1350
+ # A deterministic upstream rejection is the client's problem to fix, so
1351
+ # it keeps its own status instead of being flattened into a 503 that
1352
+ # invites the client to retry something that will never succeed.
1353
+ if _outcome.get("deterministic"):
1354
+ return JSONResponse(
1355
+ {"error": _upstream_error_message(_outcome)},
1356
+ status_code=_outcome.get("status") or 400,
1357
+ )
1358
+ base = "all keys failed (pool saturated)" if _live == 0 else "all keys failed"
1359
+ return JSONResponse(
1360
+ {"error": _upstream_error_message(_outcome, base)},
1361
+ status_code=503,
1362
+ )
1363
+
1364
+ state.stats.success += 1
1365
+ state.stats.record_key_usage(used_key, ok=True)
1366
+ if used_idx is not None:
1367
+ state.log_cb(f"✔ key[{used_idx}] OK")
1368
+
1369
+ chat_data = resp.json()
1370
+ await resp.aclose()
1371
+ state.end_in_flight(used_key) # release the load-balancer claim
1372
+
1373
+ responses_data = _chat_response_to_responses(chat_data, model_override or chat_payload["model"])
1374
+ return JSONResponse(responses_data)