fix(agent): record the serving downstream provider in stream-drop diagnostics
Relay routing is re-rolled per request and the winning downstream is reported only inside the delta chunk bodies, so the header snapshot in agent/stream_diag.py could not attribute a mid-stream drop to a provider. The per-attempt diag now carries serving_provider (first non-empty chunk-body provider), log_stream_retry prints it, and the post_api_request payload exposes it as upstream_provider for plugins auditing route compliance. Fixes #90216. (cherry picked from commit f0cf4fe4f4b47e384532008b1bf56ad259c87dc2)
This commit is contained in:
@@ -2884,6 +2884,9 @@ 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 ───────────────────────────────────────────
|
||||
|
||||
@@ -3776,6 +3779,10 @@ 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"]
|
||||
|
||||
|
||||
|
||||
@@ -24,7 +24,35 @@ STREAM_DIAG_HEADERS = (
|
||||
|
||||
def stream_diag_init() -> Dict[str, Any]:
|
||||
"""Fresh per-attempt diagnostic dict; mutated in place by the streaming functions and read by the retry block."""
|
||||
return {"started_at": time.time(), "first_chunk_at": None, "chunks": 0, "bytes": 0, "headers": {}, "http_status": None}
|
||||
return {
|
||||
"started_at": time.time(), "first_chunk_at": None, "chunks": 0, "bytes": 0,
|
||||
"headers": {}, "http_status": None, "serving_provider": None,
|
||||
}
|
||||
|
||||
|
||||
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:
|
||||
pass
|
||||
|
||||
|
||||
def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response: Any) -> None:
|
||||
@@ -71,10 +99,11 @@ def flatten_exception_chain(error: BaseException) -> str:
|
||||
|
||||
|
||||
def _diag_fields(diag: Optional[Dict[str, Any]]) -> tuple:
|
||||
"""(http_status, bytes, chunks, elapsed, ttfb, upstream) for the retry log line; ``-`` when unknown."""
|
||||
"""(http_status, bytes, chunks, elapsed, ttfb, serving_provider, upstream) for the retry log line;
|
||||
``-`` when unknown."""
|
||||
_bytes = _chunks = 0
|
||||
_elapsed = 0.0
|
||||
_ttfb = _headers_repr = _http_status = "-"
|
||||
_ttfb = _headers_repr = _http_status = _serving_provider = "-"
|
||||
if isinstance(diag, dict):
|
||||
try:
|
||||
_now = time.time()
|
||||
@@ -90,9 +119,11 @@ def _diag_fields(diag: Optional[Dict[str, Any]]) -> tuple:
|
||||
_headers_repr = " ".join(f"{k}={v}" for k, v in headers.items())
|
||||
if diag.get("http_status") is not None:
|
||||
_http_status = str(diag.get("http_status"))
|
||||
if diag.get("serving_provider"):
|
||||
_serving_provider = str(diag["serving_provider"])
|
||||
except Exception:
|
||||
pass
|
||||
return _http_status, _bytes, _chunks, _elapsed, _ttfb, _headers_repr
|
||||
return _http_status, _bytes, _chunks, _elapsed, _ttfb, _serving_provider, _headers_repr
|
||||
|
||||
|
||||
def log_stream_retry(
|
||||
@@ -101,7 +132,11 @@ def log_stream_retry(
|
||||
) -> None:
|
||||
"""Structured WARNING to ``agent.log`` for a transient stream drop + retry, always logged regardless of
|
||||
UI verbosity. With *diag*, also records upstream headers, HTTP status, bytes/chunks, elapsed and TTFB on
|
||||
the dying attempt — enough to tell "one CF edge / downstream provider" from "random across runs"."""
|
||||
the dying attempt — enough to tell "one CF edge / downstream provider" from "random across runs".
|
||||
|
||||
The ``serving_provider`` field names the downstream that actually served the attempt (relays re-roll it
|
||||
per request and report it only in the chunk body), so a drop can be attributed to a provider even when
|
||||
the response carried no ``x-openrouter-provider`` header."""
|
||||
try:
|
||||
try:
|
||||
_summary = agent._summarize_api_error(error)
|
||||
@@ -116,7 +151,8 @@ def log_stream_retry(
|
||||
|
||||
logger.warning(
|
||||
"Stream %s on attempt %s/%s — retrying. subagent_id=%s depth=%s provider=%s base_url=%s "
|
||||
"error_type=%s error=%s chain=%s http_status=%s bytes=%d chunks=%d elapsed=%.2fs ttfb=%s upstream=[%s]",
|
||||
"error_type=%s error=%s chain=%s http_status=%s bytes=%d chunks=%d elapsed=%.2fs ttfb=%s "
|
||||
"serving_provider=%s upstream=[%s]",
|
||||
kind, attempt, max_attempts,
|
||||
getattr(agent, "_subagent_id", None) or "-", getattr(agent, "_delegate_depth", 0),
|
||||
agent.provider or "-", agent.base_url or "-",
|
||||
@@ -217,6 +253,7 @@ __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,6 +104,8 @@ 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,6 +83,9 @@ 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,6 +144,7 @@ _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,6 +509,7 @@ 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")
|
||||
|
||||
|
||||
154
tests/agent/test_stream_serving_provider.py
Normal file
154
tests/agent/test_stream_serving_provider.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""Regression tests for the serving-provider field in stream diagnostics (#90216).
|
||||
|
||||
OpenRouter-style relays re-roll the downstream provider per request — three consecutive calls
|
||||
for one model came back served by Alibaba, Novita and StreamLake — and report the winner only
|
||||
inside the delta chunk bodies (``{"provider": "Novita", ...}``). Those responses carry no
|
||||
``x-openrouter-provider`` header (just ``cf-ray`` / ``server: cloudflare``), so the existing
|
||||
header snapshot cannot attribute a mid-stream drop to a provider at all.
|
||||
|
||||
Contract under test:
|
||||
|
||||
- The first non-empty top-level ``provider`` in a chunk body lands in the per-attempt diag dict
|
||||
as ``serving_provider``; later chunks must not rewrite that attribution.
|
||||
- ``log_stream_retry`` names it on the drop line (``serving_provider=-`` when unknown).
|
||||
- A successful streamed attempt stashes it on the agent and the ``post_api_request`` hook
|
||||
payload carries it as ``upstream_provider`` — refreshed per attempt, never leaked from a
|
||||
prior call.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from run_agent import AIAgent
|
||||
from tests.agent.test_run_agent import ( # noqa: F401 (_make_tool_defs used by the agent fixture)
|
||||
_make_tool_defs,
|
||||
_mock_response,
|
||||
)
|
||||
from tests.agent.test_first_chunk_at_hook import ( # noqa: F401 (shared fixture + harness)
|
||||
_make_stream_chunk,
|
||||
_run_with_hooks,
|
||||
agent,
|
||||
)
|
||||
|
||||
|
||||
def _make_agent() -> AIAgent:
|
||||
"""Standalone agent for the diag/log-level tests (no conversation loop involved)."""
|
||||
return AIAgent(
|
||||
api_key="test-key",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
quiet_mode=True,
|
||||
skip_context_files=True,
|
||||
skip_memory=True,
|
||||
)
|
||||
|
||||
|
||||
def _chunk_with_provider(content=None, finish_reason=None, model=None, provider=None):
|
||||
"""Streaming chunk plus the relay's per-chunk ``provider`` field."""
|
||||
chunk = _make_stream_chunk(content=content, finish_reason=finish_reason, model=model)
|
||||
if provider is not None:
|
||||
chunk.provider = provider
|
||||
return chunk
|
||||
|
||||
|
||||
# ── Diag dict level ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_diag_records_first_serving_provider_from_chunk_body():
|
||||
agent_ = _make_agent()
|
||||
diag = AIAgent._stream_diag_init()
|
||||
assert diag["serving_provider"] is None
|
||||
|
||||
agent_._stream_diag_note_serving_provider(diag, _chunk_with_provider(content="Hi", provider="Novita"))
|
||||
assert diag["serving_provider"] == "Novita"
|
||||
|
||||
# The attempt was served by whatever produced its first chunk; a per-chunk rewrite would
|
||||
# just track the last one and lose the drop attribution.
|
||||
agent_._stream_diag_note_serving_provider(diag, _chunk_with_provider(content="!", provider="StreamLake"))
|
||||
assert diag["serving_provider"] == "Novita"
|
||||
|
||||
|
||||
def test_diag_reads_provider_from_model_extra_and_ignores_chunks_without_one():
|
||||
"""The OpenAI SDK parks unknown top-level fields in pydantic ``model_extra``."""
|
||||
agent_ = _make_agent()
|
||||
diag = AIAgent._stream_diag_init()
|
||||
|
||||
agent_._stream_diag_note_serving_provider(diag, _chunk_with_provider(content="plain"))
|
||||
assert diag["serving_provider"] is None
|
||||
|
||||
agent_._stream_diag_note_serving_provider(diag, SimpleNamespace(model_extra={"provider": "Alibaba"}))
|
||||
assert diag["serving_provider"] == "Alibaba"
|
||||
|
||||
# Empty / non-string values are not a serving provider either.
|
||||
other = AIAgent._stream_diag_init()
|
||||
agent_._stream_diag_note_serving_provider(other, SimpleNamespace(provider=" "))
|
||||
assert other["serving_provider"] is None
|
||||
|
||||
|
||||
def test_log_stream_retry_names_the_serving_provider(caplog):
|
||||
agent_ = _make_agent()
|
||||
agent_.provider = "openrouter"
|
||||
|
||||
diag = AIAgent._stream_diag_init()
|
||||
diag["serving_provider"] = "Novita"
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
agent_._log_stream_retry(
|
||||
kind="drop", error=ConnectionError("peer closed"), attempt=2, max_attempts=3,
|
||||
mid_tool_call=False, diag=diag,
|
||||
)
|
||||
|
||||
msg = next(r.getMessage() for r in caplog.records if "Stream drop" in r.getMessage())
|
||||
assert "serving_provider=Novita" in msg
|
||||
|
||||
caplog.clear()
|
||||
unknown = AIAgent._stream_diag_init()
|
||||
with caplog.at_level(logging.WARNING):
|
||||
agent_._log_stream_retry(
|
||||
kind="drop", error=ConnectionError("peer closed"), attempt=2, max_attempts=3,
|
||||
mid_tool_call=False, diag=unknown,
|
||||
)
|
||||
msg = next(r.getMessage() for r in caplog.records if "Stream drop" in r.getMessage())
|
||||
assert "serving_provider=-" in msg
|
||||
|
||||
|
||||
# ── Conversation loop / hook payload level ───────────────────────────────
|
||||
|
||||
|
||||
class TestServingProviderReachesPostApiRequest:
|
||||
"""``upstream_provider`` on the post_api_request payload (plugin route auditing)."""
|
||||
|
||||
@patch("run_agent.AIAgent._create_request_openai_client")
|
||||
@patch("run_agent.AIAgent._close_request_openai_client")
|
||||
def test_streamed_attempt_reports_the_serving_provider(self, _mock_close, mock_create, agent):
|
||||
chunks = [
|
||||
_chunk_with_provider(content="Hello", provider="Novita"),
|
||||
_chunk_with_provider(content=" world", finish_reason="stop", model="test-model", provider="Novita"),
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat.completions.create.return_value = iter(chunks)
|
||||
mock_create.return_value = mock_client
|
||||
# A registered stream consumer forces the streaming path even though agent.client is a Mock.
|
||||
agent.stream_delta_callback = lambda _text: None
|
||||
|
||||
result, post = _run_with_hooks(agent)
|
||||
|
||||
assert result["final_response"] == "Hello world"
|
||||
assert len(post) == 1
|
||||
assert post[0]["upstream_provider"] == "Novita"
|
||||
assert agent._last_serving_provider == "Novita"
|
||||
|
||||
def test_non_streamed_attempt_reports_no_upstream_provider(self, agent):
|
||||
"""A stale value from an earlier call must not leak into the next payload."""
|
||||
agent._last_serving_provider = "Novita" # as a prior streamed attempt left it
|
||||
agent.client.chat.completions.create.return_value = _mock_response(
|
||||
content="Done", finish_reason="stop"
|
||||
)
|
||||
|
||||
result, post = _run_with_hooks(agent)
|
||||
|
||||
assert result["final_response"] == "Done"
|
||||
assert len(post) == 1
|
||||
assert post[0]["upstream_provider"] is None
|
||||
Reference in New Issue
Block a user