From 67115d5add3d58c90f56b38bba145190b8f483d6 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 9 Nov 2026 21:01:41 +0530 Subject: [PATCH 1/5] test: expect provider headers on successful managed responses Managed buffered and streaming successes drop the provider's request-id and rate limit headers, while failures keep them. This test fails until they are relayed. Signed-off-by: Somay-kousis --- .../cli/tests/coverage/shared/server_tests.rs | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index 1ded4a659..f448fd220 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -7196,6 +7196,103 @@ async fn unmanaged_routes_preserve_streaming_and_error_responses_without_session provider_task.abort(); } +#[tokio::test] +async fn managed_success_responses_keep_provider_headers() { + use crate::gateway::websocket::tests::serve; + let provider = Router::new().fallback(|request: Request| async move { + let body = axum::body::to_bytes(request.into_body(), usize::MAX) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&body).unwrap(); + let streaming = body["stream"] == json!(true); + let builder = Response::builder() + .status(StatusCode::OK) + .header("request-id", "req_mock_1") + .header("anthropic-ratelimit-unified-status", "allowed_warning") + .header("set-cookie", "a=1") + .header("set-cookie", "b=2") + .header("x-nemo-relay-worker-route-failure", "1"); + if streaming { + builder + .header("content-type", "text/event-stream; charset=utf-8") + .body(Body::from( + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-test\",\"content\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n\ + data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}\n\n\ + data: {\"type\":\"message_stop\"}\n\n", + )) + .unwrap() + } else { + builder + .header("content-type", "application/json; charset=utf-8") + .body(Body::from( + json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-test", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }) + .to_string(), + )) + .unwrap() + } + }); + let (upstream, provider_task) = serve(provider).await; + let mut config = test_config(); + config.anthropic_base_url = upstream; + let state = AppState::new(config); + let app = router_with_state(state.clone()); + for streaming in [false, true] { + let response = app + .clone() + .oneshot( + Request::post("/v1/messages") + .header("content-type", "application/json") + .header("x-api-key", "sk-ant-test") + .body(Body::from( + json!({ + "model": "claude-test", + "max_tokens": 1, + "stream": streaming, + "messages": [{"role": "user", "content": "hi"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK, "streaming={streaming}"); + let headers = response.headers(); + assert_eq!(headers["request-id"], "req_mock_1", "streaming={streaming}"); + assert_eq!( + headers["anthropic-ratelimit-unified-status"], "allowed_warning", + "streaming={streaming}" + ); + assert_eq!(headers.get_all("set-cookie").iter().count(), 2); + // A provider must not be able to raise Relay's own worker route failure signal. + assert!( + headers.get("x-nemo-relay-worker-route-failure").is_none(), + "streaming={streaming}" + ); + // The gateway re-encodes the body, so the content type stays its own. + assert_eq!( + headers.get_all("content-type").iter().collect::>(), + [if streaming { + "text/event-stream" + } else { + "application/json" + }], + "streaming={streaming}" + ); + response.into_body().collect().await.unwrap(); + } + state.sessions.close_all("test_shutdown").await.unwrap(); + provider_task.abort(); +} + #[tokio::test] async fn websocket_tunnel_blocks_idle_shutdown_until_disconnect() { use crate::gateway::websocket::tests::{provider, serve}; From 3bfca0f3e80b076b6e38bcb3735c842a1eda8f2a Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 9 Nov 2026 21:03:27 +0530 Subject: [PATCH 2/5] fix: relay provider headers on successful managed responses Managed buffered and streaming successes built their headers from an empty map, so clients lost request-id and rate limit headers that failures keep. Capture each successful attempt's headers and always relay one attempt's: for buffered calls the attempt whose JSON was returned, for streams the attempt whose stream yielded the first event, else the latest success. Relay's own x-nemo-relay-* headers, content-type and content-encoding are not taken from the provider. The two concurrent-next tests now expect the selected attempt's header instead of none. Signed-off-by: Somay-kousis --- crates/cli/src/gateway/mod.rs | 152 ++++++++++++++++-- .../tests/coverage/shared/gateway_tests.rs | 47 ++++++ .../cli/tests/coverage/shared/server_tests.rs | 137 +++++++++++++++- 3 files changed, 317 insertions(+), 19 deletions(-) diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 264833a87..6c5825996 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -36,7 +36,7 @@ use nemo_relay_runtime::api::llm::{ }; use nemo_relay_runtime::api::runtime::{ LlmCodecIdentity, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmStreamExecutionNextFn, - TASK_SCOPE_STACK, + LlmStreamInner, TASK_SCOPE_STACK, }; use nemo_relay_runtime::codec::response_codec as build_response_codec; use nemo_relay_types::codec::request::AnnotatedLlmRequest; @@ -465,6 +465,122 @@ impl CapturedUpstreamFailures { type CapturedUpstreamFailuresRef = Arc; +/// Provider headers from the successful attempts of one gateway call. +/// +/// A successful response always carries one attempt's provider headers, matching failures. +/// A buffered success uses the attempt whose JSON the runtime returned. A streaming success +/// uses the attempt whose stream most recently yielded its first event, which is the stream the +/// gateway reads right before it sends headers. Middleware that rewrites a buffered response +/// leaves no exact match, so the latest successful attempt is used. +#[derive(Default)] +struct CapturedUpstreamSuccesses { + entries: Mutex>, + first_event: Mutex>, +} + +struct CapturedUpstreamSuccess { + response: Option, + headers: HeaderMap, +} + +impl CapturedUpstreamSuccesses { + fn capture(&self, response: Option, headers: HeaderMap) -> usize { + let mut entries = self + .entries + .lock() + .expect("captured upstream successes lock poisoned"); + entries.push(CapturedUpstreamSuccess { response, headers }); + entries.len() - 1 + } + + fn mark_first_event(&self, index: usize) { + *self + .first_event + .lock() + .expect("captured upstream successes lock poisoned") = Some(index); + } + + /// Returns the client headers for the response, with the gateway's own content type. + fn response_headers(&self, selected: Option<&Value>, content_type: &'static str) -> HeaderMap { + let mut entries = self + .entries + .lock() + .expect("captured upstream successes lock poisoned"); + let first_event = *self + .first_event + .lock() + .expect("captured upstream successes lock poisoned"); + let index = selected + .and_then(|selected| { + entries + .iter() + .rposition(|entry| entry.response.as_ref() == Some(selected)) + }) + .or(first_event) + .or_else(|| entries.len().checked_sub(1)); + let mut headers = index.map_or_else(HeaderMap::new, |index| { + std::mem::take(&mut entries[index].headers) + }); + // Relay's own control headers (for example the worker route failure signal) must + // never be taken from a provider, as every worker forwarding path already ensures. + let private: Vec<_> = headers + .keys() + .filter(|name| name.as_str().starts_with("x-nemo-relay-")) + .cloned() + .collect(); + for name in private { + headers.remove(name); + } + // The gateway re-encodes the body, so the provider's body headers no longer apply. + headers.remove(http::header::CONTENT_TYPE); + headers.remove(http::header::CONTENT_ENCODING); + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static(content_type), + ); + headers + } +} + +type CapturedUpstreamSuccessesRef = Arc; + +/// Records which attempt's stream yielded its first event, without changing close semantics. +struct FirstEventMarked { + inner: LlmJsonStream, + mark: Option<(CapturedUpstreamSuccessesRef, usize)>, +} + +impl futures_util::Stream for FirstEventMarked { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + let polled = std::pin::Pin::new(&mut this.inner).poll_next(cx); + if let std::task::Poll::Ready(Some(Ok(_))) = &polled + && let Some((successes, index)) = this.mark.take() + { + successes.mark_first_event(index); + } + polled + } +} + +impl LlmStreamInner for FirstEventMarked { + fn terminalize(self: std::pin::Pin<&mut Self>) { + self.get_mut().inner.terminalize(); + } + + fn close( + self: std::pin::Pin<&mut Self>, + ) -> std::pin::Pin> + Send + '_>> + { + Box::pin(self.get_mut().inner.close()) + } +} + // Runs the managed pipeline for a prepared gateway request. Streaming and non-streaming branches // share the same prep + codec dispatch but diverge in how the runtime drives the upstream call. async fn run_managed_gateway( @@ -656,11 +772,13 @@ async fn run_managed_buffered( operational: OperationalContext, ) -> Result, CliError> { let upstream_failures = Arc::new(CapturedUpstreamFailures::default()); + let upstream_successes = Arc::new(CapturedUpstreamSuccesses::default()); let func = build_buffered_func( state.clone(), &prepared, codecs.source_body, upstream_failures.clone(), + upstream_successes.clone(), operational.clone(), ); let GatewayCallPrep { @@ -694,11 +812,8 @@ async fn run_managed_buffered( match result { Ok(response_json) => { operational::upstream_completed(&operational, false); - let mut headers = HeaderMap::new(); - headers.insert( - http::header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ); + let headers = + upstream_successes.response_headers(Some(&response_json), "application/json"); let response_body = Body::from(response_json.to_string()); if record_task_response { state @@ -738,6 +853,7 @@ fn build_buffered_func( prepared: &PreparedGatewayRequest, source_body: SourceRequestBody, upstream_failures: CapturedUpstreamFailuresRef, + upstream_successes: CapturedUpstreamSuccessesRef, operational: OperationalContext, ) -> LlmExecutionNextFn { let http = upstream_client(&state, prepared.client_named_upstream).clone(); @@ -759,6 +875,7 @@ fn build_buffered_func( let body_bytes = body_bytes.clone(); let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); + let upstream_successes = upstream_successes.clone(); let operational = operational.clone(); let source_body = source_body.clone(); Box::pin(async move { @@ -831,7 +948,10 @@ fn build_buffered_func( ); } match serde_json::from_slice::(&bytes) { - Ok(response) => Ok(response), + Ok(response) => { + upstream_successes.capture(Some(response.clone()), response_headers); + Ok(response) + } Err(_) if retry_aware => { Err(FlowError::Upstream(invalid_json_failure(&response_headers))) } @@ -870,11 +990,13 @@ async fn run_managed_streaming( ) }); let upstream_failures = Arc::new(CapturedUpstreamFailures::default()); + let upstream_successes = Arc::new(CapturedUpstreamSuccesses::default()); let func = build_streaming_func( state.clone(), &prepared, codecs.source_body, upstream_failures.clone(), + upstream_successes.clone(), operational.clone(), ); let provider_route = prepared.provider; @@ -947,11 +1069,6 @@ async fn run_managed_streaming( return Err(translate_runtime_error(error)); } }; - let mut headers = HeaderMap::new(); - headers.insert( - http::header::CONTENT_TYPE, - HeaderValue::from_static("text/event-stream"), - ); guard.operational = operational; guard.response_ready = Some(response_ready); let body = match prepare_streaming_response_body(json_stream, provider_route, guard).await { @@ -963,6 +1080,9 @@ async fn run_managed_streaming( return Err(translate_runtime_error(error)); } }; + // Read after the first event, so the headers belong to the stream the client receives. + // Retries later in the stream can't change headers the client already has. + let headers = upstream_successes.response_headers(None, "text/event-stream"); build_response(StatusCode::OK, headers, body) } @@ -1031,6 +1151,7 @@ fn build_streaming_func( prepared: &PreparedGatewayRequest, source_body: SourceRequestBody, upstream_failures: CapturedUpstreamFailuresRef, + upstream_successes: CapturedUpstreamSuccessesRef, operational: OperationalContext, ) -> LlmStreamExecutionNextFn { let http = upstream_client(&state, prepared.client_named_upstream).clone(); @@ -1052,6 +1173,7 @@ fn build_streaming_func( let body_bytes = body_bytes.clone(); let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); + let upstream_successes = upstream_successes.clone(); let operational = operational.clone(); let source_body = source_body.clone(); Box::pin(async move { @@ -1123,8 +1245,12 @@ fn build_streaming_func( }), ); } + let index = upstream_successes.capture(None, response_headers); let json_stream = sse_json_stream(response, operational); - Ok(json_stream) + Ok(LlmJsonStream::from_closeable(FirstEventMarked { + inner: json_stream, + mark: Some((upstream_successes, index)), + })) }) }) } diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index aa6f0eee7..cc003d913 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -1556,6 +1556,7 @@ async fn streaming_provider_error_does_not_poison_the_next_request() { &prepared, SourceRequestBody::default(), Arc::new(CapturedUpstreamFailures::default()), + Arc::new(CapturedUpstreamSuccesses::default()), crate::operational::OperationalContext::new(), ); let mut request_headers = Map::new(); @@ -1628,6 +1629,7 @@ async fn buffered_body_read_failure_stays_structured() { &prepared, SourceRequestBody::default(), Arc::new(CapturedUpstreamFailures::default()), + Arc::new(CapturedUpstreamSuccesses::default()), crate::operational::OperationalContext::new(), ); let mut request_headers = Map::new(); @@ -1683,6 +1685,7 @@ async fn buffered_invalid_json_becomes_safe_upstream_failure() { &prepared, SourceRequestBody::default(), Arc::new(CapturedUpstreamFailures::default()), + Arc::new(CapturedUpstreamSuccesses::default()), crate::operational::OperationalContext::new(), ); let mut request_headers = Map::new(); @@ -2863,6 +2866,50 @@ fn captured_upstream_failures_reject_ambiguous_attempts_without_consuming_them() assert!(failures.take(&selected).is_some()); } +/// A success always carries one attempt's headers: the one whose JSON was returned, else the latest. +#[test] +fn captured_upstream_successes_select_the_returned_attempt() { + let attempt = |id: &'static str| { + let mut headers = HeaderMap::new(); + headers.insert("request-id", HeaderValue::from_static(id)); + headers.insert( + "content-type", + HeaderValue::from_static("application/x-provider"), + ); + headers.insert("content-encoding", HeaderValue::from_static("gzip")); + headers + }; + + let none = CapturedUpstreamSuccesses::default(); + let headers = none.response_headers(None, "application/json"); + assert_eq!(headers.len(), 1); + assert_eq!(headers["content-type"], "application/json"); + + let one = CapturedUpstreamSuccesses::default(); + one.capture(None, attempt("req_1")); + let headers = one.response_headers(None, "text/event-stream"); + assert_eq!(headers["request-id"], "req_1"); + assert_eq!(headers["content-type"], "text/event-stream"); + assert_eq!(headers.get_all("content-type").iter().count(), 1); + assert!(headers.get("content-encoding").is_none()); + + // Fan-out: the returned JSON picks its own attempt, even when it finished first. + let first = json!({"id": "first"}); + let second = json!({"id": "second"}); + let fan_out = CapturedUpstreamSuccesses::default(); + fan_out.capture(Some(first.clone()), attempt("req_1")); + fan_out.capture(Some(second), attempt("req_2")); + let headers = fan_out.response_headers(Some(&first), "application/json"); + assert_eq!(headers["request-id"], "req_1"); + + // A rewritten response matches no attempt, so the latest success is used. + let rewritten = CapturedUpstreamSuccesses::default(); + rewritten.capture(Some(first), attempt("req_1")); + rewritten.capture(Some(json!({"id": "second"})), attempt("req_2")); + let headers = rewritten.response_headers(Some(&json!({"id": "rewritten"})), "application/json"); + assert_eq!(headers["request-id"], "req_2"); +} + #[tokio::test] async fn streaming_gateway_call_guard_finishes_when_body_is_dropped() { let manager = SessionManager::new(GatewayConfig::default()); diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index f448fd220..1bf0bdb8b 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -5400,6 +5400,21 @@ async fn spawn_concurrent_attempt_upstream(release_losing: Arc) -> Te if attempt == "losing" { release_losing.acquire().await.unwrap().forget(); } + if attempt == "losing" && request["losing_succeeds"].as_bool() == Some(true) { + let body = if streaming { + CONCURRENT_LOSING_CHAT_SSE.to_string() + } else { + let mut response = concurrent_selected_chat_response(); + response["id"] = json!("chatcmpl-losing"); + response.to_string() + }; + return Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "application/x-losing-success") + .header("x-upstream-attempt", attempt) + .body(Body::from(body)) + .unwrap(); + } let (status, content_type, body) = match (attempt, streaming, both_fail) { ("selected", false, true) => ( @@ -5543,9 +5558,14 @@ async fn gateway_concurrent_next_uses_canonical_selected_buffered_response() { .and_then(|value| value.to_str().ok()), Some("application/json") ); - assert!( - response.headers().get("x-upstream-attempt").is_none(), - "managed success must not inherit metadata from an arbitrary upstream attempt" + assert_eq!( + response + .headers() + .get_all("x-upstream-attempt") + .iter() + .collect::>(), + ["selected"], + "managed success must carry the selected attempt's headers, never the losing attempt's" ); let body = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!( @@ -5628,9 +5648,14 @@ async fn gateway_concurrent_next_uses_canonical_selected_streaming_response() { .and_then(|value| value.to_str().ok()), Some("text/event-stream") ); - assert!( - response.headers().get("x-upstream-attempt").is_none(), - "managed stream success must not inherit metadata from an arbitrary upstream attempt" + assert_eq!( + response + .headers() + .get_all("x-upstream-attempt") + .iter() + .collect::>(), + ["selected"], + "managed stream success must carry the selected attempt's headers, never the losing attempt's" ); let body = response.into_body().collect().await.unwrap().to_bytes(); let body = std::str::from_utf8(&body).unwrap(); @@ -5638,6 +5663,106 @@ async fn gateway_concurrent_next_uses_canonical_selected_streaming_response() { assert!(!body.contains("losing attempt"), "{body}"); } +/// Both concurrent attempts succeed and the later one finishes last; the returned attempt's +/// headers must still win over the latest success. +#[tokio::test] +async fn gateway_concurrent_successes_keep_the_selected_attempt_headers() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = crate::test_support::plugin_host::test_close_plugin_host(); + + const BUFFERED_NAME: &str = "cli-test-concurrent-buffered-successes"; + const STREAMING_NAME: &str = "cli-test-concurrent-streaming-successes"; + const MARKER: &str = "select one of two successful concurrent attempts"; + let _ = deregister_llm_execution_intercept(BUFFERED_NAME); + let _buffered_cleanup = LlmExecutionInterceptCleanup(BUFFERED_NAME); + let _ = deregister_llm_stream_execution_intercept(STREAMING_NAME); + let _streaming_cleanup = LlmStreamExecutionInterceptCleanup(STREAMING_NAME); + let release_losing = Arc::new(Semaphore::new(0)); + let buffered_release = release_losing.clone(); + register_llm_execution_intercept( + BUFFERED_NAME, + 1, + Arc::new(move |_name, request, _context, next| { + let release_losing = buffered_release.clone(); + Box::pin(async move { + if request.content["messages"][0]["content"].as_str() != Some(MARKER) { + return next(request).await; + } + let mut selected_request = request.clone(); + selected_request.content["attempt"] = json!("selected"); + let mut losing_request = request; + losing_request.content["attempt"] = json!("losing"); + losing_request.content["losing_succeeds"] = json!(true); + let selected = next(selected_request).await; + release_losing.add_permits(1); + assert!(next(losing_request).await.is_ok()); + selected + }) + }), + ) + .unwrap(); + let streaming_release = release_losing.clone(); + register_llm_stream_execution_intercept( + STREAMING_NAME, + 1, + Arc::new(move |_name, request, _context, next| { + let release_losing = streaming_release.clone(); + Box::pin(async move { + if request.content["messages"][0]["content"].as_str() != Some(MARKER) { + return next(request).await; + } + let mut selected_request = request.clone(); + selected_request.content["attempt"] = json!("selected"); + let mut losing_request = request; + losing_request.content["attempt"] = json!("losing"); + losing_request.content["losing_succeeds"] = json!(true); + let selected = next(selected_request).await; + release_losing.add_permits(1); + assert!(next(losing_request).await.is_ok()); + selected + }) + }), + ) + .unwrap(); + + let upstream = spawn_concurrent_attempt_upstream(release_losing).await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + for streaming in [false, true] { + let response = router(config.clone()) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": MARKER}], + "stream": streaming + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK, "streaming={streaming}"); + assert_eq!( + response + .headers() + .get_all("x-upstream-attempt") + .iter() + .collect::>(), + ["selected"], + "streaming={streaming}" + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(!body.contains("losing"), "streaming={streaming}: {body}"); + } +} + #[tokio::test] async fn gateway_concurrent_next_relays_selected_buffered_failure() { let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; From 66d3b02c9bb72397cd5f3e38a37403cd223ba2f5 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 9 Nov 2026 21:27:56 +0530 Subject: [PATCH 3/5] fix: drop provider digest headers from re-encoded success responses content-digest and repr-digest describe the provider's bytes, which the gateway no longer sends. Tests now also cover the private header and digest drops, the first-event selection for streams, and the exact set-cookie values. Signed-off-by: Somay-kousis --- crates/cli/src/gateway/mod.rs | 2 ++ .../cli/tests/coverage/shared/gateway_tests.rs | 17 +++++++++++++++++ .../cli/tests/coverage/shared/server_tests.rs | 17 ++++++++++++++++- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 6c5825996..9dc1f8136 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -534,6 +534,8 @@ impl CapturedUpstreamSuccesses { // The gateway re-encodes the body, so the provider's body headers no longer apply. headers.remove(http::header::CONTENT_TYPE); headers.remove(http::header::CONTENT_ENCODING); + headers.remove("content-digest"); + headers.remove("repr-digest"); headers.insert( http::header::CONTENT_TYPE, HeaderValue::from_static(content_type), diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index cc003d913..7d4e647ff 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -2877,6 +2877,11 @@ fn captured_upstream_successes_select_the_returned_attempt() { HeaderValue::from_static("application/x-provider"), ); headers.insert("content-encoding", HeaderValue::from_static("gzip")); + headers.insert("content-digest", HeaderValue::from_static("sha-256=:AAAA:")); + headers.insert( + "x-nemo-relay-worker-route-failure", + HeaderValue::from_static("pass-through"), + ); headers }; @@ -2892,6 +2897,18 @@ fn captured_upstream_successes_select_the_returned_attempt() { assert_eq!(headers["content-type"], "text/event-stream"); assert_eq!(headers.get_all("content-type").iter().count(), 1); assert!(headers.get("content-encoding").is_none()); + assert!(headers.get("content-digest").is_none()); + assert!(headers.get("x-nemo-relay-worker-route-failure").is_none()); + + // Streams: the attempt whose stream yielded the first event wins over capture order. + for (marked, expected) in [(0, "req_1"), (1, "req_2")] { + let streams = CapturedUpstreamSuccesses::default(); + streams.capture(None, attempt("req_1")); + streams.capture(None, attempt("req_2")); + streams.mark_first_event(marked); + let headers = streams.response_headers(None, "text/event-stream"); + assert_eq!(headers["request-id"], expected); + } // Fan-out: the returned JSON picks its own attempt, even when it finished first. let first = json!({"id": "first"}); diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index 1bf0bdb8b..db4a8d9d7 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -7336,6 +7336,8 @@ async fn managed_success_responses_keep_provider_headers() { .header("anthropic-ratelimit-unified-status", "allowed_warning") .header("set-cookie", "a=1") .header("set-cookie", "b=2") + .header("content-encoding", "identity") + .header("repr-digest", "sha-256=:AAAA:") .header("x-nemo-relay-worker-route-failure", "1"); if streaming { builder @@ -7396,7 +7398,20 @@ async fn managed_success_responses_keep_provider_headers() { headers["anthropic-ratelimit-unified-status"], "allowed_warning", "streaming={streaming}" ); - assert_eq!(headers.get_all("set-cookie").iter().count(), 2); + assert_eq!( + headers.get_all("set-cookie").iter().collect::>(), + ["a=1", "b=2"], + "streaming={streaming}" + ); + // The body is re-encoded, so the provider's encoding and digest no longer describe it. + assert!( + headers.get("content-encoding").is_none(), + "streaming={streaming}" + ); + assert!( + headers.get("repr-digest").is_none(), + "streaming={streaming}" + ); // A provider must not be able to raise Relay's own worker route failure signal. assert!( headers.get("x-nemo-relay-worker-route-failure").is_none(), From 07b5fb8870f79257248095b4a8f0f5e32f2bcce4 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 9 Nov 2026 23:09:49 +0530 Subject: [PATCH 4/5] fix: pick success headers by the content the gateway returns Headers followed the stream polled last and the last matching body, so middleware that read both streams, or two attempts with the same body, could give the client the wrong attempt's headers. Key each attempt by its JSON or first stream event and match what the gateway actually returns. Several matches, or a rewritten response with several successes, keep only the headers every candidate sent with the same values. Signed-off-by: Somay-kousis --- crates/cli/src/gateway/mod.rs | 137 ++++++++++------ .../tests/coverage/shared/gateway_tests.rs | 125 ++++++++++----- .../cli/tests/coverage/shared/server_tests.rs | 147 ++++++++++++++++++ 3 files changed, 327 insertions(+), 82 deletions(-) diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 9dc1f8136..3579088a4 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -467,60 +467,82 @@ type CapturedUpstreamFailuresRef = Arc; /// Provider headers from the successful attempts of one gateway call. /// -/// A successful response always carries one attempt's provider headers, matching failures. -/// A buffered success uses the attempt whose JSON the runtime returned. A streaming success -/// uses the attempt whose stream most recently yielded its first event, which is the stream the -/// gateway reads right before it sends headers. Middleware that rewrites a buffered response -/// leaves no exact match, so the latest successful attempt is used. +/// Each attempt is keyed by what it produced: its JSON for a buffered call, its first event for +/// a stream. The client response uses the attempts whose key equals what the gateway actually +/// returns. When several attempts match, or a rewritten response matches none, only headers +/// that every candidate sent with the same values are relayed, so a header never comes from +/// the wrong attempt. #[derive(Default)] struct CapturedUpstreamSuccesses { entries: Mutex>, - first_event: Mutex>, } struct CapturedUpstreamSuccess { - response: Option, + key: Option, headers: HeaderMap, } impl CapturedUpstreamSuccesses { - fn capture(&self, response: Option, headers: HeaderMap) -> usize { + fn capture(&self, key: Option, headers: HeaderMap) -> usize { let mut entries = self .entries .lock() .expect("captured upstream successes lock poisoned"); - entries.push(CapturedUpstreamSuccess { response, headers }); + entries.push(CapturedUpstreamSuccess { key, headers }); entries.len() - 1 } - fn mark_first_event(&self, index: usize) { - *self - .first_event + fn set_key(&self, index: usize, key: Value) { + if let Some(entry) = self + .entries .lock() - .expect("captured upstream successes lock poisoned") = Some(index); + .expect("captured upstream successes lock poisoned") + .get_mut(index) + { + entry.key = Some(key); + } } - /// Returns the client headers for the response, with the gateway's own content type. - fn response_headers(&self, selected: Option<&Value>, content_type: &'static str) -> HeaderMap { + /// Returns the client headers for `returned`, with the gateway's own content type. + fn response_headers(&self, returned: Option<&Value>, content_type: &'static str) -> HeaderMap { let mut entries = self .entries .lock() .expect("captured upstream successes lock poisoned"); - let first_event = *self - .first_event - .lock() - .expect("captured upstream successes lock poisoned"); - let index = selected - .and_then(|selected| { - entries - .iter() - .rposition(|entry| entry.response.as_ref() == Some(selected)) + let mut candidates: Vec = returned + .map(|returned| { + (0..entries.len()) + .filter(|&index| entries[index].key.as_ref() == Some(returned)) + .collect() }) - .or(first_event) - .or_else(|| entries.len().checked_sub(1)); - let mut headers = index.map_or_else(HeaderMap::new, |index| { - std::mem::take(&mut entries[index].headers) - }); + .unwrap_or_default(); + if candidates.is_empty() { + candidates = (0..entries.len()).collect(); + } + let mut headers = match candidates.as_slice() { + [] => HeaderMap::new(), + [only] => std::mem::take(&mut entries[*only].headers), + [first, rest @ ..] => { + let mut shared = HeaderMap::new(); + for name in entries[*first].headers.keys() { + let values: Vec<_> = entries[*first].headers.get_all(name).iter().collect(); + let agreed = rest.iter().all(|&other| { + entries[other] + .headers + .get_all(name) + .iter() + .collect::>() + == values + }); + if agreed { + for value in values { + shared.append(name.clone(), value.clone()); + } + } + } + shared + } + }; // Relay's own control headers (for example the worker route failure signal) must // never be taken from a provider, as every worker forwarding path already ensures. let private: Vec<_> = headers @@ -546,13 +568,25 @@ impl CapturedUpstreamSuccesses { type CapturedUpstreamSuccessesRef = Arc; -/// Records which attempt's stream yielded its first event, without changing close semantics. -struct FirstEventMarked { +type FirstEventObserver = Box; + +/// Hands the first successful event of a stream to an observer, without changing close +/// semantics. +struct FirstEventTap { inner: LlmJsonStream, - mark: Option<(CapturedUpstreamSuccessesRef, usize)>, + observer: Option, +} + +impl FirstEventTap { + fn wrap(inner: LlmJsonStream, observer: FirstEventObserver) -> LlmJsonStream { + LlmJsonStream::from_closeable(Self { + inner, + observer: Some(observer), + }) + } } -impl futures_util::Stream for FirstEventMarked { +impl futures_util::Stream for FirstEventTap { type Item = Result; fn poll_next( @@ -561,16 +595,16 @@ impl futures_util::Stream for FirstEventMarked { ) -> std::task::Poll> { let this = self.get_mut(); let polled = std::pin::Pin::new(&mut this.inner).poll_next(cx); - if let std::task::Poll::Ready(Some(Ok(_))) = &polled - && let Some((successes, index)) = this.mark.take() + if let std::task::Poll::Ready(Some(Ok(event))) = &polled + && let Some(observer) = this.observer.take() { - successes.mark_first_event(index); + observer(event); } polled } } -impl LlmStreamInner for FirstEventMarked { +impl LlmStreamInner for FirstEventTap { fn terminalize(self: std::pin::Pin<&mut Self>) { self.get_mut().inner.terminalize(); } @@ -1073,6 +1107,16 @@ async fn run_managed_streaming( }; guard.operational = operational; guard.response_ready = Some(response_ready); + let returned_first_event = Arc::new(Mutex::new(None::)); + let record_returned = Arc::clone(&returned_first_event); + let json_stream = FirstEventTap::wrap( + json_stream, + Box::new(move |event| { + *record_returned + .lock() + .expect("returned first event lock poisoned") = Some(event.clone()); + }), + ); let body = match prepare_streaming_response_body(json_stream, provider_route, guard).await { Ok(body) => body, Err(error) => { @@ -1082,9 +1126,14 @@ async fn run_managed_streaming( return Err(translate_runtime_error(error)); } }; - // Read after the first event, so the headers belong to the stream the client receives. - // Retries later in the stream can't change headers the client already has. - let headers = upstream_successes.response_headers(None, "text/event-stream"); + // Match the first event the client receives, so replayed or pre-read streams still pick + // the attempt that produced it. Later retries can't change headers already sent. + let returned_first_event = returned_first_event + .lock() + .expect("returned first event lock poisoned") + .take(); + let headers = + upstream_successes.response_headers(returned_first_event.as_ref(), "text/event-stream"); build_response(StatusCode::OK, headers, body) } @@ -1249,10 +1298,10 @@ fn build_streaming_func( } let index = upstream_successes.capture(None, response_headers); let json_stream = sse_json_stream(response, operational); - Ok(LlmJsonStream::from_closeable(FirstEventMarked { - inner: json_stream, - mark: Some((upstream_successes, index)), - })) + Ok(FirstEventTap::wrap( + json_stream, + Box::new(move |event| upstream_successes.set_key(index, event.clone())), + )) }) }) } diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index 7d4e647ff..2ae61d420 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -2866,32 +2866,46 @@ fn captured_upstream_failures_reject_ambiguous_attempts_without_consuming_them() assert!(failures.take(&selected).is_some()); } -/// A success always carries one attempt's headers: the one whose JSON was returned, else the latest. -#[test] -fn captured_upstream_successes_select_the_returned_attempt() { - let attempt = |id: &'static str| { - let mut headers = HeaderMap::new(); - headers.insert("request-id", HeaderValue::from_static(id)); - headers.insert( - "content-type", - HeaderValue::from_static("application/x-provider"), - ); - headers.insert("content-encoding", HeaderValue::from_static("gzip")); - headers.insert("content-digest", HeaderValue::from_static("sha-256=:AAAA:")); - headers.insert( - "x-nemo-relay-worker-route-failure", - HeaderValue::from_static("pass-through"), - ); - headers - }; +fn success_headers(id: &'static str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert("request-id", HeaderValue::from_static(id)); + headers.insert( + "content-type", + HeaderValue::from_static("application/x-provider"), + ); + headers.insert("content-encoding", HeaderValue::from_static("gzip")); + headers.insert("content-digest", HeaderValue::from_static("sha-256=:AAAA:")); + headers.insert( + "x-nemo-relay-worker-route-failure", + HeaderValue::from_static("pass-through"), + ); + headers +} + +fn shared_success_headers(id: &'static str) -> HeaderMap { + let mut headers = success_headers(id); + headers.insert( + "anthropic-ratelimit-unified-status", + HeaderValue::from_static("allowed"), + ); + headers +} + +fn assert_only_shared_headers(headers: &HeaderMap) { + assert!(headers.get("request-id").is_none()); + assert_eq!(headers["anthropic-ratelimit-unified-status"], "allowed"); +} +/// Provider body and control headers never reach the client, and the gateway owns content-type. +#[test] +fn captured_upstream_successes_filter_provider_headers() { let none = CapturedUpstreamSuccesses::default(); let headers = none.response_headers(None, "application/json"); assert_eq!(headers.len(), 1); assert_eq!(headers["content-type"], "application/json"); let one = CapturedUpstreamSuccesses::default(); - one.capture(None, attempt("req_1")); + one.capture(None, success_headers("req_1")); let headers = one.response_headers(None, "text/event-stream"); assert_eq!(headers["request-id"], "req_1"); assert_eq!(headers["content-type"], "text/event-stream"); @@ -2899,32 +2913,67 @@ fn captured_upstream_successes_select_the_returned_attempt() { assert!(headers.get("content-encoding").is_none()); assert!(headers.get("content-digest").is_none()); assert!(headers.get("x-nemo-relay-worker-route-failure").is_none()); +} - // Streams: the attempt whose stream yielded the first event wins over capture order. - for (marked, expected) in [(0, "req_1"), (1, "req_2")] { - let streams = CapturedUpstreamSuccesses::default(); - streams.capture(None, attempt("req_1")); - streams.capture(None, attempt("req_2")); - streams.mark_first_event(marked); - let headers = streams.response_headers(None, "text/event-stream"); +/// A success carries the headers of the attempt that produced what was returned. +#[test] +fn captured_upstream_successes_select_the_returned_attempt() { + // The returned body picks its own attempt in either capture order, buffered or streamed. + let first = json!({"id": "first"}); + let second = json!({"id": "second"}); + for (returned, expected) in [(&first, "req_1"), (&second, "req_2")] { + let fan_out = CapturedUpstreamSuccesses::default(); + fan_out.capture(Some(first.clone()), success_headers("req_1")); + let index = fan_out.capture(None, success_headers("req_2")); + fan_out.set_key(index, second.clone()); + let headers = fan_out.response_headers(Some(returned), "application/json"); assert_eq!(headers["request-id"], expected); } - // Fan-out: the returned JSON picks its own attempt, even when it finished first. + // A rewritten response matches no attempt; one success still lends its headers. + let rewritten_one = CapturedUpstreamSuccesses::default(); + rewritten_one.capture(Some(first), success_headers("req_1")); + let headers = + rewritten_one.response_headers(Some(&json!({"id": "rewritten"})), "application/json"); + assert_eq!(headers["request-id"], "req_1"); +} + +/// Ambiguous matches relay only the headers every candidate agrees on. +#[test] +fn captured_upstream_successes_keep_only_agreed_headers_when_ambiguous() { let first = json!({"id": "first"}); let second = json!({"id": "second"}); - let fan_out = CapturedUpstreamSuccesses::default(); - fan_out.capture(Some(first.clone()), attempt("req_1")); - fan_out.capture(Some(second), attempt("req_2")); - let headers = fan_out.response_headers(Some(&first), "application/json"); - assert_eq!(headers["request-id"], "req_1"); - // A rewritten response matches no attempt, so the latest success is used. - let rewritten = CapturedUpstreamSuccesses::default(); - rewritten.capture(Some(first), attempt("req_1")); - rewritten.capture(Some(json!({"id": "second"})), attempt("req_2")); - let headers = rewritten.response_headers(Some(&json!({"id": "rewritten"})), "application/json"); - assert_eq!(headers["request-id"], "req_2"); + // Identical bodies with different headers. + let identical = CapturedUpstreamSuccesses::default(); + identical.capture(Some(first.clone()), shared_success_headers("req_1")); + identical.capture(Some(first.clone()), shared_success_headers("req_2")); + assert_only_shared_headers(&identical.response_headers(Some(&first), "application/json")); + + // A rewritten response with several successes. + let rewritten_many = CapturedUpstreamSuccesses::default(); + rewritten_many.capture(Some(first), shared_success_headers("req_1")); + rewritten_many.capture(Some(second), shared_success_headers("req_2")); + assert_only_shared_headers( + &rewritten_many.response_headers(Some(&json!({"id": "rewritten"})), "application/json"), + ); + + // Streams whose first events are identical. + let same_first_event = CapturedUpstreamSuccesses::default(); + for id in ["req_1", "req_2"] { + let index = same_first_event.capture(None, shared_success_headers(id)); + same_first_event.set_key(index, json!({"type": "message_start"})); + } + assert_only_shared_headers( + &same_first_event + .response_headers(Some(&json!({"type": "message_start"})), "text/event-stream"), + ); + + // An empty stream returns no first event. + let empty_stream = CapturedUpstreamSuccesses::default(); + empty_stream.capture(None, shared_success_headers("req_1")); + empty_stream.capture(None, shared_success_headers("req_2")); + assert_only_shared_headers(&empty_stream.response_headers(None, "text/event-stream")); } #[tokio::test] diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index db4a8d9d7..e7de5736a 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -5403,6 +5403,9 @@ async fn spawn_concurrent_attempt_upstream(release_losing: Arc) -> Te if attempt == "losing" && request["losing_succeeds"].as_bool() == Some(true) { let body = if streaming { CONCURRENT_LOSING_CHAT_SSE.to_string() + } else if request["same_body"].as_bool() == Some(true) { + // An upstream cache can answer both attempts with the same body. + concurrent_selected_chat_response().to_string() } else { let mut response = concurrent_selected_chat_response(); response["id"] = json!("chatcmpl-losing"); @@ -5412,6 +5415,7 @@ async fn spawn_concurrent_attempt_upstream(release_losing: Arc) -> Te .status(StatusCode::OK) .header(header::CONTENT_TYPE, "application/x-losing-success") .header("x-upstream-attempt", attempt) + .header("x-upstream-shared", "same") .body(Body::from(body)) .unwrap(); } @@ -5464,6 +5468,7 @@ async fn spawn_concurrent_attempt_upstream(release_losing: Arc) -> Te .status(status) .header(header::CONTENT_TYPE, content_type) .header("x-upstream-attempt", attempt) + .header("x-upstream-shared", "same") .body(Body::from(body)) .unwrap() } @@ -5763,6 +5768,148 @@ async fn gateway_concurrent_successes_keep_the_selected_attempt_headers() { } } +/// Middleware reads the first event of both streams, then returns the selected stream with its +/// first event replayed. The headers must follow the returned content, not the last stream read. +#[tokio::test] +async fn gateway_streaming_headers_follow_the_returned_stream_after_both_are_read() { + use futures_util::StreamExt; + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = crate::test_support::plugin_host::test_close_plugin_host(); + + const INTERCEPT_NAME: &str = "cli-test-concurrent-streaming-pre-read"; + const MARKER: &str = "read both streams before returning one"; + let _ = deregister_llm_stream_execution_intercept(INTERCEPT_NAME); + let _cleanup = LlmStreamExecutionInterceptCleanup(INTERCEPT_NAME); + let release_losing = Arc::new(Semaphore::new(0)); + let release_from_intercept = release_losing.clone(); + register_llm_stream_execution_intercept( + INTERCEPT_NAME, + 1, + Arc::new(move |_name, request, _context, next| { + let release_losing = release_from_intercept.clone(); + Box::pin(async move { + if request.content["messages"][0]["content"].as_str() != Some(MARKER) { + return next(request).await; + } + let mut selected_request = request.clone(); + selected_request.content["attempt"] = json!("selected"); + let mut losing_request = request; + losing_request.content["attempt"] = json!("losing"); + losing_request.content["losing_succeeds"] = json!(true); + let mut selected = next(selected_request).await?; + release_losing.add_permits(1); + let mut losing = next(losing_request).await?; + let selected_first = selected.next().await.expect("selected first event")?; + let losing_first = losing.next().await.expect("losing first event")?; + assert!(losing_first.get("losing").is_some()); + drop(losing); + let replayed = + futures_util::stream::once(async move { Ok(selected_first) }).chain(selected); + Ok(nemo_relay_runtime::api::runtime::LlmJsonStream::new( + replayed, + )) + }) + }), + ) + .unwrap(); + + let upstream = spawn_concurrent_attempt_upstream(release_losing).await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let response = router(config) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": MARKER}], + "stream": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response + .headers() + .get_all("x-upstream-attempt") + .iter() + .collect::>(), + ["selected"] + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(body.contains("selected attempt"), "{body}"); +} + +/// Two attempts return the same body with different headers, as an upstream cache can. No +/// attempt-specific header may be guessed; headers both attempts agree on are kept. +#[tokio::test] +async fn gateway_identical_buffered_successes_keep_only_shared_headers() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = crate::test_support::plugin_host::test_close_plugin_host(); + + const INTERCEPT_NAME: &str = "cli-test-concurrent-buffered-identical"; + const MARKER: &str = "two attempts with the same body"; + let _ = deregister_llm_execution_intercept(INTERCEPT_NAME); + let _cleanup = LlmExecutionInterceptCleanup(INTERCEPT_NAME); + let release_losing = Arc::new(Semaphore::new(0)); + let release_from_intercept = release_losing.clone(); + register_llm_execution_intercept( + INTERCEPT_NAME, + 1, + Arc::new(move |_name, request, _context, next| { + let release_losing = release_from_intercept.clone(); + Box::pin(async move { + if request.content["messages"][0]["content"].as_str() != Some(MARKER) { + return next(request).await; + } + let mut selected_request = request.clone(); + selected_request.content["attempt"] = json!("selected"); + let mut losing_request = request; + losing_request.content["attempt"] = json!("losing"); + losing_request.content["losing_succeeds"] = json!(true); + losing_request.content["same_body"] = json!(true); + let selected = next(selected_request).await; + release_losing.add_permits(1); + assert!(next(losing_request).await.is_ok()); + selected + }) + }), + ) + .unwrap(); + + let upstream = spawn_concurrent_attempt_upstream(release_losing).await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let response = router(config) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": MARKER}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers().get("x-upstream-attempt").is_none()); + assert_eq!(response.headers()["x-upstream-shared"], "same"); +} + #[tokio::test] async fn gateway_concurrent_next_relays_selected_buffered_failure() { let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; From 77a55dd49304d30f1aa701d5ca0254245181cdf7 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 9 Nov 2026 23:32:33 +0530 Subject: [PATCH 5/5] test: cover the first-event tap and the pre-read stream headers The tap observes only the first successful event and passes terminalize and close through. The pre-read stream test also checks the shared header, the gateway content type, and that no losing content reaches the client. Signed-off-by: Somay-kousis --- .../tests/coverage/shared/gateway_tests.rs | 69 +++++++++++++++++++ .../cli/tests/coverage/shared/server_tests.rs | 6 ++ 2 files changed, 75 insertions(+) diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index 2ae61d420..824b813f8 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -2976,6 +2976,75 @@ fn captured_upstream_successes_keep_only_agreed_headers_when_ambiguous() { assert_only_shared_headers(&empty_stream.response_headers(None, "text/event-stream")); } +/// Inner stream that records explicit lifecycle calls. +struct RecordingStream { + items: std::collections::VecDeque>, + closed: Arc, + terminalized: Arc, +} + +impl futures_util::Stream for RecordingStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(self.get_mut().items.pop_front()) + } +} + +impl LlmStreamInner for RecordingStream { + fn terminalize(self: std::pin::Pin<&mut Self>) { + self.terminalized + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + + fn close( + self: std::pin::Pin<&mut Self>, + ) -> std::pin::Pin> + Send + '_>> + { + self.closed + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Box::pin(async { Ok(()) }) + } +} + +/// The tap sees only the first successful event, and leaves the stream's lifecycle untouched. +#[tokio::test] +async fn first_event_tap_observes_the_first_ok_event_and_delegates_lifecycle() { + use futures_util::StreamExt; + let closed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let terminalized = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let inner = LlmJsonStream::from_closeable(RecordingStream { + items: [ + Err(FlowError::Internal("decode".into())), + Ok(json!({"event": 1})), + Ok(json!({"event": 2})), + ] + .into(), + closed: Arc::clone(&closed), + terminalized: Arc::clone(&terminalized), + }); + let observed = Arc::new(Mutex::new(Vec::new())); + let record = Arc::clone(&observed); + let mut tapped = FirstEventTap::wrap( + inner, + Box::new(move |event| record.lock().unwrap().push(event.clone())), + ); + + assert!(tapped.next().await.unwrap().is_err()); + assert!(observed.lock().unwrap().is_empty()); + assert_eq!(tapped.next().await.unwrap().unwrap(), json!({"event": 1})); + assert_eq!(tapped.next().await.unwrap().unwrap(), json!({"event": 2})); + assert_eq!(*observed.lock().unwrap(), [json!({"event": 1})]); + + tapped.terminalize(); + tapped.close().await.unwrap(); + assert_eq!(terminalized.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(closed.load(std::sync::atomic::Ordering::SeqCst), 1); +} + #[tokio::test] async fn streaming_gateway_call_guard_finishes_when_body_is_dropped() { let manager = SessionManager::new(GatewayConfig::default()); diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index e7de5736a..1ea8a3515 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -5843,9 +5843,15 @@ async fn gateway_streaming_headers_follow_the_returned_stream_after_both_are_rea .collect::>(), ["selected"] ); + assert_eq!(response.headers()["x-upstream-shared"], "same"); + assert_eq!( + response.headers()[header::CONTENT_TYPE], + "text/event-stream" + ); let body = response.into_body().collect().await.unwrap().to_bytes(); let body = std::str::from_utf8(&body).unwrap(); assert!(body.contains("selected attempt"), "{body}"); + assert!(!body.contains("losing"), "{body}"); } /// Two attempts return the same body with different headers, as an upstream cache can. No