diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 52d9642010..65641e9a0b 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -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"] diff --git a/agent/stream_diag.py b/agent/stream_diag.py index cb1b42e25c..66c691f478 100644 --- a/agent/stream_diag.py +++ b/agent/stream_diag.py @@ -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", diff --git a/agent/turn_api_request.py b/agent/turn_api_request.py index b00934cc72..5c339f8cef 100644 --- a/agent/turn_api_request.py +++ b/agent/turn_api_request.py @@ -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. diff --git a/agent/turn_response_intake.py b/agent/turn_response_intake.py index 67374daa78..43182ac143 100644 --- a/agent/turn_response_intake.py +++ b/agent/turn_response_intake.py @@ -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), diff --git a/hermes_cli/hooks.py b/hermes_cli/hooks.py index 6c0f3e64ae..68e1a6322b 100644 --- a/hermes_cli/hooks.py +++ b/hermes_cli/hooks.py @@ -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, diff --git a/run_agent.py b/run_agent.py index 83f121f968..ed7acb06fe 100644 --- a/run_agent.py +++ b/run_agent.py @@ -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") diff --git a/tests/agent/test_stream_serving_provider.py b/tests/agent/test_stream_serving_provider.py new file mode 100644 index 0000000000..338654dbaa --- /dev/null +++ b/tests/agent/test_stream_serving_provider.py @@ -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