test(integration): streamed /v1/responses success callbacks keep provider response headers (Pylon #7775)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/_support/wire.py

@@ -7,6 +7,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from queue import SimpleQueue
+from types import MappingProxyType
 from typing import Final
 
 
@@ -26,6 +27,7 @@ class Reply:
     chunks: tuple[bytes, ...] | None = None
     abort_after: int | None = None
     gate_after_first: threading.Event | None = None
+    headers: Mapping[str, str] = MappingProxyType({})
 
 
 @dataclass(frozen=True, slots=True)
@@ -64,6 +66,8 @@ def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None
                 reply = Reply(status=500)
             self.send_response(reply.status)
             self.send_header("content-type", reply.content_type)
+            for name, value in reply.headers.items():
+                self.send_header(name, value)
             if reply.chunks is None:
                 self.send_header("content-length", str(len(reply.body)))
             else:

tests/integration/contracts.json

@@ -1764,6 +1764,9 @@
     ],
     "tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
       "other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
+    ],
+    "tests/integration/observability/test_callback_delivery.py::test_streamed_responses_success_callback_carries_provider_apim_request_id": [
+      "other.observability.callbacks.streamed_responses_events_carry_provider_response_headers"
     ]
   },
   "browser": {

tests/integration/observability/test_callback_delivery.py

@@ -133,7 +133,7 @@ def test_concurrent_success_and_failure_join_callbacks_and_rows_without_credenti
                     assert "synthetic callback failure" in json.dumps(event["error_information"])
                 rows: Final = eventually(
                     lambda identity=event["id"]: read_rows(
-                        'SELECT request_id, spend, prompt_tokens, completion_tokens, request_tags '
+                        "SELECT request_id, spend, prompt_tokens, completion_tokens, request_tags "
                         'FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
                         (identity,),
                     ),
@@ -152,3 +152,115 @@ def test_concurrent_success_and_failure_join_callbacks_and_rows_without_credenti
                     assert rows[0]["prompt_tokens"] == event["prompt_tokens"]
                 else:
                     assert event["prompt_tokens"] == event["completion_tokens"] == rows[0]["completion_tokens"] == 0
+
+
+def _responses_frames(identity: str, text: str) -> tuple[bytes, ...]:
+    output: Final = [
+        {
+            "type": "message",
+            "id": f"msg_{identity}",
+            "status": "completed",
+            "role": "assistant",
+            "content": [{"type": "output_text", "text": text, "annotations": []}],
+        }
+    ]
+    completed: Final = {
+        "id": identity,
+        "object": "response",
+        "created_at": 1,
+        "status": "completed",
+        "model": "gpt-4o-mini",
+        "output": output,
+        "usage": {
+            "input_tokens": 11,
+            "output_tokens": 4,
+            "total_tokens": 15,
+            "input_tokens_details": {"cached_tokens": 0},
+            "output_tokens_details": {"reasoning_tokens": 0},
+        },
+    }
+    events: Final = (
+        {"type": "response.created", "response": {**completed, "status": "in_progress", "output": [], "usage": None}},
+        {
+            "type": "response.output_text.delta",
+            "item_id": f"msg_{identity}",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": text,
+        },
+        {"type": "response.completed", "response": completed},
+    )
+    return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events)
+
+
+@pytest.mark.covers("other.observability.callbacks.streamed_responses_events_carry_provider_response_headers")
+def test_streamed_responses_success_callback_carries_provider_apim_request_id(gateway: Gateway, tmp_path: Path) -> None:
+    marker: Final = "resp_" + uuid.uuid4().hex
+    correlation: Final = "azure-correlation-" + marker
+    region: Final = "East US 2"
+    secret: Final = "synthetic-provider-secret-" + marker
+    sink_secret: Final = "synthetic-sink-secret-" + marker
+
+    def upstream(request: Request) -> Reply:
+        assert request.target.endswith("/responses"), request.target
+        assert request.headers["authorization"] == f"Bearer {secret}"
+        assert json.loads(request.body) == {
+            "model": "gpt-4o-mini",
+            "input": "header control " + marker,
+            "stream": True,
+        }, request.body
+        return Reply(
+            content_type="text/event-stream",
+            chunks=_responses_frames(marker, "streamed control"),
+            headers={"apim-request-id": correlation, "x-ms-region": region},
+        )
+
+    def sink(request: Request) -> Reply:
+        assert request.headers["authorization"] == f"Bearer {sink_secret}"
+        return Reply()
+
+    with wire_server(upstream) as provider, wire_server(sink) as endpoint:
+        config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+        config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1})
+        path: Final = tmp_path / "callbacks.yaml"
+        path.write_text(yaml.safe_dump(config))
+        with (
+            owned_proxy(
+                gateway,
+                tmp_path,
+                {
+                    "GENERIC_LOGGER_ENDPOINT": endpoint.url,
+                    "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}",
+                },
+                config=path,
+            ) as candidate,
+            candidate.scenario() as scenario,
+        ):
+            model: Final = scenario.model(api_base=provider.url + "/v1", api_key=secret)
+            response: Final = candidate.request(
+                "POST", "/v1/responses", {"model": model, "input": "header control " + marker, "stream": True}
+            )
+            assert response.status_code == 200, response.text
+            assert f'"item_id":"msg_{marker}"' in response.text, response.text
+            assert '"type":"response.completed"' in response.text, response.text
+            assert len(provider.drain()) == 1
+            batches = []
+
+            def delivered() -> tuple[dict, ...]:
+                batches.extend(endpoint.drain())
+                return tuple(
+                    event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model
+                )
+
+            events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10)
+            assert (events[0]["status"], events[0]["stream"], events[0]["call_type"]) == ("success", True, "aresponses")
+            additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {}
+            provider_headers: Final = {
+                name: value
+                for name, value in additional_headers.items()
+                if name in ("llm_provider-apim-request-id", "llm_provider-x-ms-region")
+            }
+            assert provider_headers == {
+                "llm_provider-apim-request-id": correlation,
+                "llm_provider-x-ms-region": region,
+            }, json.dumps(events[0]["hidden_params"])