From a89964b8035c97d1490c524bee5c368341b6c567 Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Fri, 7 Aug 2026 16:43:44 -0700 Subject: [PATCH 1/2] fix: preserve successful terminal stream status Signed-off-by: Ajay Thorve --- crates/core/src/stream.rs | 63 ++++++++++++----- crates/core/tests/integration/stream_tests.rs | 68 +++++++++++++++++++ crates/core/tests/unit/stream_tests.rs | 8 +++ 3 files changed, 121 insertions(+), 18 deletions(-) diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 4679b2751..3be672a2f 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -48,7 +48,9 @@ use crate::api::runtime::{LlmSanitizeResponseContext, LlmSanitizeResponseFn}; use crate::api::shared::{ metadata_with_otel_error, metadata_with_otel_status, snapshot_event_sanitizers, }; -use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider}; +use crate::codec::response::{ + AnnotatedLlmResponse, FinishReason, attach_estimated_cost_for_provider, +}; use crate::codec::traits::LlmResponseCodec; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -86,6 +88,19 @@ pub struct LlmStreamWrapper { terminal_result: Option>, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum StreamTermination { + Complete, + Failed, + Dropped, +} + +impl StreamTermination { + const fn is_interrupted(self) -> bool { + !matches!(self, Self::Complete) + } +} + impl LlmStreamWrapper { /// Create a new `LlmStreamWrapper` around the given raw stream. /// @@ -196,33 +211,28 @@ impl LlmStreamWrapper { self.handle .optimization_recorder .close_for_finalization(None); - self.finalization = self.emit_end_event(metadata, true, background_thread); + self.finalization = + self.emit_end_event(metadata, StreamTermination::Dropped, background_thread); } - fn finish_with_status( - &mut self, - status_code: &'static str, - status_message: Option, - interrupted: bool, - ) { + fn finish_cleanly(&mut self) { if self.ended { return; } self.ended = true; self.inner.terminalize(); - let metadata = - metadata_with_otel_status(self.metadata.clone(), status_code, status_message); - self.finalization = self.emit_end_event(metadata, interrupted, false); + let metadata = metadata_with_otel_status(self.metadata.clone(), "OK", None); + self.finalization = self.emit_end_event(metadata, StreamTermination::Complete, false); } - fn finish_with_error(&mut self, error: &FlowError, interrupted: bool) { + fn finish_with_error(&mut self, error: &FlowError) { if self.ended { return; } self.ended = true; self.inner.terminalize(); let metadata = metadata_with_otel_error(self.metadata.clone(), error); - self.finalization = self.emit_end_event(metadata, interrupted, false); + self.finalization = self.emit_end_event(metadata, StreamTermination::Failed, false); } /// Emit the LLM END event with aggregated response data. @@ -232,7 +242,7 @@ impl LlmStreamWrapper { fn emit_end_event( &mut self, metadata: Option, - interrupted: bool, + termination: StreamTermination, background_thread: bool, ) -> Option> { // The finalizer below runs on the caller's Tokio runtime. Register a @@ -286,7 +296,14 @@ impl LlmStreamWrapper { }) }) .flatten(); - let interruption = (interrupted + let metadata = if termination == StreamTermination::Dropped + && has_authoritative_successful_completion(annotated_response.as_ref()) + { + metadata_with_otel_status(metadata, "OK", None) + } else { + metadata + }; + let interruption = (termination.is_interrupted() && !has_authoritative_final_usage(annotated_response.as_ref())) .then_some("stream_interrupted"); handle @@ -460,19 +477,19 @@ impl Stream for LlmStreamWrapper { match (this.collector)(raw_chunk.clone()) { Ok(()) => Poll::Ready(Some(Ok(raw_chunk))), Err(e) => { - this.finish_with_error(&e, true); + this.finish_with_error(&e); this.terminal_result = Some(Err(e)); self.poll_next(cx) } } } Poll::Ready(Some(Err(e))) => { - this.finish_with_error(&e, true); + this.finish_with_error(&e); this.terminal_result = Some(Err(e)); self.poll_next(cx) } Poll::Ready(None) => { - this.finish_with_status("OK", None, false); + this.finish_cleanly(); self.poll_next(cx) } Poll::Pending => Poll::Pending, @@ -513,6 +530,16 @@ fn has_authoritative_final_usage(response: Option<&AnnotatedLlmResponse>) -> boo }) } +fn has_authoritative_successful_completion(response: Option<&AnnotatedLlmResponse>) -> bool { + has_authoritative_final_usage(response) + && response.is_some_and(|response| { + response + .finish_reason + .as_ref() + .is_some_and(|reason| !matches!(reason, FinishReason::Unknown(_))) + }) +} + fn llm_chunk_mark_data(chunk_index: u64, raw_chunk: &Json) -> Json { if let Some(data) = summarize_openai_chat_chunk(chunk_index, raw_chunk) { return data; diff --git a/crates/core/tests/integration/stream_tests.rs b/crates/core/tests/integration/stream_tests.rs index 3f3ab78ca..0b66a0ab2 100644 --- a/crates/core/tests/integration/stream_tests.rs +++ b/crates/core/tests/integration/stream_tests.rs @@ -18,7 +18,9 @@ use nemo_relay::api::optimization::LlmOptimizationRecorder; use nemo_relay::api::runtime::global_context; use nemo_relay::api::runtime::{LlmJsonStream, LlmStreamInner, NemoRelayContextState}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::codec::openai_responses::{OpenAIResponsesCodec, OpenAIResponsesStreamingCodec}; use nemo_relay::codec::optimization::LlmOptimizationContribution; +use nemo_relay::codec::streaming::StreamingCodec; use nemo_relay::error::FlowError; use nemo_relay::error::Result; use nemo_relay::json::Json; @@ -471,6 +473,72 @@ async fn test_stream_wrapper_drop_emits_end_event_for_partial_stream() { deregister_subscriber("stream_drop_end_test").unwrap(); } +#[tokio::test] +async fn dropped_stream_after_terminal_response_emits_success() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + + let events = Arc::new(Mutex::new(Vec::new())); + let captured = events.clone(); + register_subscriber( + "stream_terminal_drop_status_test", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let terminal_event = json!({ + "type": "response.completed", + "response": { + "id": "resp_complete", + "model": "gpt-5", + "status": "completed", + "output": [{ + "type": "message", + "content": [{"type": "output_text", "text": "done"}] + }], + "usage": {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12} + } + }); + let streaming_codec = OpenAIResponsesStreamingCodec::new(); + let collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({"input": "finish"}), + }; + let handle = llm_call( + LlmCallParams::builder() + .name("openai.responses") + .request(&request) + .attributes(LlmAttributes::STREAMING) + .build(), + ) + .unwrap(); + let mut wrapper = LlmStreamWrapper::new( + make_stream(vec![Ok(terminal_event.clone())]), + handle, + collector, + finalizer, + None, + None, + Some(Arc::new(OpenAIResponsesCodec)), + ); + + assert_eq!(wrapper.next().await.unwrap().unwrap(), terminal_event); + drop(wrapper); + + let events = captured_snapshot(&events); + let end_event = events + .iter() + .find(|event| is_llm_end(event)) + .expect("expected END event after the terminal response was dropped"); + let metadata = end_event.metadata().unwrap(); + assert_eq!(metadata["otel.status_code"], json!("OK")); + assert!(metadata.get("otel.status_description").is_none()); + + deregister_subscriber("stream_terminal_drop_status_test").unwrap(); +} + #[tokio::test] async fn stream_termination_modes_close_accounting_without_losing_evidence() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/unit/stream_tests.rs b/crates/core/tests/unit/stream_tests.rs index 0eb68537d..e3cd90e1d 100644 --- a/crates/core/tests/unit/stream_tests.rs +++ b/crates/core/tests/unit/stream_tests.rs @@ -35,6 +35,14 @@ fn partial_stream_usage_is_not_treated_as_authoritative_without_terminal_evidenc ..partial }; assert!(has_authoritative_final_usage(Some(&terminal))); + assert!(has_authoritative_successful_completion(Some(&terminal))); + + let failed = AnnotatedLlmResponse { + finish_reason: Some(FinishReason::Unknown("failed".to_string())), + ..terminal + }; + assert!(has_authoritative_final_usage(Some(&failed))); + assert!(!has_authoritative_successful_completion(Some(&failed))); } #[test] From c736871cf1d6f42d08cfa1b755848c01f34b9614 Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Fri, 7 Aug 2026 17:00:03 -0700 Subject: [PATCH 2/2] test: clarify terminal stream outcome semantics Signed-off-by: Ajay Thorve --- crates/core/src/stream.rs | 4 ++-- crates/core/tests/unit/stream_tests.rs | 23 +++++++++++++++-------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 3be672a2f..d2defea3e 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -297,7 +297,7 @@ impl LlmStreamWrapper { }) .flatten(); let metadata = if termination == StreamTermination::Dropped - && has_authoritative_successful_completion(annotated_response.as_ref()) + && has_authoritative_terminal_outcome(annotated_response.as_ref()) { metadata_with_otel_status(metadata, "OK", None) } else { @@ -530,7 +530,7 @@ fn has_authoritative_final_usage(response: Option<&AnnotatedLlmResponse>) -> boo }) } -fn has_authoritative_successful_completion(response: Option<&AnnotatedLlmResponse>) -> bool { +fn has_authoritative_terminal_outcome(response: Option<&AnnotatedLlmResponse>) -> bool { has_authoritative_final_usage(response) && response.is_some_and(|response| { response diff --git a/crates/core/tests/unit/stream_tests.rs b/crates/core/tests/unit/stream_tests.rs index e3cd90e1d..548690bab 100644 --- a/crates/core/tests/unit/stream_tests.rs +++ b/crates/core/tests/unit/stream_tests.rs @@ -30,19 +30,26 @@ fn partial_stream_usage_is_not_treated_as_authoritative_without_terminal_evidenc }; assert!(!has_authoritative_final_usage(Some(&partial))); - let terminal = AnnotatedLlmResponse { - finish_reason: Some(FinishReason::Complete), - ..partial - }; - assert!(has_authoritative_final_usage(Some(&terminal))); - assert!(has_authoritative_successful_completion(Some(&terminal))); + for finish_reason in [ + FinishReason::Complete, + FinishReason::Length, + FinishReason::ToolUse, + FinishReason::ContentFilter, + ] { + let terminal = AnnotatedLlmResponse { + finish_reason: Some(finish_reason), + ..partial.clone() + }; + assert!(has_authoritative_final_usage(Some(&terminal))); + assert!(has_authoritative_terminal_outcome(Some(&terminal))); + } let failed = AnnotatedLlmResponse { finish_reason: Some(FinishReason::Unknown("failed".to_string())), - ..terminal + ..partial }; assert!(has_authoritative_final_usage(Some(&failed))); - assert!(!has_authoritative_successful_completion(Some(&failed))); + assert!(!has_authoritative_terminal_outcome(Some(&failed))); } #[test]