diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 264833a87..3579088a4 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,158 @@ impl CapturedUpstreamFailures { type CapturedUpstreamFailuresRef = Arc; +/// Provider headers from the successful attempts of one gateway call. +/// +/// 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>, +} + +struct CapturedUpstreamSuccess { + key: Option, + headers: HeaderMap, +} + +impl CapturedUpstreamSuccesses { + fn capture(&self, key: Option, headers: HeaderMap) -> usize { + let mut entries = self + .entries + .lock() + .expect("captured upstream successes lock poisoned"); + entries.push(CapturedUpstreamSuccess { key, headers }); + entries.len() - 1 + } + + fn set_key(&self, index: usize, key: Value) { + if let Some(entry) = self + .entries + .lock() + .expect("captured upstream successes lock poisoned") + .get_mut(index) + { + entry.key = Some(key); + } + } + + /// 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 mut candidates: Vec = returned + .map(|returned| { + (0..entries.len()) + .filter(|&index| entries[index].key.as_ref() == Some(returned)) + .collect() + }) + .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 + .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.remove("content-digest"); + headers.remove("repr-digest"); + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static(content_type), + ); + headers + } +} + +type CapturedUpstreamSuccessesRef = Arc; + +type FirstEventObserver = Box; + +/// Hands the first successful event of a stream to an observer, without changing close +/// semantics. +struct FirstEventTap { + inner: LlmJsonStream, + observer: Option, +} + +impl FirstEventTap { + fn wrap(inner: LlmJsonStream, observer: FirstEventObserver) -> LlmJsonStream { + LlmJsonStream::from_closeable(Self { + inner, + observer: Some(observer), + }) + } +} + +impl futures_util::Stream for FirstEventTap { + 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(event))) = &polled + && let Some(observer) = this.observer.take() + { + observer(event); + } + polled + } +} + +impl LlmStreamInner for FirstEventTap { + 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 +808,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 +848,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 +889,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 +911,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 +984,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 +1026,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,13 +1105,18 @@ 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 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) => { @@ -963,6 +1126,14 @@ async fn run_managed_streaming( return Err(translate_runtime_error(error)); } }; + // 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) } @@ -1031,6 +1202,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 +1224,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 +1296,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(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 aa6f0eee7..824b813f8 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,185 @@ fn captured_upstream_failures_reject_ambiguous_attempts_without_consuming_them() assert!(failures.take(&selected).is_some()); } +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, 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"); + 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()); +} + +/// 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); + } + + // 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"}); + + // 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")); +} + +/// 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 1ded4a659..1ea8a3515 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -5400,6 +5400,25 @@ 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 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"); + response.to_string() + }; + return Response::builder() + .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(); + } let (status, content_type, body) = match (attempt, streaming, both_fail) { ("selected", false, true) => ( @@ -5449,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() } @@ -5543,9 +5563,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 +5653,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 +5668,254 @@ 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}"); + } +} + +/// 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"] + ); + 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 +/// 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; @@ -7196,6 +7474,118 @@ 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("content-encoding", "identity") + .header("repr-digest", "sha-256=:AAAA:") + .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().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(), + "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};