fix(stream_diag): feed the existing chunk-body upstream_provider into diag + hook payload (#90216)
Drop the duplicate chunk-reading helper, run_agent facade forward and _last_serving_provider agent state; the chat-completions loop already captures chunk.provider, so stamp it on the per-attempt diag there and read the hook's upstream_provider from the assembled response.provider.
This commit is contained in:
@@ -148,6 +148,8 @@ class ApiRequestHooksMixin:
|
||||
return self._sanitize_hook_payload(
|
||||
{
|
||||
"model": getattr(response, "model", None),
|
||||
# Downstream that served the call (relays re-roll it per request; #90216).
|
||||
"upstream_provider": getattr(response, "provider", None),
|
||||
"finish_reason": finish_reason,
|
||||
"assistant_message": {
|
||||
"role": getattr(assistant_message, "role", "assistant"),
|
||||
|
||||
@@ -2884,9 +2884,6 @@ class _StreamingCall(StreamingWaitMonitor):
|
||||
diag["first_chunk_at"] = self.last_chunk_time["t"]
|
||||
# Delta-length estimate: ~3x cheaper than repr() per chunk.
|
||||
diag["bytes"] = int(diag.get("bytes", 0)) + _estimate_chunk_bytes(chunk)
|
||||
# Relays re-roll the serving provider per request and report it only in the chunk body,
|
||||
# so attribute a drop to the downstream that served it, not to the aggregator (#90216).
|
||||
self.agent._stream_diag_note_serving_provider(diag, chunk)
|
||||
|
||||
# ── chat_completions wire ───────────────────────────────────────────
|
||||
|
||||
@@ -3078,6 +3075,7 @@ class _StreamingCall(StreamingWaitMonitor):
|
||||
response_id = chunk.id
|
||||
if upstream_provider is None and isinstance(getattr(chunk, "provider", None), str) and chunk.provider:
|
||||
upstream_provider = chunk.provider # OpenRouter stamps who served
|
||||
_diag["serving_provider"] = upstream_provider.strip()[:64] # attribute a mid-stream drop (#90216)
|
||||
if not chunk.choices:
|
||||
usage, finish_reason = self._choiceless_chunk(chunk, finish_reason)
|
||||
usage_obj = usage or usage_obj
|
||||
@@ -3779,10 +3777,6 @@ class _StreamingCall(StreamingWaitMonitor):
|
||||
# Propagate first-chunk timing for the ``post_api_request`` hook.
|
||||
if isinstance(self.clients.diag, dict) and self.clients.diag.get("first_chunk_at"):
|
||||
self.agent._last_api_first_chunk_at = float(self.clients.diag["first_chunk_at"])
|
||||
# Same per-attempt stash for the downstream that actually served the stream: relays re-roll it
|
||||
# per request, and plugins auditing route compliance cannot get it from anywhere else (#90216).
|
||||
if isinstance(self.clients.diag, dict) and self.clients.diag.get("serving_provider"):
|
||||
self.agent._last_serving_provider = str(self.clients.diag["serving_provider"])
|
||||
return self.result["response"]
|
||||
|
||||
|
||||
|
||||
@@ -30,34 +30,6 @@ def stream_diag_init() -> Dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def stream_diag_note_serving_provider(diag: Dict[str, Any], chunk: Any) -> None:
|
||||
"""Record which downstream provider actually served this attempt, from a delta chunk body.
|
||||
|
||||
On OpenRouter-style relays the provider is re-rolled per request and reported only inside the
|
||||
chunk JSON (``"provider": "Novita"``) — those responses carry no ``x-openrouter-provider``
|
||||
header, just ``cf-ray`` / ``server: cloudflare``, so the header snapshot cannot attribute a
|
||||
mid-stream drop to a downstream. First non-empty value wins: the attempt was served by whatever
|
||||
produced its first chunk, and later chunks must not overwrite that. Best-effort, never raises.
|
||||
"""
|
||||
if not isinstance(diag, dict) or diag.get("serving_provider"):
|
||||
return
|
||||
try:
|
||||
value = getattr(chunk, "provider", None)
|
||||
if value is None: # unknown top-level fields land in pydantic model_extra on the OpenAI SDK
|
||||
extra = getattr(chunk, "model_extra", None)
|
||||
if isinstance(extra, dict):
|
||||
value = extra.get("provider")
|
||||
if value is None and isinstance(chunk, dict):
|
||||
value = chunk.get("provider")
|
||||
if isinstance(value, str) and value.strip():
|
||||
diag["serving_provider"] = value.strip()[:64] # keep log lines bounded
|
||||
except Exception:
|
||||
# Deliberately swallowed: this is a best-effort annotation on the diagnostics path and
|
||||
# must never break streaming. Logged at DEBUG so a reader can tell it apart from a
|
||||
# missed error-handling gap.
|
||||
logger.debug("stream_diag: could not read serving provider from chunk", exc_info=True)
|
||||
|
||||
|
||||
def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response: Any) -> None:
|
||||
"""Snapshot headers + HTTP status at stream open (so they survive a drop before the first chunk). Best-effort."""
|
||||
if http_response is None or not isinstance(diag, dict):
|
||||
@@ -256,7 +228,6 @@ __all__ = [
|
||||
"connect_exhausted_notice",
|
||||
"buffer_connect_exhausted_notice",
|
||||
"stream_diag_init",
|
||||
"stream_diag_note_serving_provider",
|
||||
"stream_diag_capture_response",
|
||||
"flatten_exception_chain",
|
||||
"log_stream_retry",
|
||||
|
||||
@@ -104,8 +104,6 @@ def build_api_request(
|
||||
agent._reset_stream_delivery_tracking()
|
||||
# Per-attempt first-chunk timestamp so a stale value never leaks into post_api_request.
|
||||
agent._last_api_first_chunk_at = None
|
||||
# Same for the downstream provider that served the attempt (relays re-roll it per request).
|
||||
agent._last_serving_provider = None
|
||||
# api_messages was built for the primary; a fallback (DeepSeek / Kimi / MiMo) may
|
||||
# require reasoning_content — re-apply the echo-back pad (idempotent) and re-render
|
||||
# the prompt-cache decoration for the current provider.
|
||||
|
||||
@@ -83,9 +83,6 @@ def _fire_post_api_request_hook(
|
||||
# First stream chunk time (epoch s); None if not streamed / no chunk.
|
||||
# TTFB = first_chunk_at - started_at.
|
||||
first_chunk_at=getattr(agent, "_last_api_first_chunk_at", None),
|
||||
# Downstream that actually served the stream ("Novita"): relays re-roll it per
|
||||
# request and report it only in the chunk body, never in a response header.
|
||||
upstream_provider=getattr(agent, "_last_serving_provider", None),
|
||||
finish_reason=finish_reason,
|
||||
message_count=len(api_messages),
|
||||
response_model=getattr(response, "model", None),
|
||||
|
||||
@@ -144,7 +144,6 @@ _DEFAULT_PAYLOADS = {
|
||||
"base_url": "https://api.anthropic.com", "api_mode": "anthropic_messages",
|
||||
"api_call_count": 1, "api_duration": 1.234,
|
||||
"started_at": 1756000000.0, "ended_at": 1756000001.234, "first_chunk_at": 1756000000.512,
|
||||
"upstream_provider": "Novita",
|
||||
"finish_reason": "stop", "message_count": 4, "response_model": "claude-sonnet-4-6",
|
||||
"usage": {"input_tokens": 2048, "output_tokens": 512},
|
||||
"assistant_content_chars": 1200, "assistant_tool_call_count": 0,
|
||||
|
||||
@@ -509,7 +509,6 @@ class AIAgent(
|
||||
return {"messages": stripped_messages, "items": stripped_items}
|
||||
|
||||
_stream_diag_init = _forward_static("agent.stream_diag", "stream_diag_init")
|
||||
_stream_diag_note_serving_provider = _forward_static("agent.stream_diag", "stream_diag_note_serving_provider")
|
||||
_stream_diag_capture_response = _forward("agent.stream_diag", "stream_diag_capture_response")
|
||||
_flatten_exception_chain = _forward_static("agent.stream_diag", "flatten_exception_chain")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user