diff --git a/crates/adaptive/tests/integration/runtime_integration_tests.rs b/crates/adaptive/tests/integration/runtime_integration_tests.rs index 5b68723a2..16da0444f 100644 --- a/crates/adaptive/tests/integration/runtime_integration_tests.rs +++ b/crates/adaptive/tests/integration/runtime_integration_tests.rs @@ -15,7 +15,9 @@ use nemo_relay::api::llm::{ }; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn}; +use nemo_relay::api::runtime::{ + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, +}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::api::tool::tool_call_execute; use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; @@ -775,9 +777,7 @@ impl Plugin for HeaderPlugin { } chunks.push(Ok(chunk)); } - let stream = Box::pin(tokio_stream::iter(chunks)) - as Pin> + Send>>; - Ok(stream) + Ok(LlmJsonStream::new(tokio_stream::iter(chunks))) }) }), )?; @@ -853,10 +853,7 @@ async fn test_top_level_plugin_registers_request_and_execution_intercepts() { let llm_stream_func: LlmStreamExecutionNextFn = Arc::new(|_req: LlmRequest| { Box::pin(async move { let chunks = vec![Ok(json!({"streamed": true}))]; - Ok(Box::pin(tokio_stream::iter(chunks)) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(chunks))) }) }); let collected = Arc::new(StdMutex::new(Vec::new())); diff --git a/crates/adaptive/tests/unit/acg_component_tests.rs b/crates/adaptive/tests/unit/acg_component_tests.rs index 04e601381..fcf09c01a 100644 --- a/crates/adaptive/tests/unit/acg_component_tests.rs +++ b/crates/adaptive/tests/unit/acg_component_tests.rs @@ -17,8 +17,7 @@ use crate::config::AcgComponentConfig; use crate::storage::memory::InMemoryBackend; use crate::storage::traits::StorageBackendDyn; use nemo_relay::api::llm::LlmRequest; -use nemo_relay::api::runtime::LlmExecutionNextFn; -use nemo_relay::api::runtime::LlmStreamExecutionNextFn; +use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn}; use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; use serde_json::{Value, json}; use tokio_stream::StreamExt; @@ -646,10 +645,9 @@ async fn acg_component_stream_execution_intercept_rewrites_streaming_requests() ); let next: LlmStreamExecutionNextFn = Arc::new(|req| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(req.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(req.content), + ]))) }) }); @@ -1295,10 +1293,9 @@ async fn acg_component_stream_execution_intercept_passes_original_request_when_t ); let next: LlmStreamExecutionNextFn = Arc::new(|req| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(req.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(req.content), + ]))) }) }); diff --git a/crates/adaptive/tests/unit/runtime_features_tests.rs b/crates/adaptive/tests/unit/runtime_features_tests.rs index 2a45418ec..5a53c7734 100644 --- a/crates/adaptive/tests/unit/runtime_features_tests.rs +++ b/crates/adaptive/tests/unit/runtime_features_tests.rs @@ -27,6 +27,7 @@ use nemo_relay::api::registry::{ register_llm_execution_intercept, register_llm_request_intercept, register_llm_stream_execution_intercept, register_tool_execution_intercept, }; +use nemo_relay::api::runtime::LlmJsonStream; use nemo_relay::api::runtime::ToolExecutionNextFn; use nemo_relay::api::runtime::global_context; use nemo_relay::api::runtime::{ @@ -746,14 +747,9 @@ async fn registration_context_registers_all_supported_callback_types() { 7, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin< - Box< - dyn tokio_stream::Stream< - Item = nemo_relay::error::Result, - > + Send, - >, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) @@ -878,8 +874,7 @@ async fn acg_feature_registers_execution_and_stream_intercepts() { let stream_next: LlmStreamExecutionNextFn = Arc::new(|request| { Box::pin(async move { - let stream: nemo_relay::api::runtime::LlmJsonStream = - Box::pin(tokio_stream::iter(vec![Ok(request.content)])); + let stream = LlmJsonStream::new(tokio_stream::iter(vec![Ok(request.content)])); Ok(stream) }) }); diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index edbbf1cee..0eefae373 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -534,7 +534,7 @@ fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream { Err(error) => yield Err(error), } }; - Box::pin(stream) + LlmJsonStream::new(stream) } // Re-encodes a runtime JSON stream as `text/event-stream` frames for the downstream client. Event diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index 246073441..1024590ff 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -1776,7 +1776,7 @@ async fn streaming_gateway_call_guard_finishes_when_body_is_dropped() { .await .unwrap(); - let stream: LlmJsonStream = Box::pin(futures_util::stream::pending::< + let stream = LlmJsonStream::new(futures_util::stream::pending::< std::result::Result, >()); let body = client_sse_body( @@ -1856,7 +1856,7 @@ fn streaming_gateway_call_guard_finishes_without_a_current_runtime() { (manager, prep) }); let final_response = json!({ "output_text": "streamed final" }); - let stream: LlmJsonStream = Box::pin(futures_util::stream::pending::< + let stream = LlmJsonStream::new(futures_util::stream::pending::< std::result::Result, >()); let body = client_sse_body( @@ -1934,7 +1934,7 @@ async fn streaming_body_records_final_response_for_turn_output() { let session_id = prep.session_id.clone(); let owner_subagent_id = prep.owner_subagent_id.clone(); let final_response = json!({ "output_text": "streamed final" }); - let stream: LlmJsonStream = Box::pin(futures_util::stream::empty::< + let stream = LlmJsonStream::new(futures_util::stream::empty::< std::result::Result, >()); let body = client_sse_body( diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 9e289f948..1a6e66dac 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -1173,7 +1173,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu response_codec, lifecycle_subscribers, ); - Ok(Box::pin(wrapper) as LlmJsonStream) + Ok(LlmJsonStream::from_closeable(wrapper)) } Err(error) => { let end_metadata = diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 7d5a12882..8b02d998b 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -12,8 +12,8 @@ pub mod subscriber_dispatcher; pub use callbacks::{ EventSanitizeFn, EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, - LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn, - ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamInner, + ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; pub use global::global_context; pub use scope_stack::{ diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index ad81932d0..21fe853f1 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -11,6 +11,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use tokio_stream::Stream; @@ -240,7 +241,85 @@ pub type LlmExecutionFn = Arc< + Sync, >; /// Stream of JSON chunks produced by the managed streaming LLM pipeline. -pub type LlmJsonStream = Pin> + Send>>; +/// +/// In addition to ordinary stream polling, managed streams provide an explicit +/// asynchronous close operation. A successful close means the producer has +/// released its resources; subsequent polls return no more chunks. +pub struct LlmJsonStream { + inner: Pin>, +} + +impl LlmJsonStream { + /// Wrap a stream whose producer has no asynchronous teardown work. + pub fn new(stream: S) -> Self + where + S: Stream> + Send + 'static, + { + Self { + inner: Box::pin(DefaultLlmStream { + stream: Some(Box::pin(stream)), + }), + } + } + + /// Wrap a stream that implements explicit asynchronous teardown. + pub fn from_closeable(stream: S) -> Self + where + S: LlmStreamInner + 'static, + { + Self { + inner: Box::pin(stream), + } + } + + /// Stop the producer and wait for its cleanup to complete. + pub async fn close(&mut self) -> Result<()> { + self.inner.as_mut().close().await + } +} + +impl Stream for LlmJsonStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.get_mut().inner.as_mut().poll_next(cx) + } +} + +/// Internal close-aware stream implementation. +pub trait LlmStreamInner: Stream> + Send { + /// Stop the producer and wait for cleanup. Implementations must be idempotent. + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>>; +} + +struct DefaultLlmStream { + stream: Option>>, +} + +impl Stream for DefaultLlmStream +where + S: Stream> + Send, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.stream.as_mut() { + Some(stream) => stream.as_mut().poll_next(cx), + None => Poll::Ready(None), + } + } +} + +impl LlmStreamInner for DefaultLlmStream +where + S: Stream> + Send, +{ + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + self.get_mut().stream.take(); + Box::pin(async { Ok(()) }) + } +} /// Per-chunk collector used by the streaming LLM runtime. /// /// # Parameters diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 70d49aa49..988e5fb89 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -82,7 +82,7 @@ impl std::fmt::Display for UpstreamFailure { /// /// Each variant represents a distinct failure mode that callers can match on /// to determine the appropriate recovery strategy. -#[derive(Debug, Error)] +#[derive(Clone, Debug, Error)] pub enum FlowError { /// A resource with the given name is already registered. /// diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 79c9d82eb..1aeffbd1a 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -2159,11 +2159,11 @@ fn native_stream_to_relay_stream( next_ctx: Option, callback_user_data: Option>, ) -> FlowResult { - Ok(Box::pin(NativeRelayLlmStream::from_raw( + Ok(LlmJsonStream::new(NativeRelayLlmStream::from_raw( raw, next_ctx, callback_user_data, - )?) as LlmJsonStream) + )?)) } fn drop_native_stream(mut raw: NemoRelayNativeLlmStreamV1) { diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 9183eaaa8..ed967482f 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -1745,7 +1745,9 @@ impl WorkerPluginCallback { } guard.finish(); }); - Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx))) + Ok(LlmJsonStream::new( + tokio_stream::wrappers::ReceiverStream::new(rx), + )) } fn base_request( diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index 485467ba2..04b99c01c 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -4,22 +4,27 @@ use std::collections::HashMap; use std::env; use std::ffi::{OsStr, OsString}; +use std::future::Future; use std::io::{BufRead, BufReader, Write}; +use std::pin::Pin; use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, mpsc as std_mpsc}; +use std::task::{Context, Poll}; use std::thread; use std::time::Duration; use serde::Deserialize; use serde_json::json; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tokio::task::JoinHandle; use tokio_stream::StreamExt; use tokio_stream::wrappers::ReceiverStream; use crate::api::llm::LlmRequest; -use crate::api::runtime::{LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, ToolExecutionFn}; +use crate::api::runtime::{ + LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, LlmStreamInner, ToolExecutionFn, +}; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; use crate::codec::resolve::{ProviderSurface, request_codec, response_codec}; use crate::error::{FlowError, Result as FlowResult}; @@ -279,6 +284,8 @@ impl LocalGuardrailsRuntime { .spawn_stream_monitor(messages, text_rx, Arc::clone(&blocked))?; let codec = *self.codec()?; + let (cancel, cancel_rx) = watch::channel(false); + let (closed, closed_rx) = watch::channel(None); tokio::spawn(async move { forward_guarded_provider_stream( provider_stream, @@ -287,11 +294,17 @@ impl LocalGuardrailsRuntime { chunk_tx, monitor, blocked, + cancel_rx, + closed, ) .await; }); - Ok(Box::pin(ReceiverStream::new(chunk_rx)) as LlmJsonStream) + Ok(LlmJsonStream::from_closeable(GuardedProviderStream { + receiver: ReceiverStream::new(chunk_rx), + cancel, + closed: closed_rx, + })) } fn codec(&self) -> FlowResult<&LocalGuardrailsCodec> { @@ -1056,6 +1069,45 @@ fn local_violation(message: impl Into) -> FlowError { FlowError::Internal(message.into()) } +struct GuardedProviderStream { + receiver: ReceiverStream>, + cancel: watch::Sender, + closed: watch::Receiver>>, +} + +impl tokio_stream::Stream for GuardedProviderStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.receiver).poll_next(cx) + } +} + +impl Drop for GuardedProviderStream { + fn drop(&mut self) { + self.cancel.send_replace(true); + } +} + +impl LlmStreamInner for GuardedProviderStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let this = self.get_mut(); + this.cancel.send_replace(true); + this.receiver.close(); + while this.receiver.as_mut().try_recv().is_ok() {} + let mut closed = this.closed.clone(); + Box::pin(async move { + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + FlowError::Internal("guarded stream cleanup task ended early".into()) + })?; + } + closed.borrow().clone().expect("close state checked above") + }) + } +} + +#[allow(clippy::too_many_arguments)] async fn forward_guarded_provider_stream( mut provider_stream: LlmJsonStream, codec: LocalGuardrailsCodec, @@ -1063,39 +1115,68 @@ async fn forward_guarded_provider_stream( chunk_tx: mpsc::Sender>, monitor: JoinHandle>, blocked: Arc>>, + mut cancel: watch::Receiver, + closed: watch::Sender>>, ) { - while let Some(item) = provider_stream.next().await { + let mut monitor = Some(monitor); + loop { + if *cancel.borrow() { + break; + } + let item = tokio::select! { + _ = cancel.changed() => break, + item = provider_stream.next() => item, + }; + let Some(item) = item else { + break; + }; let chunk = match item { Ok(chunk) => chunk, Err(err) => { let _ = chunk_tx.send(Err(err)).await; let _ = text_tx.send(None).await; - let _ = monitor.await; - return; + let _ = monitor.take().expect("monitor available").await; + break; } }; if let Some(message) = blocked_message(&blocked) { let _ = chunk_tx.send(Err(streaming_output_blocked(message))).await; let _ = text_tx.send(None).await; - let _ = monitor.await; - return; + let _ = monitor.take().expect("monitor available").await; + break; } if let Some(text) = extract_stream_text(codec, &chunk) && text_tx.send(Some(text)).await.is_err() { - send_stream_monitor_error(monitor, &chunk_tx, &blocked).await; - return; + send_stream_monitor_error( + monitor.take().expect("monitor available"), + &chunk_tx, + &blocked, + ) + .await; + break; } - if chunk_tx.send(Ok(chunk)).await.is_err() { + let sent = tokio::select! { + _ = cancel.changed() => break, + sent = chunk_tx.send(Ok(chunk)) => sent, + }; + if sent.is_err() { let _ = text_tx.send(None).await; - let _ = monitor.await; - return; + let _ = monitor.take().expect("monitor available").await; + break; } } let _ = text_tx.send(None).await; - let _ = send_stream_monitor_error(monitor, &chunk_tx, &blocked).await; + if *cancel.borrow() { + if let Some(monitor) = monitor.take() { + monitor.abort(); + } + } else if let Some(monitor) = monitor.take() { + let _ = send_stream_monitor_error(monitor, &chunk_tx, &blocked).await; + } + closed.send_replace(Some(provider_stream.close().await)); } async fn send_stream_monitor_error( diff --git a/crates/core/src/plugins/nemo_guardrails/remote.rs b/crates/core/src/plugins/nemo_guardrails/remote.rs index 2ed41dd0a..b6eeeadc4 100644 --- a/crates/core/src/plugins/nemo_guardrails/remote.rs +++ b/crates/core/src/plugins/nemo_guardrails/remote.rs @@ -182,7 +182,7 @@ impl RemoteBackendRuntime { let (tx, rx) = mpsc::channel(16); self.spawn_stream_decoder(response, status, parent.clone(), tx); - Ok(Box::pin(ReceiverStream::new(rx)) as LlmJsonStream) + Ok(LlmJsonStream::new(ReceiverStream::new(rx))) } fn emit_remote_start(&self, parent: &Option, stream: bool) { diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 6d1d13dfb..1124b3d0d 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -25,6 +25,7 @@ //! aggregated response then flows through sanitize response guardrails before //! being included in the END event. +use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -37,7 +38,9 @@ use crate::api::llm::emit_optimization_marks; use crate::api::optimization::finalize_optimization_summary; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; -use crate::api::runtime::{EventSubscriberFn, ScopeStackHandle, current_scope_stack}; +use crate::api::runtime::{ + EventSubscriberFn, LlmJsonStream, LlmStreamInner, ScopeStackHandle, current_scope_stack, +}; use crate::api::shared::metadata_with_otel_status; use crate::api::shared::sanitize_event_with_scope_stack; use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider}; @@ -50,16 +53,19 @@ use serde_json::Map; /// /// 1. Passes each chunk to the user-supplied **collector** closure. /// If the collector returns `Err`, the stream terminates with that error. -/// 2. On stream exhaustion, calls the **finalizer** to produce an aggregated -/// [`Json`] response, runs sanitize response guardrails on it, then emits -/// the LLM END event. +/// 2. On stream exhaustion or explicit close, calls the **finalizer** to +/// produce an aggregated [`Json`] response, runs sanitize response +/// guardrails on it, then emits the LLM END event. Explicit close marks the +/// end event as interrupted and waits for producer cleanup. /// /// This type is returned by [`crate::api::llm::llm_stream_call_execute`] and -/// is usually consumed as an ordinary async stream. The wrapper preserves the -/// originating scope stack so end-of-stream bookkeeping still uses the correct -/// scope-local middleware and subscribers even when polling happens elsewhere. +/// is usually consumed as an ordinary async stream. Consumers that stop early +/// should call [`LlmJsonStream::close`] to perform deterministic cleanup. The +/// wrapper preserves the originating scope stack so end-of-stream bookkeeping +/// still uses the correct scope-local middleware and subscribers even when +/// polling happens elsewhere. pub struct LlmStreamWrapper { - inner: Pin> + Send>>, + inner: LlmJsonStream, handle: LlmHandle, scope_stack: ScopeStackHandle, collector: Box Result<()> + Send>, @@ -69,6 +75,7 @@ pub struct LlmStreamWrapper { subscribers: Vec, chunk_index: u64, ended: bool, + close_result: Option>, } impl LlmStreamWrapper { @@ -94,7 +101,7 @@ impl LlmStreamWrapper { /// # Returns /// A new [`LlmStreamWrapper`] ready to be polled. pub fn new( - inner: Pin> + Send>>, + inner: LlmJsonStream, handle: LlmHandle, collector: Box Result<()> + Send>, finalizer: Box Json + Send>, @@ -124,7 +131,7 @@ impl LlmStreamWrapper { } pub(crate) fn new_managed( - inner: Pin> + Send>>, + inner: LlmJsonStream, handle: LlmHandle, collector: Box Result<()> + Send>, finalizer: Box Json + Send>, @@ -144,6 +151,7 @@ impl LlmStreamWrapper { subscribers, chunk_index: 0, ended: false, + close_result: None, } } @@ -308,7 +316,7 @@ impl Stream for LlmStreamWrapper { } // Poll the inner stream - match this.inner.as_mut().poll_next(cx) { + match Pin::new(&mut this.inner).poll_next(cx) { Poll::Ready(Some(Ok(raw_chunk))) => { let chunk_index = this.chunk_index; this.chunk_index += 1; @@ -337,6 +345,24 @@ impl Stream for LlmStreamWrapper { } } +impl LlmStreamInner for LlmStreamWrapper { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let this = self.get_mut(); + Box::pin(async move { + if let Some(result) = &this.close_result { + return result.clone(); + } + let result = this.inner.close().await; + this.finish(); + this.close_result = Some(result.clone()); + this.close_result + .as_ref() + .expect("close result was just stored") + .clone() + }) + } +} + fn has_authoritative_final_usage(response: Option<&AnnotatedLlmResponse>) -> bool { response.is_some_and(|response| { response.finish_reason.is_some() diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index c508cc484..a8eb0536d 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -5,7 +5,6 @@ #![allow(clippy::await_holding_lock)] -use std::pin::Pin; use std::sync::{Arc, Mutex}; use chrono::{DateTime, TimeDelta, Utc}; @@ -51,7 +50,9 @@ use nemo_relay::api::registry::{ }; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn}; +use nemo_relay::api::runtime::{ + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, +}; use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::ScopeType; use nemo_relay::api::scope::{event, pop_scope, push_scope}; @@ -67,7 +68,6 @@ use nemo_relay::api::tool::{ use nemo_relay::error::{FlowError, Result}; use nemo_relay::json::Json; use serde_json::{Map, json}; -use tokio_stream::Stream; static TEST_MUTEX: Mutex<()> = Mutex::new(()); @@ -763,8 +763,9 @@ fn failing_llm_exec() -> LlmExecutionNextFn { fn noop_llm_stream_exec() -> LlmStreamExecutionNextFn { Arc::new(|request| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin> + Send>>) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }) } @@ -775,8 +776,7 @@ fn fixed_llm_stream_exec(chunks: Vec) -> LlmStreamExecutionNextFn { let chunks = chunks.clone(); Box::pin(async move { let items = chunks.iter().cloned().map(Ok).collect::>(); - Ok(Box::pin(tokio_stream::iter(items)) - as Pin> + Send>>) + Ok(LlmJsonStream::new(tokio_stream::iter(items))) }) }) } @@ -908,8 +908,9 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin> + Send>>) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) @@ -1153,8 +1154,9 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin> + Send>>) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index b24105e78..5acc08b63 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -2753,7 +2753,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { assert_middleware_callback_locks_are_free(); Box::pin(async move { let stream = tokio_stream::iter(vec![Ok(json!({"chunk": true}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }); let tracked = callbacks.clone(); @@ -3738,10 +3738,9 @@ async fn test_stream_optimization_mark_uses_the_llm_captured_sanitizer_scope() { }) .func(Arc::new(|_| { Box::pin(async { - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": "done"}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": "done" + }))]))) }) })) .collector(Box::new(|_| Ok(()))) @@ -4130,7 +4129,7 @@ async fn test_llm_stream_start_emits_before_short_circuit_execution_intercept() .unwrap() .insert("phase".into(), json!("execution")); let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "short-circuited"}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }), ) @@ -4142,7 +4141,7 @@ async fn test_llm_stream_start_emits_before_short_circuit_execution_intercept() oc.store(true, Ordering::SeqCst); Box::pin(async move { let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }); diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index ca1f3cdec..b7a64ed0d 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -554,13 +554,13 @@ async fn sdk_cdylib_registers_tool_request_intercept() { }) .func(Arc::new(|request| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![ + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ Ok(json!({ "stream_chunk": 1, "request": request.content, })), Ok(json!({ "stream_chunk": 2 })), - ])) as LlmJsonStream) + ]))) }) })) .collector(Box::new(move |chunk| { diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 53f0c7b14..b44e68056 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -7,12 +7,10 @@ #![allow(clippy::await_holding_lock)] -use std::pin::Pin; use std::sync::{Arc, Mutex}; use futures::StreamExt; use serde_json::json; -use tokio_stream::Stream; use nemo_relay::api::event::{Event, PendingMarkSpec, ScopeCategory}; use nemo_relay::api::llm::{ @@ -329,9 +327,9 @@ fn noop_exec_fn() -> LlmExecutionNextFn { fn noop_stream_exec_fn() -> LlmStreamExecutionNextFn { Arc::new(|_req| { Box::pin(async { - let stream: Pin> + Send>> = - Box::pin(futures::stream::once(async { Ok(json!({"chunk": 1})) })); - Ok(stream) + Ok(LlmJsonStream::new(futures::stream::once(async { + Ok(json!({"chunk": 1})) + }))) }) }) } @@ -596,11 +594,7 @@ async fn test_stream_codec_rejects_raw_content_mutation_before_lifecycle() { let called = provider_called.clone(); let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { *called.lock().unwrap() = true; - Box::pin(async { - let stream: Pin> + Send>> = - Box::pin(futures::stream::empty()); - Ok(stream) - }) + Box::pin(async { Ok(LlmJsonStream::new(futures::stream::empty())) }) }); let (codec, _, _) = make_tracking_codec("stream_codec_raw_read_only"); let result = llm_stream_call_execute( @@ -1710,10 +1704,9 @@ async fn managed_buffered_and_streaming_close_price_the_committed_route_not_resp assert!(record_llm_optimization_contribution( routed_model_contribution() )); - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": "done"}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": "done" + }))]))) }) })) .collector(Box::new(|_| Ok(()))) diff --git a/crates/core/tests/integration/stream_tests.rs b/crates/core/tests/integration/stream_tests.rs index 3dbbee616..8cd48ecc3 100644 --- a/crates/core/tests/integration/stream_tests.rs +++ b/crates/core/tests/integration/stream_tests.rs @@ -5,15 +5,18 @@ #![allow(clippy::await_holding_lock)] +use std::future::Future; use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; use nemo_relay::api::event::{Event, ScopeCategory}; use nemo_relay::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; use nemo_relay::api::llm::{LlmCallParams, llm_call}; use nemo_relay::api::optimization::LlmOptimizationRecorder; -use nemo_relay::api::runtime::NemoRelayContextState; 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::optimization::LlmOptimizationContribution; use nemo_relay::error::FlowError; @@ -51,8 +54,54 @@ fn make_optimized_llm_handle(name: &str, producer: &str) -> (LlmHandle, LlmOptim (handle, recorder) } -fn make_stream(items: Vec>) -> Pin> + Send>> { - Box::pin(tokio_stream::iter(items)) +fn make_stream(items: Vec>) -> LlmJsonStream { + LlmJsonStream::new(tokio_stream::iter(items)) +} + +struct CloseTrackingStream { + stream: Pin> + Send>>, + close_calls: Arc, + close_error: Option, + closed: bool, +} + +impl Stream for CloseTrackingStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.closed { + Poll::Ready(None) + } else { + self.stream.as_mut().poll_next(cx) + } + } +} + +impl LlmStreamInner for CloseTrackingStream { + fn close(mut self: Pin<&mut Self>) -> Pin> + Send + '_>> { + self.closed = true; + self.close_calls.fetch_add(1, Ordering::SeqCst); + let close_error = self.close_error.clone(); + Box::pin(async move { close_error.map_or(Ok(()), Err) }) + } +} + +struct DropTrackingStream { + drops: Arc, +} + +impl Stream for DropTrackingStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Pending + } +} + +impl Drop for DropTrackingStream { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } } fn captured_snapshot(items: &Arc>>) -> Vec { @@ -105,6 +154,107 @@ async fn test_stream_wrapper_basic_chunks() { assert_eq!(chunks[1]["token"], "world"); } +#[tokio::test] +async fn explicit_close_finalizes_once_and_exhausts_the_managed_stream() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "explicit_close_finalizes_once", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let close_calls = Arc::new(AtomicUsize::new(0)); + let finalizer_calls = Arc::new(AtomicUsize::new(0)); + let finalizer_count = Arc::clone(&finalizer_calls); + let inner = LlmJsonStream::from_closeable(CloseTrackingStream { + stream: Box::pin(tokio_stream::iter(vec![ + Ok(json!({"chunk": "first"})), + Ok(json!({"chunk": "unread"})), + ])), + close_calls: Arc::clone(&close_calls), + close_error: None, + closed: false, + }); + let wrapper = LlmStreamWrapper::new( + inner, + make_llm_handle("explicit-close"), + Box::new(|_| Ok(())), + Box::new(move || { + finalizer_count.fetch_add(1, Ordering::SeqCst); + json!({"partial": true}) + }), + None, + None, + None, + ); + let mut stream = LlmJsonStream::from_closeable(wrapper); + + assert_eq!(stream.next().await.unwrap().unwrap()["chunk"], "first"); + stream.close().await.unwrap(); + stream.close().await.unwrap(); + + assert_eq!(close_calls.load(Ordering::SeqCst), 1); + assert_eq!(finalizer_calls.load(Ordering::SeqCst), 1); + assert!(stream.next().await.is_none()); + + flush_subscribers().unwrap(); + let end_events = events.lock().unwrap(); + assert_eq!( + end_events.iter().filter(|event| is_llm_end(event)).count(), + 1 + ); + assert!(deregister_subscriber("explicit_close_finalizes_once").unwrap()); +} + +#[tokio::test] +async fn explicit_close_caches_cleanup_errors() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + + let close_calls = Arc::new(AtomicUsize::new(0)); + let inner = LlmJsonStream::from_closeable(CloseTrackingStream { + stream: Box::pin(tokio_stream::empty()), + close_calls: Arc::clone(&close_calls), + close_error: Some(FlowError::NotFound("producer cleanup failed".into())), + closed: false, + }); + let wrapper = LlmStreamWrapper::new( + inner, + make_llm_handle("explicit-close-error"), + Box::new(|_| Ok(())), + Box::new(|| Json::Null), + None, + None, + None, + ); + let mut stream = LlmJsonStream::from_closeable(wrapper); + + for _ in 0..2 { + let error = stream.close().await.expect_err("close should fail"); + assert!(error.to_string().contains("producer cleanup failed")); + assert!(matches!(error, FlowError::NotFound(_))); + } + assert_eq!(close_calls.load(Ordering::SeqCst), 1); + assert!(stream.next().await.is_none()); +} + +#[tokio::test] +async fn default_stream_close_drops_the_wrapped_stream() { + let drops = Arc::new(AtomicUsize::new(0)); + let mut stream = LlmJsonStream::new(DropTrackingStream { + drops: Arc::clone(&drops), + }); + + stream.close().await.unwrap(); + + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(stream.next().await.is_none()); +} + #[tokio::test] async fn test_stream_wrapper_passthrough() { let _lock = TEST_MUTEX.lock().unwrap(); @@ -132,7 +282,7 @@ async fn test_stream_wrapper_empty_stream() { let _lock = TEST_MUTEX.lock().unwrap(); reset_global(); - let inner: Pin> + Send>> = Box::pin(tokio_stream::empty()); + let inner = LlmJsonStream::new(tokio_stream::empty()); let handle = make_llm_handle("test_llm"); let (collector, finalizer, _collected) = make_collector_finalizer(); let mut wrapper = LlmStreamWrapper::new(inner, handle, collector, finalizer, None, None, None); diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs index 8ac588a53..ea3f447c7 100644 --- a/crates/core/tests/integration/worker_plugin_tests.rs +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -13,7 +13,7 @@ use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, }; -use nemo_relay::api::runtime::{TASK_SCOPE_STACK, create_scope_stack}; +use nemo_relay::api::runtime::{LlmJsonStream, TASK_SCOPE_STACK, create_scope_stack}; use nemo_relay::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, }; @@ -343,7 +343,7 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { "chunk": 1, "request": request.content, }); - Ok(Box::pin(tokio_stream::iter(vec![Ok(first)])) as _) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(first)]))) }) })) .collector(Box::new(|_chunk| Ok(()))) @@ -595,7 +595,7 @@ async fn worker_llm_stream_open_error_surfaces_to_host() { .func(Arc::new(|request| { Box::pin(async move { let chunk = json!({ "request": request.content }); - Ok(Box::pin(tokio_stream::iter(vec![Ok(chunk)])) as _) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(chunk)]))) }) })) .collector(Box::new(|_chunk| Ok(()))) diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index a6669254a..7b31169a6 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -550,9 +550,7 @@ async fn callback_stream_transport_error_surfaces_to_host_stream() { "stream_transport_error", "model", valid_llm_request(), - Arc::new(|_request| { - Box::pin(async { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) - }), + Arc::new(|_request| Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) })), ) .await .expect("host stream should be returned"); @@ -605,9 +603,7 @@ async fn callback_stream_stops_when_host_receiver_is_dropped() { "stream_receiver_drop", "model", valid_llm_request(), - Arc::new(|_request| { - Box::pin(async { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) - }), + Arc::new(|_request| Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) })), ) .await .expect("host stream should be returned"); @@ -962,9 +958,7 @@ async fn dropping_host_stream_sends_explicit_worker_cancellation() { "cancel_stream", "model", valid_llm_request(), - Arc::new(|_request| { - Box::pin(async { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) - }), + Arc::new(|_request| Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) })), ) .await .expect("host stream should be returned"); @@ -1502,9 +1496,9 @@ async fn host_runtime_service_covers_continuation_errors_and_stream_items() { let stream_continuation = state .insert_continuation(Continuation::LlmStream(Arc::new(|_request| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Err(FlowError::Internal( - "stream item failed".into(), - ))])) as LlmJsonStream) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Err( + FlowError::Internal("stream item failed".into()), + )]))) }) }))) .expect("stream continuation should insert"); @@ -1541,7 +1535,7 @@ async fn host_runtime_service_covers_continuation_errors_and_stream_items() { let stream_continuation = state .insert_continuation(Continuation::LlmStream(Arc::new(|_request| { - Box::pin(async move { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) + Box::pin(async move { Ok(LlmJsonStream::new(tokio_stream::empty())) }) }))) .expect("stream continuation should insert"); let invalid_stream_request = match service diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index ca541caa3..5dd2e3914 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -281,10 +281,9 @@ fn managed_and_streaming_calls_cull_event_inputs_and_annotations_with_real_codec .func(Arc::new(|request| { Box::pin(async move { assert_eq!(request.content["messages"].as_array().unwrap().len(), 4); - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) }) })) .collector(Box::new(|_| Ok(()))) @@ -384,10 +383,9 @@ fn projection_encode_failures_do_not_block_managed_or_streaming_calls() { .func(Arc::new(|request| { Box::pin(async move { assert_eq!(request.content["messages"].as_array().unwrap().len(), 4); - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) }) })) .collector(Box::new(|_| Ok(()))) @@ -855,10 +853,9 @@ fn llm_stream_call_execute_adds_otel_status_metadata_to_end_events() { .request(request()) .func(Arc::new(|_request| { Box::pin(async { - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) }) })) .collector(Box::new(|_chunk| Ok(()))) @@ -916,9 +913,9 @@ fn llm_stream_call_execute_adds_otel_error_metadata_to_failed_end_events() { .request(request()) .func(Arc::new(|_request| { Box::pin(async { - Ok(Box::pin(tokio_stream::iter(vec![Err(FlowError::Internal( - "stream boom".to_string(), - ))])) as LlmJsonStream) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Err( + FlowError::Internal("stream boom".to_string()), + )]))) }) })) .collector(Box::new(|_chunk| Ok(()))) @@ -939,10 +936,9 @@ fn llm_stream_call_execute_adds_otel_error_metadata_to_failed_end_events() { .request(request()) .func(Arc::new(|_request| { Box::pin(async { - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) - as LlmJsonStream, - ) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) }) })) .collector(Box::new(|_chunk| { diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 3fc87180c..3e3243b16 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -953,7 +953,7 @@ async fn native_stream_adapter_covers_chunks_end_errors_and_cancellation() { #[tokio::test] async fn relay_stream_adapter_covers_poll_end_error_and_cancel() { - let stream: LlmJsonStream = Box::pin(tokio_stream::iter(vec![ + let stream = LlmJsonStream::new(tokio_stream::iter(vec![ Ok(json!({"chunk": 1})), Err(FlowError::Internal("stream failed".into())), ])); @@ -975,7 +975,7 @@ async fn relay_stream_adapter_covers_poll_end_error_and_cancel() { ); drop_native_stream(raw); - let stream: LlmJsonStream = Box::pin(tokio_stream::empty()); + let stream = LlmJsonStream::new(tokio_stream::empty()); let raw = relay_stream_to_native_stream(stream); let poll = raw.next.unwrap(); out = ptr::null_mut(); @@ -999,7 +999,7 @@ async fn relay_stream_adapter_covers_poll_end_error_and_cancel() { ); unsafe { drop_relay_llm_stream(ptr::null_mut()) }; - let stream: LlmJsonStream = Box::pin(tokio_stream::empty()); + let stream = LlmJsonStream::new(tokio_stream::empty()); let raw = relay_stream_to_native_stream(stream); let state = unsafe { &*(raw.user_data as *const NativeHostLlmStream) }; let mutex = state.stream.clone(); @@ -1024,7 +1024,9 @@ fn native_stream_continuation_covers_success_and_error() { let next: LlmStreamExecutionNextFn = Arc::new(|_request| { Box::pin(async { - Ok(Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) as LlmJsonStream) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) }) }); let next_ctx = Box::into_raw(Box::new(next)) as *mut c_void; diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index b24c0591b..995a21fbf 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -13,8 +13,8 @@ use tokio::sync::Notify; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::llm::{llm_conditional_execution, llm_request_intercepts}; -use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; +use crate::api::runtime::{LlmJsonStream, NemoRelayContextState}; use crate::api::tool::tool_conditional_execution; use crate::error::FlowError; @@ -913,10 +913,9 @@ fn test_plugin_registration_context_covers_all_registration_helpers() { 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) @@ -1635,10 +1634,9 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) @@ -1649,10 +1647,9 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ), @@ -1765,10 +1762,9 @@ fn test_plugin_registration_context_maps_deregistration_errors() { 1, Arc::new(|_name, request, _next| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(request.content)])) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok( + request.content + )]))) }) }), ) diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs index ec38fdb4c..5fa7418e2 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs @@ -464,7 +464,7 @@ async fn stream_monitor_records_blocked_message() { #[tokio::test(flavor = "current_thread")] async fn guarded_provider_stream_reports_block_after_forwarded_chunks() { - let provider_stream: LlmJsonStream = Box::pin(tokio_stream::iter(vec![Ok(json!({ + let provider_stream = LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ "choices": [{"delta": {"content": "blocked"}}], }))])); let (text_tx, mut text_rx) = mpsc::channel::>(8); @@ -483,6 +483,8 @@ async fn guarded_provider_stream_reports_block_after_forwarded_chunks() { } Ok(()) }); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); forward_guarded_provider_stream( provider_stream, @@ -491,6 +493,8 @@ async fn guarded_provider_stream_reports_block_after_forwarded_chunks() { chunk_tx, monitor, blocked, + cancel_rx, + closed, ) .await; diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs index 330e87480..9cdfc5397 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs @@ -435,7 +435,7 @@ async fn remote_initialization_installs_stream_execution_intercept() { called.store(true, Ordering::SeqCst); Box::pin(async move { let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }); @@ -616,7 +616,7 @@ async fn remote_streaming_http_errors_are_reported_and_marked() { called.store(true, Ordering::SeqCst); Box::pin(async move { let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }); @@ -773,7 +773,7 @@ async fn remote_streaming_malformed_chunk_is_reported_and_marked() { let func: LlmStreamExecutionNextFn = Arc::new(move |_req| { Box::pin(async move { let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(Box::pin(stream) as LlmJsonStream) + Ok(LlmJsonStream::new(stream)) }) }); diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 652d6c108..af178bac7 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -905,6 +905,18 @@ NemoRelayStatus nemo_relay_llm_stream_call_execute(const char *name, const struct FfiCodecHandle *response_codec, struct FfiStream **out); +/** + * Stop a stream producer and wait for cleanup to complete. + * + * This operation is idempotent. It does not free `stream`; callers must still + * release the handle with [`nemo_relay_stream_free`]. + * + * # Safety + * `stream` must be a valid `FfiStream` pointer returned by + * `nemo_relay_llm_stream_call_execute`, or null. + */ +NemoRelayStatus nemo_relay_stream_close(struct FfiStream *stream); + /** * Poll the next chunk from a streaming LLM response. Blocks until a chunk is * available. diff --git a/crates/ffi/src/api/llm.rs b/crates/ffi/src/api/llm.rs index ad311248e..2cacc1c7b 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -443,6 +443,45 @@ pub unsafe extern "C" fn nemo_relay_llm_call_execute( pub struct FfiStream { pub(crate) receiver: tokio::sync::Mutex>>, + pub(crate) cancel: tokio::sync::watch::Sender, + pub(crate) closed: tokio::sync::watch::Receiver>>, +} + +impl Drop for FfiStream { + fn drop(&mut self) { + self.cancel.send_replace(true); + } +} + +async fn forward_stream_to_channel( + mut stream: nemo_relay::api::runtime::LlmJsonStream, + tx: tokio::sync::mpsc::Sender>, + mut cancel: tokio::sync::watch::Receiver, + closed: tokio::sync::watch::Sender>>, +) { + loop { + if *cancel.borrow() { + break; + } + let item = tokio::select! { + _ = cancel.changed() => break, + item = stream.next() => item, + }; + let Some(item) = item else { + break; + }; + tokio::select! { + _ = cancel.changed() => break, + result = tx.send(item) => { + if result.is_err() { + break; + } + } + } + } + closed.send_replace(Some( + stream.close().await.map_err(|error| error.to_string()), + )); } /// Execute a streaming LLM call end-to-end. Conditional-execution guardrails @@ -553,16 +592,18 @@ pub unsafe extern "C" fn nemo_relay_llm_stream_call_execute( match result { Ok(rust_stream) => { let (tx, rx) = tokio::sync::mpsc::channel(32); - tokio_runtime().spawn(async move { - let mut stream = rust_stream; - while let Some(item) = stream.next().await { - if tx.send(item).await.is_err() { - break; - } - } - }); + let (cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, closed_rx) = tokio::sync::watch::channel(None); + tokio_runtime().spawn(forward_stream_to_channel( + rust_stream, + tx, + cancel_rx, + closed, + )); let ffi_stream = Box::new(FfiStream { receiver: tokio::sync::Mutex::new(rx), + cancel, + closed: closed_rx, }); unsafe { *out = Box::into_raw(ffi_stream) }; NemoRelayStatus::Ok @@ -571,6 +612,42 @@ pub unsafe extern "C" fn nemo_relay_llm_stream_call_execute( } } +/// Stop a stream producer and wait for cleanup to complete. +/// +/// This operation is idempotent. It does not free `stream`; callers must still +/// release the handle with [`nemo_relay_stream_free`]. +/// +/// # Safety +/// `stream` must be a valid `FfiStream` pointer returned by +/// `nemo_relay_llm_stream_call_execute`, or null. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_stream_close(stream: *mut FfiStream) -> NemoRelayStatus { + clear_last_error(); + if stream.is_null() { + set_last_error("null pointer argument"); + return NemoRelayStatus::NullPointer; + } + let stream = unsafe { &*stream }; + let result = tokio_runtime().block_on(async { + stream.cancel.send_replace(true); + let mut closed = stream.closed.clone(); + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + nemo_relay::error::FlowError::Internal("stream close task ended early".into()) + })?; + } + let result = closed.borrow().clone().expect("close state checked above"); + let mut receiver = stream.receiver.lock().await; + receiver.close(); + while receiver.try_recv().is_ok() {} + result.map_err(nemo_relay::error::FlowError::Internal) + }); + match result { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_error(&error), + } +} + /// Poll the next chunk from a streaming LLM response. Blocks until a chunk is /// available. /// diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 955e8fed7..9b88276e6 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -23,12 +23,12 @@ use std::sync::Arc; use libc::c_char; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; -use tokio_stream::{Stream, StreamExt}; +use tokio_stream::StreamExt; use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; @@ -493,12 +493,8 @@ pub fn wrap_llm_stream_exec_intercept_fn( &str, LlmRequest, LlmStreamExecutionNextFn, - ) -> Pin< - Box< - dyn Future> + Send>>>> - + Send, - >, - > + Send + ) -> Pin> + Send>> + + Send + Sync, > { let ud = make_user_data(user_data, free_fn); @@ -558,7 +554,7 @@ pub fn wrap_llm_stream_exec_intercept_fn( )?; unsafe { nemo_relay_string_free_internal(result_ptr) }; let stream = tokio_stream::once(Ok(result)); - Ok(Box::pin(stream) as Pin> + Send>>) + Ok(LlmJsonStream::new(stream)) }) }, ) @@ -752,15 +748,7 @@ pub fn wrap_llm_stream_exec_fn( user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> Box< - dyn Fn( - LlmRequest, - ) -> Pin< - Box< - dyn Future> + Send>>>> - + Send, - >, - > + Send - + Sync, + dyn Fn(LlmRequest) -> Pin> + Send>> + Send + Sync, > { let ud = make_user_data(user_data, free_fn); Box::new(move |request: LlmRequest| { @@ -775,7 +763,7 @@ pub fn wrap_llm_stream_exec_fn( // The C callback returns the full response as a single JSON value for stream // We emit it as a single-item stream let stream = tokio_stream::once(Ok(result)); - Ok(Box::pin(stream) as Pin> + Send>>) + Ok(LlmJsonStream::new(stream)) }) }) } diff --git a/crates/ffi/tests/integration/callable_extra_tests.rs b/crates/ffi/tests/integration/callable_extra_tests.rs index d82c158f0..b7f68845f 100644 --- a/crates/ffi/tests/integration/callable_extra_tests.rs +++ b/crates/ffi/tests/integration/callable_extra_tests.rs @@ -168,7 +168,9 @@ fn test_callable_extra_trampoline_and_helper_paths() { let empty_next: LlmStreamExecutionNextFn = Arc::new(|request| { Box::pin(async move { assert_eq!(request.content, Json::Null); - Ok(Box::pin(tokio_stream::empty()) as Pin> + Send>>) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::empty(), + )) }) }); let mut empty_stream = runtime diff --git a/crates/ffi/tests/integration/main.rs b/crates/ffi/tests/integration/main.rs index 1aacee267..baed8d570 100644 --- a/crates/ffi/tests/integration/main.rs +++ b/crates/ffi/tests/integration/main.rs @@ -10,7 +10,7 @@ use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, Too use nemo_relay::api::scope::{ScopeAttributes, ScopeHandle, ScopeType}; use nemo_relay::api::tool::{ToolAttributes, ToolHandle}; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; -use nemo_relay::error::{FlowError, Result}; +use nemo_relay::error::FlowError; use nemo_relay_ffi::api::*; use nemo_relay_ffi::callable::*; use nemo_relay_ffi::convert::*; @@ -19,9 +19,7 @@ use nemo_relay_ffi::types::*; use nemo_relay_ffi::{api, convert, error}; use serde_json::{Value as Json, json}; use std::ffi::{CStr, CString}; -use std::pin::Pin; use std::sync::{Arc, Mutex}; -use tokio_stream::Stream; static TEST_MUTEX: Mutex<()> = Mutex::new(()); diff --git a/crates/ffi/tests/unit/api/coverage_sweeps_tests.rs b/crates/ffi/tests/unit/api/coverage_sweeps_tests.rs index 989ffbe70..3c0381eb6 100644 --- a/crates/ffi/tests/unit/api/coverage_sweeps_tests.rs +++ b/crates/ffi/tests/unit/api/coverage_sweeps_tests.rs @@ -1866,6 +1866,8 @@ fn test_ffi_stream_next_reports_error_items() { let stream = Box::into_raw(Box::new(FfiStream { receiver: tokio::sync::Mutex::new(rx), + cancel: tokio::sync::watch::channel(false).0, + closed: tokio::sync::watch::channel(Some(Ok(()))).1, })); unsafe { diff --git a/crates/ffi/tests/unit/api/execution_tests.rs b/crates/ffi/tests/unit/api/execution_tests.rs index 066b871c2..0d158544d 100644 --- a/crates/ffi/tests/unit/api/execution_tests.rs +++ b/crates/ffi/tests/unit/api/execution_tests.rs @@ -250,6 +250,8 @@ fn test_ffi_llm_stream_execute_response_codec_defaults_and_error_paths() { assert_eq!(nemo_relay_stream_next(stream, &mut chunk), 1); let stream_chunk = returned_json(chunk); assert_eq!(stream_chunk["id"], json!("chatcmpl-ffi")); + assert_eq!(nemo_relay_stream_close(stream), NemoRelayStatus::Ok); + assert_eq!(nemo_relay_stream_close(stream), NemoRelayStatus::Ok); assert_eq!(nemo_relay_stream_next(stream, &mut chunk), 0); assert!(lock_unpoisoned(collected_chunks()).is_empty()); assert_eq!(*lock_unpoisoned(finalizer_calls()), 0); diff --git a/crates/ffi/tests/unit/callable_tests.rs b/crates/ffi/tests/unit/callable_tests.rs index 4791b2a46..22199aeba 100644 --- a/crates/ffi/tests/unit/callable_tests.rs +++ b/crates/ffi/tests/unit/callable_tests.rs @@ -471,10 +471,9 @@ fn test_wrap_llm_exec_stream_and_event_callbacks() { wrap_llm_stream_exec_intercept_fn(llm_exec_short_circuit_cb, std::ptr::null_mut(), None); let next_stream: LlmStreamExecutionNextFn = Arc::new(|_request| { Box::pin(async { - Ok( - Box::pin(tokio_stream::iter(vec![Ok(json!({"ignored": true}))])) - as Pin> + Send>>, - ) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(vec![Ok(json!({"ignored": true}))]), + )) }) }); let mut intercepted_stream = runtime @@ -487,10 +486,11 @@ fn test_wrap_llm_exec_stream_and_event_callbacks() { wrap_llm_stream_exec_intercept_fn(llm_exec_intercept_cb, std::ptr::null_mut(), None); let next_stream: LlmStreamExecutionNextFn = Arc::new(|request| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(json!({ - "model": request.content["model"].clone() - }))])) - as Pin> + Send>>) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(vec![Ok(json!({ + "model": request.content["model"].clone() + }))]), + )) }) }); let mut intercepted_stream = runtime diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index e4cf8d982..1a7cc8a5e 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -16,7 +16,8 @@ use std::pin::Pin; use std::ptr; use std::sync::Arc; use std::sync::Mutex as StdMutex; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::task::{Context, Poll}; use chrono::{DateTime, Utc}; use napi::bindgen_prelude::*; @@ -25,13 +26,14 @@ use napi::{JsFunction, JsObject, JsUnknown, NapiRaw, NapiValue}; use napi_derive::napi; use serde::Deserialize; use serde_json::Value as Json; -use tokio_stream::StreamExt; +use tokio_stream::{Stream, StreamExt}; use nemo_relay::api::llm as core_llm_api; use nemo_relay::api::llm::{LlmAttributes, LlmRequest}; use nemo_relay::api::registry as core_registry_api; use nemo_relay::api::runtime::{ - EventSanitizeFn, LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn, + EventSanitizeFn, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, LlmStreamInner, + ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, @@ -333,17 +335,43 @@ fn build_openinference_config( static NEXT_STREAM_ID: AtomicU64 = AtomicU64::new(0); type StreamSender = tokio::sync::mpsc::UnboundedSender>; -type RustJsonStream = std::pin::Pin> + Send>>; +type RustJsonStream = LlmJsonStream; -static STREAM_CHANNELS: std::sync::LazyLock>> = +struct StreamChannel { + sender: StreamSender, + cancelled: AtomicBool, + closed: tokio::sync::watch::Sender>>, +} + +static STREAM_CHANNELS: std::sync::LazyLock>>> = std::sync::LazyLock::new(|| StdMutex::new(HashMap::new())); -fn register_stream_channel(id: u64, tx: StreamSender) { - STREAM_CHANNELS.lock().unwrap().insert(id, tx); +fn register_stream_channel( + id: u64, + tx: StreamSender, +) -> tokio::sync::watch::Receiver>> { + let (closed, closed_rx) = tokio::sync::watch::channel(None); + STREAM_CHANNELS.lock().unwrap().insert( + id, + Arc::new(StreamChannel { + sender: tx, + cancelled: AtomicBool::new(false), + closed, + }), + ); + closed_rx +} + +fn finish_stream_channel(id: u64, result: std::result::Result<(), String>) { + if let Some(channel) = STREAM_CHANNELS.lock().unwrap().remove(&id) { + channel.closed.send_replace(Some(result)); + } } -fn remove_stream_channel(id: u64) { - STREAM_CHANNELS.lock().unwrap().remove(&id); +fn cancel_stream_channel(id: u64) { + if let Some(channel) = STREAM_CHANNELS.lock().unwrap().get(&id) { + channel.cancelled.store(true, Ordering::Release); + } } fn ensure_stream_callback_queued(id: u64, status: napi::Status) -> FlowResult<()> { @@ -351,7 +379,12 @@ fn ensure_stream_callback_queued(id: u64, status: napi::Status) -> FlowResult<() return Ok(()); } - remove_stream_channel(id); + finish_stream_channel( + id, + Err(format!( + "failed to queue JS stream producer callback: {status:?}" + )), + ); Err(FlowError::Internal(format!( "failed to queue JS stream producer callback: {status:?}", ))) @@ -360,12 +393,72 @@ fn ensure_stream_callback_queued(id: u64, status: napi::Status) -> FlowResult<() async fn forward_stream_to_channel( mut stream: RustJsonStream, tx: tokio::sync::mpsc::Sender>, + mut cancel: tokio::sync::watch::Receiver, + closed: tokio::sync::watch::Sender>>, ) { - while let Some(item) = stream.next().await { - if tx.send(item).await.is_err() { + loop { + if *cancel.borrow() { + break; + } + let item = tokio::select! { + _ = cancel.changed() => break, + item = stream.next() => item, + }; + let Some(item) = item else { break; + }; + tokio::select! { + _ = cancel.changed() => break, + result = tx.send(item) => { + if result.is_err() { + break; + } + } } } + closed.send_replace(Some( + stream.close().await.map_err(|error| error.to_string()), + )); +} + +struct NodePushStream { + receiver: tokio_stream::wrappers::UnboundedReceiverStream>, + stream_id: u64, + closed: tokio::sync::watch::Receiver>>, +} + +impl Stream for NodePushStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.receiver).poll_next(cx) + } +} + +impl Drop for NodePushStream { + fn drop(&mut self) { + cancel_stream_channel(self.stream_id); + } +} + +impl LlmStreamInner for NodePushStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let stream_id = self.stream_id; + let mut closed = self.get_mut().closed.clone(); + Box::pin(async move { + cancel_stream_channel(stream_id); + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + FlowError::Internal("JS stream cleanup task ended early".into()) + })?; + } + closed + .borrow() + .clone() + .expect("close state checked above") + .map_err(FlowError::Internal) + }) + } } /// Push a chunk into the stream identified by `streamId`. @@ -373,8 +466,8 @@ async fn forward_stream_to_channel( #[napi] pub fn push_stream_chunk(stream_id: f64, chunk: Json) -> bool { let id = stream_id as u64; - if let Some(tx) = STREAM_CHANNELS.lock().unwrap().get(&id) { - tx.send(Ok(chunk)).is_ok() + if let Some(channel) = STREAM_CHANNELS.lock().unwrap().get(&id) { + !channel.cancelled.load(Ordering::Acquire) && channel.sender.send(Ok(chunk)).is_ok() } else { false } @@ -385,7 +478,7 @@ pub fn push_stream_chunk(stream_id: f64, chunk: Json) -> bool { #[napi] pub fn end_stream(stream_id: f64) { let id = stream_id as u64; - remove_stream_channel(id); + finish_stream_channel(id, Ok(())); } #[allow(clippy::enum_variant_names)] @@ -2098,8 +2191,10 @@ pub fn llm_call_execute_async( /// /// The optional `collector` callback is invoked with each intercepted chunk as JSON, /// allowing the caller to accumulate chunks for aggregation. The optional `finalizer` -/// callback is invoked once when the stream is exhausted and must return a JSON value -/// representing the aggregated response. +/// callback is invoked once when the stream is exhausted or closed early and +/// must return a JSON value representing the aggregated response. Consumers +/// that stop reading early must await `stream.close()` to wait for producer +/// cleanup and surface cleanup errors. #[allow(clippy::too_many_arguments)] #[napi(ts_return_type = "Promise")] pub fn llm_stream_call_execute( @@ -2143,7 +2238,7 @@ pub fn llm_stream_call_execute( let default_fn: LlmStreamExecutionNextFn = std::sync::Arc::new(move |req: LlmRequest| { let stream_id = NEXT_STREAM_ID.fetch_add(1, Ordering::Relaxed); let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - register_stream_channel(stream_id, tx); + let closed = register_stream_channel(stream_id, tx); // Serialize the LlmRequest to JSON and wrap with streamId so JS can extract both let req_json = serde_json::to_value(&req).unwrap_or(Json::Null); @@ -2159,11 +2254,11 @@ pub fn llm_stream_call_execute( Box::pin(async move { ensure_stream_callback_queued(stream_id, call_status)?; - let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx); - Ok(Box::pin(stream) - as std::pin::Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::from_closeable(NodePushStream { + receiver: tokio_stream::wrappers::UnboundedReceiverStream::new(rx), + stream_id, + closed, + })) }) }); @@ -2197,10 +2292,19 @@ pub fn llm_stream_call_execute( .map_err(to_napi_err)?; let (tx, rx) = tokio::sync::mpsc::channel(32); - tokio::spawn(forward_stream_to_channel(rust_stream, tx)); + let (cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, closed_rx) = tokio::sync::watch::channel(None); + tokio::spawn(forward_stream_to_channel( + rust_stream, + tx, + cancel_rx, + closed, + )); Ok(LlmStream { receiver: tokio::sync::Mutex::new(rx), + cancel, + closed: closed_rx, }) }) .await diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 8c5853bab..c888bb1a5 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; @@ -781,15 +781,8 @@ pub fn wrap_js_llm_stream_exec_intercept_fn( &str, LlmRequest, LlmStreamExecutionNextFn, - ) -> Pin< - Box< - dyn Future< - Output = Result< - Pin> + Send>>, - >, - > + Send, - >, - > + Send + ) -> Pin> + Send>> + + Send + Sync, > { Arc::new( @@ -818,10 +811,7 @@ pub fn wrap_js_llm_stream_exec_intercept_fn( value => vec![Ok(value)], }; let stream = tokio_stream::iter(chunks); - Ok(Box::pin(stream) - as Pin< - Box> + Send>, - >) + Ok(LlmJsonStream::new(stream)) }) }, ) diff --git a/crates/node/src/stream.rs b/crates/node/src/stream.rs index 3cac321cc..c9308d66a 100644 --- a/crates/node/src/stream.rs +++ b/crates/node/src/stream.rs @@ -19,6 +19,8 @@ use serde_json::Value as Json; #[napi] pub struct LlmStream { pub(crate) receiver: tokio::sync::Mutex>>, + pub(crate) cancel: tokio::sync::watch::Sender, + pub(crate) closed: tokio::sync::watch::Receiver>>, } #[napi] @@ -37,4 +39,27 @@ impl LlmStream { Some(Err(e)) => Err(napi::Error::from_reason(e.to_string())), } } + + /// Stop the producer and wait for its cleanup to complete. + #[napi] + pub async fn close(&self) -> Result<()> { + self.cancel.send_replace(true); + let mut closed = self.closed.clone(); + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + napi::Error::from_reason("stream close task ended before releasing the producer") + })?; + } + let result = closed.borrow().clone().expect("close state checked above"); + let mut receiver = self.receiver.lock().await; + receiver.close(); + while receiver.try_recv().is_ok() {} + result.map_err(napi::Error::from_reason) + } +} + +impl Drop for LlmStream { + fn drop(&mut self) { + self.cancel.send_replace(true); + } } diff --git a/crates/node/tests/typed_tests.mjs b/crates/node/tests/typed_tests.mjs index 2e2f9b650..843675a7b 100644 --- a/crates/node/tests/typed_tests.mjs +++ b/crates/node/tests/typed_tests.mjs @@ -609,6 +609,49 @@ describe('typedLlmStreamExecute', () => { assert.equal(collected.length, 2); }); + it('close waits for async-generator cleanup and exhausts subsequent reads', async () => { + const passthrough = new JsonPassthrough(); + let releaseProducer; + let producerReady; + const ready = new Promise((resolve) => { + producerReady = resolve; + }); + let finallyRan = false; + async function* partiallyConsumedSource() { + try { + yield { token: 'first' }; + await new Promise((resolve) => { + releaseProducer = resolve; + producerReady(); + }); + yield { token: 'unread' }; + } finally { + finallyRan = true; + } + } + + const stream = await typedLlmStreamExecute( + 'stream_close_waits_for_finally', + makeNative(), + partiallyConsumedSource, + () => {}, + () => null, + passthrough, + passthrough, + ); + assert.deepEqual(await stream.next(), { token: 'first' }); + await ready; + + const closing = stream.close(); + await Promise.resolve(); + assert.equal(finallyRan, false); + releaseProducer(); + await closing; + + assert.equal(finallyRan, true); + assert.equal(await stream.next(), null); + }); + it('stream with envelopeCodec for chunks and response', async () => { const native = makeNative(); diff --git a/crates/node/typed.js b/crates/node/typed.js index 5b9cceab5..ac8110e61 100644 --- a/crates/node/typed.js +++ b/crates/node/typed.js @@ -195,7 +195,7 @@ async function typedLlmExecute(name, request, func, responseJsonCodec, options) * @param {*} request - The LLM request object ({headers, content}). * @param {function(*): AsyncIterable} func - The streaming LLM implementation. * @param {function(TChunk): void} collector - Called with each typed chunk after intercepts. - * @param {function(): TResponse} finalizer - Called once when the stream is exhausted. + * @param {function(): TResponse} finalizer - Called once when the stream is exhausted or closed early. * @param {Codec} chunkJsonCodec - Codec for serializing/deserializing chunks. * @param {Codec} responseJsonCodec - Codec for serializing/deserializing the final response. * @param {object} [options] - Optional parameters. @@ -211,11 +211,17 @@ async function typedLlmExecute(name, request, func, responseJsonCodec, options) * @returns {Promise} A promise resolving to the native stream handle. * @remarks The JavaScript side drives async iteration and pushes each encoded * chunk back into the native stream bridge; the stream is always closed in the - * `finally` path even if the source iterator throws. + * `finally` path even if the source iterator throws. Consumers that stop + * reading early must await `stream.close()` to complete producer cleanup. */ async function typedLlmStreamExecute(name, request, func, collector, finalizer, ...streamArgs) { const [chunkJsonCodec, responseJsonCodec, options] = streamArgs; const opts = options || {}; + let iterator; + let resolveIterator; + const iteratorReady = new Promise((resolve) => { + resolveIterator = resolve; + }); // Push-based stream bridge: NAPI cannot resolve JS Promises from // call_with_return_value, so the JS side drives async generator iteration @@ -226,11 +232,25 @@ async function typedLlmStreamExecute(name, request, func, collector, finalizer, const streamId = wrapper.__nemo_relay_stream_id; (async () => { try { - for await (const typedChunk of func(req)) { - lib.pushStreamChunk(streamId, chunkJsonCodec.toJson(typedChunk)); + iterator = func(req)[Symbol.asyncIterator](); + resolveIterator(iterator); + while (true) { + const { done, value: typedChunk } = await iterator.next(); + if (done) { + break; + } + if (!lib.pushStreamChunk(streamId, chunkJsonCodec.toJson(typedChunk))) { + await iterator.return?.(); + break; + } } } finally { - lib.endStream(streamId); + resolveIterator(iterator); + try { + await iterator?.return?.(); + } finally { + lib.endStream(streamId); + } } })(); }; @@ -245,7 +265,7 @@ async function typedLlmStreamExecute(name, request, func, collector, finalizer, return collectCodecReturn(responseJsonCodec.toJson(finalizer())); }; - return await lib.llmStreamCallExecute( + const stream = await lib.llmStreamCallExecute( name, request, jsonFunc, @@ -260,6 +280,35 @@ async function typedLlmStreamExecute(name, request, func, collector, finalizer, opts.codec ? (payload) => encodeWithCodec(opts.codec, payload) : null, opts.responseCodec ? (response) => decodeResponseWithCodec(opts.responseCodec, response) : null, ); + const close = stream.close.bind(stream); + stream.close = async () => { + const closing = close(); + let iteratorError; + try { + await (await iteratorReady)?.return?.(); + } catch (error) { + iteratorError = error; + } + let closeError; + try { + await closing; + } catch (error) { + closeError = error; + } + if (iteratorError && closeError) { + throw new AggregateError( + [iteratorError, closeError], + 'stream close failed during iterator cleanup and native close', + ); + } + if (iteratorError) { + throw iteratorError; + } + if (closeError) { + throw closeError; + } + }; + return stream; } module.exports = { diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index d87370638..5abbdd933 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -13,7 +13,9 @@ use std::sync::Arc; use nemo_relay::api::llm as core_llm_api; use nemo_relay::api::llm::LlmAttributes; use nemo_relay::api::registry as core_registry_api; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn}; +use nemo_relay::api::runtime::{ + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, +}; use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, @@ -40,8 +42,7 @@ use crate::py_types::{ PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, }; -pub(crate) type RustJsonStream = - std::pin::Pin> + Send>>; +pub(crate) type RustJsonStream = LlmJsonStream; /// Convert an [`FlowError`] into a Python `RuntimeError`. fn to_py_err(e: FlowError) -> PyErr { @@ -100,12 +101,32 @@ fn py_annotated_llm_response( pub(crate) async fn forward_stream_to_channel( mut stream: RustJsonStream, tx: tokio::sync::mpsc::Sender>, + mut cancel: tokio::sync::watch::Receiver, + closed: tokio::sync::watch::Sender>>, ) { - while let Some(item) = stream.next().await { - if tx.send(item).await.is_err() { + loop { + if *cancel.borrow() { break; } + let item = tokio::select! { + _ = cancel.changed() => break, + item = stream.next() => item, + }; + let Some(item) = item else { + break; + }; + tokio::select! { + _ = cancel.changed() => break, + result = tx.send(item) => { + if result.is_err() { + break; + } + } + } } + closed.send_replace(Some( + stream.close().await.map_err(|error| error.to_string()), + )); } // --------------------------------------------------------------------------- @@ -775,8 +796,9 @@ fn llm_call_execute<'py>( /// func: An async callable ``(LlmRequest) -> AsyncIterator[Any]`` that returns JSON chunks. /// collector: A callable ``(Any) -> None`` invoked with each intercepted chunk /// (after stream execution intercepts have been applied). -/// finalizer: A callable ``() -> Any`` invoked once when the stream is exhausted. -/// Its return value is the aggregated response (converted to JSON). +/// finalizer: A callable ``() -> Any`` invoked once when the stream is exhausted +/// or explicitly closed. Its return value is the aggregated response +/// (converted to JSON). /// handle: Optional parent scope handle. /// attributes: Optional ``LlmAttributes`` bitflags. /// data: Optional JSON-serializable application data. @@ -786,6 +808,9 @@ fn llm_call_execute<'py>( /// response_codec: Optional response codec used to attach annotated response data /// to emitted end events. /// +/// Consumers that stop reading early must call ``await stream.aclose()`` to +/// wait for producer cleanup and surface cleanup errors. +/// /// Returns: /// An awaitable that resolves to an ``LlmStream`` async iterator of JSON chunks. #[pyfunction] @@ -861,10 +886,19 @@ fn llm_stream_call_execute<'py>( // Spawn a tokio task that drains the Rust stream into an mpsc channel let (tx, rx) = tokio::sync::mpsc::channel::>(32); - tokio::spawn(forward_stream_to_channel(rust_stream, tx)); + let (cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, closed_rx) = tokio::sync::watch::channel(None); + tokio::spawn(forward_stream_to_channel( + rust_stream, + tx, + cancel_rx, + closed, + )); Ok(PyLlmStream { - receiver: tokio::sync::Mutex::new(rx), + receiver: Arc::new(tokio::sync::Mutex::new(rx)), + cancel, + closed: closed_rx, }) }) .await diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index 25b1707a2..00bec91ea 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -23,16 +23,19 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + LlmStreamInner, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use nemo_relay::error::{FlowError, Result as FlowResult}; use pyo3::prelude::*; +use pyo3::types::PyDict; use serde_json::Value as Json; use tokio_stream::Stream; +use tokio_stream::wrappers::ReceiverStream; use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; @@ -120,21 +123,61 @@ fn next_async_iter_coro(async_iter: &Arc>) -> FlowResult) -> FlowResult> { +fn schedule_async_iter_task(coro: Py) -> FlowResult> { + Python::attach(|py| { + pyo3_async_runtimes::tokio::get_current_locals(py) + .and_then(|locals| { + let kwargs = PyDict::new(py); + kwargs.set_item("loop", locals.event_loop(py))?; + py.import("asyncio")? + .call_method("ensure_future", (coro,), Some(&kwargs)) + }) + .map(|task| task.unbind()) + .map_err(|e| FlowError::Internal(e.to_string())) + }) +} + +fn cancel_async_iter_task(task: &Py) -> FlowResult<()> { + Python::attach(|py| { + pyo3_async_runtimes::tokio::get_current_locals(py) + .and_then(|locals| { + let cancel = task.bind(py).getattr("cancel")?; + locals + .event_loop(py) + .call_method1("call_soon_threadsafe", (cancel,)) + }) + .map(|_| ()) + .map_err(|error| FlowError::Internal(error.to_string())) + }) +} + +enum AsyncIterTaskResult { + Item(Json), + End, + Cancelled, +} + +async fn await_async_iter_task_result(task: Py) -> FlowResult { let future = Python::attach(|py| { - pyo3_async_runtimes::tokio::into_future(coro.into_bound(py)) + pyo3_async_runtimes::tokio::into_future(task.into_bound(py)) .map_err(|e| FlowError::Internal(e.to_string())) })?; match future.await { Ok(result) => Python::attach(|py| { py_to_json(result.bind(py)) - .map(Some) + .map(AsyncIterTaskResult::Item) .map_err(|e| FlowError::Internal(e.to_string())) }), Err(error) => Python::attach(|py| { + let cancelled_error = py + .import("asyncio") + .and_then(|asyncio| asyncio.getattr("CancelledError")) + .map_err(|error| FlowError::Internal(error.to_string()))?; if error.is_instance_of::(py) { - Ok(None) + Ok(AsyncIterTaskResult::End) + } else if error.is_instance(py, &cancelled_error) { + Ok(AsyncIterTaskResult::Cancelled) } else { Err(FlowError::Internal(error.to_string())) } @@ -142,35 +185,160 @@ async fn await_async_iter_value(coro: Py) -> FlowResult> { } } +#[cfg(test)] +async fn await_async_iter_task(task: Py) -> FlowResult> { + match await_async_iter_task_result(task).await? { + AsyncIterTaskResult::Item(value) => Ok(Some(value)), + AsyncIterTaskResult::End => Ok(None), + AsyncIterTaskResult::Cancelled => Err(FlowError::Internal( + "async iterator task was cancelled".into(), + )), + } +} + +#[cfg(test)] +async fn await_async_iter_value(coro: Py) -> FlowResult> { + await_async_iter_task(schedule_async_iter_task(coro)?).await +} + +async fn close_async_iter(async_iter: &Arc>) -> FlowResult<()> { + let close = Python::attach(|py| { + let iter = async_iter.bind(py); + match iter.call_method0("aclose") { + Ok(close) => Ok(Some(close.unbind())), + Err(error) if error.is_instance_of::(py) => { + Ok(None) + } + Err(error) => Err(FlowError::Internal(error.to_string())), + } + }); + let Some(close) = close? else { + return Ok(()); + }; + let task = schedule_async_iter_task(close)?; + let future = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(task.into_bound(py)) + .map_err(|error| FlowError::Internal(error.to_string())) + })?; + future + .await + .map_err(|error| FlowError::Internal(error.to_string()))?; + Ok(()) +} + async fn forward_async_iter( async_iter: Arc>, tx: tokio::sync::mpsc::Sender>, + mut cancel: tokio::sync::watch::Receiver, + closed: tokio::sync::watch::Sender>>, ) { - loop { + let result = loop { + if *cancel.borrow() { + break close_async_iter(&async_iter).await; + } let next_value = match next_async_iter_coro(&async_iter) { - Ok(None) => break, - Ok(Some(coro)) => await_async_iter_value(coro).await, + Ok(None) => break Ok(()), + Ok(Some(coro)) => match schedule_async_iter_task(coro) { + Ok(task) => { + let task_for_future = Python::attach(|py| task.clone_ref(py)); + let mut next_value = Box::pin(await_async_iter_task_result(task_for_future)); + tokio::select! { + _ = cancel.changed() => { + let _ = cancel_async_iter_task(&task); + let next_result = next_value.await; + let close_result = close_async_iter(&async_iter).await; + break match next_result { + Err(error) => Err(error), + Ok(_) => close_result, + }; + } + value = &mut next_value => value, + } + } + Err(error) => Err(error), + }, Err(error) => Err(error), }; match next_value { - Ok(Some(value)) => { - if tx.send(Ok(value)).await.is_err() { - break; + Ok(AsyncIterTaskResult::Item(value)) => { + let sent = tokio::select! { + _ = cancel.changed() => { + break close_async_iter(&async_iter).await; + } + sent = tx.send(Ok(value)) => sent, + }; + if sent.is_err() { + break close_async_iter(&async_iter).await; } } - Ok(None) => break, + Ok(AsyncIterTaskResult::End) => break Ok(()), + Ok(AsyncIterTaskResult::Cancelled) => { + let close_result = close_async_iter(&async_iter).await; + break close_result.and(Err(FlowError::Internal( + "async iterator task was cancelled".into(), + ))); + } Err(error) => { - let _ = tx.send(Err(error)).await; - break; + let sent = tokio::select! { + _ = cancel.changed() => break close_async_iter(&async_iter).await, + sent = tx.send(Err(error)) => sent, + }; + if sent.is_err() { + break close_async_iter(&async_iter).await; + } + break close_async_iter(&async_iter).await; } } + }; + closed.send_replace(Some(result)); +} + +#[derive(Clone)] +struct PythonAsyncIteratorClose { + cancel: tokio::sync::watch::Sender, + closed: tokio::sync::watch::Receiver>>, +} + +struct PythonAsyncIteratorStream { + receiver: ReceiverStream>, + close: PythonAsyncIteratorClose, +} + +impl Stream for PythonAsyncIteratorStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.receiver).poll_next(cx) } } -fn stream_from_async_iter( - async_iter: Py, -) -> FlowResult> + Send>>> { +impl Drop for PythonAsyncIteratorStream { + fn drop(&mut self) { + self.close.cancel.send_replace(true); + } +} + +impl LlmStreamInner for PythonAsyncIteratorStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let this = self.get_mut(); + this.close.cancel.send_replace(true); + this.receiver.close(); + while this.receiver.as_mut().try_recv().is_ok() {} + let close = this.close.clone(); + Box::pin(async move { + let mut closed = close.closed; + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + FlowError::Internal("Python stream cleanup task ended early".into()) + })?; + } + closed.borrow().clone().expect("close state checked above") + }) + } +} + +fn stream_from_async_iter(async_iter: Py) -> FlowResult { let (tx, rx) = tokio::sync::mpsc::channel::>(32); let task_locals = Python::attach(|py| { pyo3_async_runtimes::tokio::get_current_locals(py) @@ -178,12 +346,20 @@ fn stream_from_async_iter( })?; let async_iter = Arc::new(async_iter); + let (cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, closed_rx) = tokio::sync::watch::channel(None); tokio::spawn(pyo3_async_runtimes::tokio::scope(task_locals, async move { - forward_async_iter(async_iter, tx).await; + forward_async_iter(async_iter, tx, cancel_rx, closed).await; })); - let stream = tokio_stream::wrappers::ReceiverStream::new(rx); - Ok(Box::pin(stream) as Pin> + Send>>) + let stream = PythonAsyncIteratorStream { + receiver: ReceiverStream::new(rx), + close: PythonAsyncIteratorClose { + cancel, + closed: closed_rx, + }, + }; + Ok(LlmJsonStream::from_closeable(stream)) } /// Wrap a Python callable `(str, Json) -> Json` for tool sanitize/intercept fns. @@ -378,18 +554,19 @@ impl PyLlmStreamNextFn { // Drain into mpsc channel and return PyLlmStream let (tx, rx) = tokio::sync::mpsc::channel::>(32); - tokio::spawn(async move { - use tokio_stream::StreamExt; - let mut stream = rust_stream; - while let Some(item) = stream.next().await { - if tx.send(item).await.is_err() { - break; - } - } - }); + let (cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, closed_rx) = tokio::sync::watch::channel(None); + tokio::spawn(crate::py_api::forward_stream_to_channel( + rust_stream, + tx, + cancel_rx, + closed, + )); Ok(crate::py_types::PyLlmStream { - receiver: tokio::sync::Mutex::new(rx), + receiver: Arc::new(tokio::sync::Mutex::new(rx)), + cancel, + closed: closed_rx, }) }) } @@ -555,13 +732,8 @@ pub fn wrap_py_llm_stream_exec_intercept_fn( &str, LlmRequest, LlmStreamExecutionNextFn, - ) -> Pin< - Box< - dyn Future< - Output = FlowResult> + Send>>>, - > + Send, - >, - > + Send + ) -> Pin> + Send>> + + Send + Sync, > { let py_fn = Arc::new(py_fn); @@ -729,15 +901,8 @@ pub fn wrap_py_llm_exec_fn( pub fn wrap_py_llm_stream_exec_fn( py_fn: Py, ) -> Box< - dyn Fn( - LlmRequest, - ) -> Pin< - Box< - dyn Future< - Output = FlowResult> + Send>>>, - > + Send, - >, - > + Send + dyn Fn(LlmRequest) -> Pin> + Send>> + + Send + Sync, > { let py_fn = std::sync::Arc::new(py_fn); @@ -780,9 +945,9 @@ pub fn wrap_py_collector_fn( /// Wrap a Python callable `() -> Any` as a finalizer for streaming LLM calls. /// -/// The finalizer is called once when the stream is fully consumed. Its return -/// value is converted from a Python object to `serde_json::Value` (Json) and -/// used as the aggregated response. +/// The finalizer is called once when the stream is fully consumed or explicitly +/// closed. Its return value is converted from a Python object to +/// `serde_json::Value` (Json) and used as the aggregated response. pub fn wrap_py_finalizer_fn(py_fn: Py) -> Box Json + Send> { Box::new(move || { Python::attach(|py| { diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index 442e7f3b4..aeca45305 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::sync::Arc; + use pyo3::prelude::*; use super::{ @@ -20,10 +22,15 @@ use nemo_relay::api::tool::ToolExecutionInterceptOutcome; /// /// Use ``async for chunk in stream:`` to consume chunks. Each chunk is a /// Python object (converted from JSON). The stream automatically emits an -/// End lifecycle event when exhausted. +/// End lifecycle event when exhausted. When stopping early, await +/// ``stream.aclose()`` to wait for producer cleanup and emit an interrupted +/// End event. #[pyclass(name = "LlmStream")] pub struct PyLlmStream { - pub receiver: tokio::sync::Mutex>>, + pub receiver: + Arc>>>, + pub cancel: tokio::sync::watch::Sender, + pub closed: tokio::sync::watch::Receiver>>, } #[pymethods] @@ -33,18 +40,10 @@ impl PyLlmStream { } pub(crate) fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { - // We need to get a reference to the receiver inside the tokio Mutex. - // Since PyLlmStream is behind a PyRef (shared), we use tokio::sync::Mutex. - let receiver_ptr = &self.receiver - as *const tokio::sync::Mutex< - tokio::sync::mpsc::Receiver>, - >; - // SAFETY: The PyLlmStream outlives this future because Python holds a reference to it. - // The tokio Mutex ensures exclusive access to the receiver. - let receiver_ref = unsafe { &*receiver_ptr }; + let receiver = Arc::clone(&self.receiver); pyo3_async_runtimes::tokio::future_into_py(py, async move { - let mut guard = receiver_ref.lock().await; + let mut guard = receiver.lock().await; let next_item = guard.recv().await; match next_item { None => Err(PyErr::new::( @@ -57,6 +56,37 @@ impl PyLlmStream { } }) } + + /// Stop forwarding chunks, wait for producer cleanup, and release the native stream. + /// + /// This operation is idempotent and raises when producer cleanup fails. + fn aclose<'py>(&self, py: Python<'py>) -> PyResult> { + let cancel = self.cancel.clone(); + let mut closed = self.closed.clone(); + let receiver = Arc::clone(&self.receiver); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + cancel.send_replace(true); + while closed.borrow().is_none() { + closed.changed().await.map_err(|_| { + PyErr::new::( + "stream close task ended before releasing the native stream", + ) + })?; + } + let result = closed.borrow().clone().expect("close state checked above"); + let mut receiver = receiver.lock().await; + receiver.close(); + while receiver.try_recv().is_ok() {} + result.map_err(PyErr::new::) + }) + } +} + +impl Drop for PyLlmStream { + fn drop(&mut self) { + self.cancel.send_replace(true); + } } // --------------------------------------------------------------------------- diff --git a/crates/python/tests/coverage/coverage_tests.rs b/crates/python/tests/coverage/coverage_tests.rs index 16093c2d5..765c3850a 100644 --- a/crates/python/tests/coverage/coverage_tests.rs +++ b/crates/python/tests/coverage/coverage_tests.rs @@ -4,13 +4,11 @@ //! Coverage tests for coverage in the NeMo Relay Python crate. use std::ffi::CString; -use std::pin::Pin; use std::sync::Arc; use pyo3::prelude::*; use pyo3::types::PyModule; use serde_json::{Value as Json, json}; -use tokio_stream::Stream; use tokio_stream::StreamExt; use uuid::Uuid; @@ -214,14 +212,16 @@ payload = {"nested": {"value": 7}, "items": [1, 2, 3]} fn test_py_api_forward_stream_to_channel_exits_when_receiver_is_dropped() { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let stream: crate::py_api::RustJsonStream = Box::pin(tokio_stream::iter(vec![ + let stream = crate::py_api::RustJsonStream::new(tokio_stream::iter(vec![ Ok(json!({"chunk": 1})), Ok(json!({"chunk": 2})), ])); let (tx, rx) = tokio::sync::mpsc::channel(1); drop(rx); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); - crate::py_api::forward_stream_to_channel(stream, tx).await; + crate::py_api::forward_stream_to_channel(stream, tx, cancel_rx, closed).await; }); } @@ -831,10 +831,9 @@ async def llm_stream_intercept(request, next): let stream_next: LlmStreamExecutionNextFn = Arc::new(|_request| { Box::pin(async move { let chunks = vec![Ok(json!({"chunk": "a"})), Ok(json!({"chunk": "b"}))]; - Ok(Box::pin(tokio_stream::iter(chunks)) - as Pin< - Box> + Send>, - >) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(chunks), + )) }) }); let mut stream = stream_intercept("llm", make_request(), stream_next) @@ -971,10 +970,9 @@ async def llm_stream_intercept_fail(request, next): let stream_next: LlmStreamExecutionNextFn = Arc::new(|_request| { Box::pin(async move { let chunks = vec![Ok(json!({"chunk": "downstream"}))]; - Ok(Box::pin(tokio_stream::iter(chunks)) - as Pin< - Box> + Send>, - >) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(chunks), + )) }) }); let mut stream = stream_intercept("llm", make_request(), stream_next) @@ -990,10 +988,9 @@ async def llm_stream_intercept_fail(request, next): wrap_py_llm_stream_exec_intercept_fn(llm_stream_intercept_fail_py); let stream_next: LlmStreamExecutionNextFn = Arc::new(|_request| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": 1}))])) - as Pin< - Box> + Send>, - >) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(vec![Ok(json!({"chunk": 1}))]), + )) }) }); let err = match failing_stream_intercept("llm", make_request(), stream_next).await { diff --git a/crates/python/tests/coverage/py_api_coverage_tests.rs b/crates/python/tests/coverage/py_api_coverage_tests.rs index fa1bb80a4..0dee9e556 100644 --- a/crates/python/tests/coverage/py_api_coverage_tests.rs +++ b/crates/python/tests/coverage/py_api_coverage_tests.rs @@ -814,13 +814,15 @@ fn to_py_err_and_forward_stream_to_channel_cover_private_helpers() { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let stream: RustJsonStream = Box::pin(tokio_stream::iter(vec![ + let stream = RustJsonStream::new(tokio_stream::iter(vec![ Ok(json!({"chunk": 1})), Ok(json!({"chunk": 2})), ])); let (tx, mut rx) = tokio::sync::mpsc::channel(2); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); - forward_stream_to_channel(stream, tx).await; + forward_stream_to_channel(stream, tx, cancel_rx, closed).await; assert_eq!(rx.recv().await.unwrap().unwrap(), json!({"chunk": 1})); assert_eq!(rx.recv().await.unwrap().unwrap(), json!({"chunk": 2})); diff --git a/crates/python/tests/coverage/py_callable_coverage_tests.rs b/crates/python/tests/coverage/py_callable_coverage_tests.rs index 2ea3fb12a..709fed86b 100644 --- a/crates/python/tests/coverage/py_callable_coverage_tests.rs +++ b/crates/python/tests/coverage/py_callable_coverage_tests.rs @@ -6,7 +6,6 @@ use super::*; use std::ffi::CString; -use std::pin::Pin; use std::sync::Arc; use pyo3::types::PyModule; @@ -261,6 +260,8 @@ fn async_iter_helpers_cover_stop_error_and_dropped_receiver_paths() { let module = load_module( py, r#" +import asyncio + class StopIter: def __anext__(self): raise StopAsyncIteration @@ -282,6 +283,19 @@ class ValueIter: return self.value return inner() +class DroppedReceiverIter: + def __init__(self, value): + self.value = value + self.closed = False + + def __anext__(self): + async def inner(): + return self.value + return inner() + + async def aclose(self): + self.closed = True + async def coro_value(): return {"value": 1} @@ -291,6 +305,9 @@ async def coro_stop(): async def coro_error(): raise RuntimeError("await boom") +async def coro_cancel(): + raise asyncio.CancelledError("user cancellation") + async def coro_non_json(): return object() "#, @@ -299,9 +316,12 @@ async def coro_non_json(): let stop_iter_cls: Py = module.getattr("StopIter").unwrap().unbind(); let error_iter_cls: Py = module.getattr("ErrorIter").unwrap().unbind(); let value_iter_cls: Py = module.getattr("ValueIter").unwrap().unbind(); + let dropped_receiver_iter_cls: Py = + module.getattr("DroppedReceiverIter").unwrap().unbind(); let coro_value_fn: Py = module.getattr("coro_value").unwrap().unbind(); let coro_stop_fn: Py = module.getattr("coro_stop").unwrap().unbind(); let coro_error_fn: Py = module.getattr("coro_error").unwrap().unbind(); + let coro_cancel_fn: Py = module.getattr("coro_cancel").unwrap().unbind(); let coro_non_json_fn: Py = module.getattr("coro_non_json").unwrap().unbind(); assert!( @@ -344,6 +364,13 @@ async def coro_non_json(): .to_string() .contains("await boom") ); + assert!( + await_async_iter_value(Python::attach(|py| coro_cancel_fn.call0(py).unwrap())) + .await + .unwrap_err() + .to_string() + .contains("cancelled") + ); assert!( await_async_iter_value(Python::attach(|py| coro_non_json_fn .call0(py) @@ -355,20 +382,28 @@ async def coro_non_json(): ); let (tx, mut rx) = tokio::sync::mpsc::channel(2); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); forward_async_iter( Arc::new(Python::attach(|py| { value_iter_cls.call1(py, (value_payload.bind(py),)).unwrap() })), tx, + cancel_rx, + closed, ) .await; assert_eq!(rx.recv().await.unwrap().unwrap(), json!({"x": 1})); assert!(rx.recv().await.is_none()); let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); forward_async_iter( Arc::new(Python::attach(|py| error_iter_cls.call0(py).unwrap())), tx, + cancel_rx, + closed, ) .await; assert!( @@ -382,15 +417,28 @@ async def coro_non_json(): let (tx, rx) = tokio::sync::mpsc::channel(1); drop(rx); + let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); + let (closed, _closed_rx) = tokio::sync::watch::channel(None); + let dropped_iter = Python::attach(|py| { + dropped_receiver_iter_cls + .call1(py, (dropped_payload.bind(py),)) + .unwrap() + }); forward_async_iter( - Arc::new(Python::attach(|py| { - value_iter_cls - .call1(py, (dropped_payload.bind(py),)) - .unwrap() - })), + Arc::new(Python::attach(|py| dropped_iter.clone_ref(py))), tx, + cancel_rx, + closed, ) .await; + assert!(Python::attach(|py| { + dropped_iter + .bind(py) + .getattr("closed") + .unwrap() + .is_truthy() + .unwrap() + })); Ok(()) }) .unwrap(); @@ -530,10 +578,9 @@ async def collect_stream(awaitable): let stream_next = PyLlmStreamNextFn { inner: Arc::new(|_| { Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": 1}))])) - as Pin< - Box> + Send>, - >) + Ok(nemo_relay::api::runtime::LlmJsonStream::new( + tokio_stream::iter(vec![Ok(json!({"chunk": 1}))]), + )) }) }), }; diff --git a/crates/python/tests/coverage/py_types_coverage_tests.rs b/crates/python/tests/coverage/py_types_coverage_tests.rs index e8783fcb9..0a8ae2452 100644 --- a/crates/python/tests/coverage/py_types_coverage_tests.rs +++ b/crates/python/tests/coverage/py_types_coverage_tests.rs @@ -786,10 +786,14 @@ async def next_item(stream): let (tx_ok, rx_ok) = tokio::sync::mpsc::channel(2); tx_ok.blocking_send(Ok(json!({"chunk": 1}))).unwrap(); drop(tx_ok); + let (cancel_ok, _cancel_ok_rx) = tokio::sync::watch::channel(false); + let (_closed_ok, closed_ok_rx) = tokio::sync::watch::channel(Some(Ok(()))); let stream_ok = pyo3::Py::new( py, PyLlmStream { - receiver: tokio::sync::Mutex::new(rx_ok), + receiver: Arc::new(tokio::sync::Mutex::new(rx_ok)), + cancel: cancel_ok, + closed: closed_ok_rx, }, ) .unwrap(); @@ -819,10 +823,14 @@ async def next_item(stream): ))) .unwrap(); drop(tx_err); + let (cancel_err, _cancel_err_rx) = tokio::sync::watch::channel(false); + let (_closed_err, closed_err_rx) = tokio::sync::watch::channel(Some(Ok(()))); let stream_err = pyo3::Py::new( py, PyLlmStream { - receiver: tokio::sync::Mutex::new(rx_err), + receiver: Arc::new(tokio::sync::Mutex::new(rx_err)), + cancel: cancel_err, + closed: closed_err_rx, }, ) .unwrap(); @@ -840,10 +848,14 @@ async def next_item(stream): let (tx_done, rx_done) = tokio::sync::mpsc::channel(1); drop(tx_done); + let (cancel_done, _cancel_done_rx) = tokio::sync::watch::channel(false); + let (_closed_done, closed_done_rx) = tokio::sync::watch::channel(Some(Ok(()))); let stream_done = pyo3::Py::new( py, PyLlmStream { - receiver: tokio::sync::Mutex::new(rx_done), + receiver: Arc::new(tokio::sync::Mutex::new(rx_done)), + cancel: cancel_done, + closed: closed_done_rx, }, ) .unwrap(); diff --git a/crates/switchyard/src/component.rs b/crates/switchyard/src/component.rs index 28c6ad21d..f4c859733 100644 --- a/crates/switchyard/src/component.rs +++ b/crates/switchyard/src/component.rs @@ -3,18 +3,20 @@ //! Switchyard plugin configuration and Relay execution integration. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use std::time::{Duration, Instant}; -use async_stream::stream; -use futures_util::{StreamExt, stream as futures_stream}; +use futures_util::StreamExt; use nemo_relay::api::event::{CategoryProfile, DataSchema, EventCategory}; use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::optimization::record_llm_optimization_contribution; -use nemo_relay::api::runtime::{LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn}; +use nemo_relay::api::runtime::{ + LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, LlmStreamInner, +}; use nemo_relay::api::scope::{EmitMarkEventParams, event}; use nemo_relay::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationKind, LlmOptimizationModel, @@ -797,9 +799,10 @@ impl SwitchyardRuntime { match first { Some(Ok(first)) => { self.record_routing_contribution(&decision, attempt, true); - let committed = - Box::pin(futures_stream::once(async move { Ok(first) }).chain(upstream)) - as LlmJsonStream; + let committed = LlmJsonStream::from_closeable(PrefixedStream { + first: Some(Ok(first)), + upstream, + }); let output = if target_protocol == inbound { committed } else { @@ -1661,23 +1664,133 @@ fn provider_error_summary(error: &FlowError) -> String { } } -fn mark_terminal_stream( - mut upstream: LlmJsonStream, +struct PrefixedStream { + first: Option>, + upstream: LlmJsonStream, +} + +impl futures_util::Stream for PrefixedStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if let Some(first) = self.first.take() { + Poll::Ready(Some(first)) + } else { + Pin::new(&mut self.upstream).poll_next(cx) + } + } +} + +impl LlmStreamInner for PrefixedStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let this = self.get_mut(); + this.first = None; + Box::pin(async move { this.upstream.close().await }) + } +} + +struct TerminalMarkedStream { + upstream: LlmJsonStream, phase: &'static str, rollout_mode: &'static str, metadata: Json, -) -> LlmJsonStream { - Box::pin(stream! { - while let Some(item) = upstream.next().await { - match item { - Ok(chunk) => yield Ok(chunk), - Err(error) => { - emit_terminal_error(&error, phase, rollout_mode, metadata.clone()); - yield Err(error); - return; + finished: bool, +} + +impl futures_util::Stream for TerminalMarkedStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.finished { + return Poll::Ready(None); + } + match Pin::new(&mut self.upstream).poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => Poll::Ready(Some(Ok(chunk))), + Poll::Ready(Some(Err(error))) => { + self.finished = true; + emit_terminal_error(&error, self.phase, self.rollout_mode, self.metadata.clone()); + Poll::Ready(Some(Err(error))) + } + Poll::Ready(None) => { + self.finished = true; + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl LlmStreamInner for TerminalMarkedStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + Box::pin(async move { self.get_mut().upstream.close().await }) + } +} + +struct TranslatedStream { + upstream: LlmJsonStream, + transcoder: StreamTranscoder, + buffered: VecDeque>, + upstream_finished: bool, +} + +impl futures_util::Stream for TranslatedStream { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + if let Some(chunk) = self.buffered.pop_front() { + return Poll::Ready(Some(chunk)); + } + if self.upstream_finished { + return Poll::Ready(None); + } + + match Pin::new(&mut self.upstream).poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => match self.transcoder.transcode(&chunk) { + Ok(chunks) => self.buffered.extend(chunks.into_iter().map(Ok)), + Err(error) => { + self.upstream_finished = true; + return Poll::Ready(Some(Err(error))); + } + }, + Poll::Ready(Some(Err(error))) => { + self.upstream_finished = true; + return Poll::Ready(Some(Err(error))); + } + Poll::Ready(None) => { + self.upstream_finished = true; + match self.transcoder.finish() { + Ok(chunks) => self.buffered.extend(chunks.into_iter().map(Ok)), + Err(error) => return Poll::Ready(Some(Err(error))), + } } + Poll::Pending => return Poll::Pending, } } + } +} + +impl LlmStreamInner for TranslatedStream { + fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { + let this = self.get_mut(); + this.buffered.clear(); + this.upstream_finished = true; + Box::pin(async move { this.upstream.close().await }) + } +} + +fn mark_terminal_stream( + upstream: LlmJsonStream, + phase: &'static str, + rollout_mode: &'static str, + metadata: Json, +) -> LlmJsonStream { + LlmJsonStream::from_closeable(TerminalMarkedStream { + upstream, + phase, + rollout_mode, + metadata, + finished: false, }) } @@ -1685,39 +1798,13 @@ fn translated_stream( source: WireProtocol, target: WireProtocol, effective_model: String, - mut upstream: LlmJsonStream, + upstream: LlmJsonStream, ) -> LlmJsonStream { - let mut transcoder = StreamTranscoder::new(source, target, effective_model); - Box::pin(stream! { - while let Some(item) = upstream.next().await { - match item { - Ok(chunk) => { - match transcoder.transcode(&chunk) { - Ok(chunks) => { - for chunk in chunks { - yield Ok(chunk); - } - } - Err(error) => { - yield Err(error); - return; - } - } - } - Err(error) => { - yield Err(error); - return; - } - } - } - match transcoder.finish() { - Ok(chunks) => { - for chunk in chunks { - yield Ok(chunk); - } - } - Err(error) => yield Err(error), - } + LlmJsonStream::from_closeable(TranslatedStream { + upstream, + transcoder: StreamTranscoder::new(source, target, effective_model), + buffered: VecDeque::new(), + upstream_finished: false, }) } diff --git a/crates/switchyard/tests/unit/component_tests.rs b/crates/switchyard/tests/unit/component_tests.rs index def694d54..796030982 100644 --- a/crates/switchyard/tests/unit/component_tests.rs +++ b/crates/switchyard/tests/unit/component_tests.rs @@ -3,8 +3,11 @@ //! Unit tests for the Switchyard Relay plugin component. +use std::future::Future; +use std::pin::Pin; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; use axum::{ Json as AxumJson, Router, @@ -12,11 +15,12 @@ use axum::{ http::StatusCode, routing::{get, post}, }; +use futures_util::{Stream, stream as futures_stream}; use nemo_relay::api::event::{Event, ScopeCategory}; use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, }; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn}; +use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, LlmStreamInner}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::codec::optimization::LlmOptimizationSummaryStatus; use nemo_relay::error::{UpstreamFailure, UpstreamFailureClass}; @@ -24,6 +28,56 @@ use nemo_relay::plugin::rollback_registrations; use super::*; +struct CloseTrackingStream { + close_calls: Arc, + closed: bool, +} + +impl Stream for CloseTrackingStream { + type Item = FlowResult; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(None) + } +} + +impl LlmStreamInner for CloseTrackingStream { + fn close( + mut self: Pin<&mut Self>, + ) -> Pin> + Send + '_>> { + self.closed = true; + self.close_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(()) }) + } +} + +fn close_tracking_stream(close_calls: Arc) -> LlmJsonStream { + LlmJsonStream::from_closeable(CloseTrackingStream { + close_calls, + closed: false, + }) +} + +fn prefixed_adapter(upstream: LlmJsonStream) -> LlmJsonStream { + LlmJsonStream::from_closeable(PrefixedStream { + first: Some(Ok(json!({"first": true}))), + upstream, + }) +} + +fn terminal_adapter(upstream: LlmJsonStream) -> LlmJsonStream { + mark_terminal_stream(upstream, "test", "enforce", json!({"route": "test"})) +} + +fn translated_adapter(upstream: LlmJsonStream) -> LlmJsonStream { + translated_stream( + WireProtocol::OpenaiChat, + WireProtocol::AnthropicMessages, + "selected".into(), + upstream, + ) +} + fn binding(protocol: WireProtocol, model: &str) -> TargetBinding { TargetBinding { model: model.into(), @@ -618,10 +672,9 @@ async fn stream_setup_retries_and_empty_streams_have_one_bounded_fallback() { class: UpstreamFailureClass::Connection, })); } - Ok(Box::pin(futures_stream::iter(vec![Ok(chat_chunk( - "ok", - json!("stop"), - ))])) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(vec![Ok( + chat_chunk("ok", json!("stop")), + )]))) }) }); let output = runtime @@ -648,7 +701,7 @@ async fn stream_setup_retries_and_empty_streams_have_one_bounded_fallback() { assert_eq!(request.content["model"], "fallback"); vec![Ok(chat_chunk("fallback", json!("stop")))] }; - Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(items))) }) }); let output = runtime @@ -687,10 +740,10 @@ async fn fallback_setup_failures_preserve_the_provider_error() { #[tokio::test] async fn translated_stream_preserves_success_and_propagates_both_error_sources() { - let source = Box::pin(futures_stream::iter(vec![ + let source = LlmJsonStream::new(futures_stream::iter(vec![ Ok(chat_chunk("hello", Json::Null)), Ok(chat_chunk("", json!("stop"))), - ])) as LlmJsonStream; + ])); let output = translated_stream( WireProtocol::OpenaiChat, WireProtocol::AnthropicMessages, @@ -706,7 +759,7 @@ async fn translated_stream_preserves_success_and_propagates_both_error_sources() })); let upstream_error = FlowError::Internal("upstream stream failed".into()); - let source = Box::pin(futures_stream::iter(vec![Err(upstream_error)])) as LlmJsonStream; + let source = LlmJsonStream::new(futures_stream::iter(vec![Err(upstream_error)])); let output = translated_stream( WireProtocol::OpenaiChat, WireProtocol::AnthropicMessages, @@ -717,9 +770,9 @@ async fn translated_stream_preserves_success_and_propagates_both_error_sources() .await; assert!(output[0].is_err()); - let malformed = Box::pin(futures_stream::iter(vec![Ok(json!({ + let malformed = LlmJsonStream::new(futures_stream::iter(vec![Ok(json!({ "choices": [{"delta": {"reasoning_content": "private"}}] - }))])) as LlmJsonStream; + }))])); let output = translated_stream( WireProtocol::OpenaiChat, WireProtocol::AnthropicMessages, @@ -731,6 +784,23 @@ async fn translated_stream_preserves_success_and_propagates_both_error_sources() assert!(output[0].is_err()); } +#[tokio::test] +async fn stream_adapters_forward_explicit_close_to_the_upstream_stream() { + for make_adapter in [ + prefixed_adapter as fn(LlmJsonStream) -> LlmJsonStream, + terminal_adapter, + translated_adapter, + ] { + let close_calls = Arc::new(AtomicUsize::new(0)); + let mut stream = make_adapter(close_tracking_stream(Arc::clone(&close_calls))); + + stream.close().await.unwrap(); + + assert_eq!(close_calls.load(Ordering::SeqCst), 1); + assert!(stream.next().await.is_none()); + } +} + #[test] fn provider_failure_reporting_covers_every_retry_class() { for (class, label, retryable) in [ @@ -1409,10 +1479,9 @@ async fn streaming_accounting_commits_on_the_first_successful_item_only() { SwitchyardRuntime::new(config(url)).unwrap(), Arc::new(|_| { Box::pin(async { - Ok(Box::pin(futures_stream::iter(vec![Ok(chat_chunk( - "ok", - json!("stop"), - ))])) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(vec![Ok( + chat_chunk("ok", json!("stop")), + )]))) }) }), ) @@ -1443,7 +1512,7 @@ async fn streaming_accounting_commits_on_the_first_successful_item_only() { } else { vec![Ok(chat_chunk("fallback", json!("stop")))] }; - Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(items))) }) }), ) @@ -1467,7 +1536,7 @@ async fn committed_stream_error_keeps_one_route_and_never_redispatches() { let seen = Arc::clone(&seen); Box::pin(async move { seen.fetch_add(1, Ordering::SeqCst); - Ok(Box::pin(futures_stream::iter(vec![ + Ok(LlmJsonStream::new(futures_stream::iter(vec![ Ok(chat_chunk("partial", Json::Null)), Err(FlowError::Upstream(UpstreamFailure { status: None, @@ -1475,7 +1544,7 @@ async fn committed_stream_error_keeps_one_route_and_never_redispatches() { headers: BTreeMap::new(), class: UpstreamFailureClass::Connection, })), - ])) as LlmJsonStream) + ]))) }) }), ) @@ -1618,7 +1687,7 @@ async fn streaming_retries_before_first_item() { } else { vec![Ok(chat_chunk("ok", json!("stop")))] }; - Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(items))) }) }); let stream = runtime @@ -1651,7 +1720,7 @@ async fn streaming_never_retries_after_first_item() { class: UpstreamFailureClass::Connection, })), ]; - Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + Ok(LlmJsonStream::new(futures_stream::iter(items))) }) }); let stream = runtime @@ -1723,9 +1792,9 @@ async fn same_protocol_targets_route_nonportable_streaming_requests() { assert_eq!(request.content["model"], "model-b"); assert_eq!(request.content["store"], true); assert_eq!(request.content["reasoning"]["effort"], "high"); - Ok(Box::pin(futures_stream::iter(vec![Ok( + Ok(LlmJsonStream::new(futures_stream::iter(vec![Ok( json!({"type": "response.output_text.delta", "delta": "ok"}), - )])) as LlmJsonStream) + )]))) }) }); let output = runtime diff --git a/docs/about-nemo-relay/release-notes/highlights.mdx b/docs/about-nemo-relay/release-notes/highlights.mdx index a0063f184..822e0d705 100644 --- a/docs/about-nemo-relay/release-notes/highlights.mdx +++ b/docs/about-nemo-relay/release-notes/highlights.mdx @@ -119,6 +119,9 @@ dynamic plugins. names, or complete request histories. - Update Rust exhaustive matches, direct struct literals, and editor metadata consumers for the new public variants and fields. +- Migrate Rust code that directly constructs `LlmJsonStream`, and update Go + call sites to handle the error returned by `LlmStream.Close`. Call an + explicit close method when a consumer stops reading an LLM stream early. - Point any automation that opens public skills by directory name at the new task-oriented catalog. diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index bc07faae0..953985ffb 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -33,6 +33,10 @@ If you're upgrading from 0.5, plan for these changes: - Update Rust code that exhaustively matches public enums or directly constructs mark types, ATOF sink types, or editor metadata types. The provided builders and constructors are the safest migration path. +- Update managed LLM stream consumers that construct `LlmJsonStream` directly + or ignore Go `LlmStream.Close` results. Use the stream constructors and + explicit close methods described in the compatibility notes when exiting a + stream early. - Replace direct references to old public skill directories with the task-oriented `nemo-relay-install`, `nemo-relay-get-started`, `nemo-relay-instrument-*`, and `nemo-relay-plugin-*` entry points. @@ -69,6 +73,10 @@ Here's what changed across the main product surfaces: - **Decision routing:** The opt-in experimental Switchyard plugin validates and routes buffered or streaming requests across OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages. +- **Managed LLM streams:** Explicit early close now deterministically stops the + producer, finalizes the partial response, and emits the interrupted end + event across Rust, Python, Node.js, Go, and C. Rust and Go callers have + source-compatibility changes; see the migration notes before upgrading. - **Agent skills:** The public catalog now follows task-oriented workflows for installation, getting started, instrumentation, plugin configuration, migration, and debugging. The guidance calls for installers pinned to diff --git a/docs/about-nemo-relay/release-notes/known-issues.mdx b/docs/about-nemo-relay/release-notes/known-issues.mdx index 95bcb9298..d17fbb98b 100644 --- a/docs/about-nemo-relay/release-notes/known-issues.mdx +++ b/docs/about-nemo-relay/release-notes/known-issues.mdx @@ -81,6 +81,28 @@ The following limitations and migration steps apply to NVIDIA NeMo Relay 0.6. ### Compatibility and Migration Notes +#### Update Managed LLM Stream Consumers + + +Managed LLM streams now have one explicit close contract across bindings. When +you stop consuming a stream before natural exhaustion, close it so Relay can +stop the producer, finalize the partial response, and emit the interrupted end +event. Natural exhaustion is unchanged. + +Rust `LlmJsonStream` is no longer the former boxed-stream type alias. Replace +direct construction and alias-dependent type usage with `LlmJsonStream::new` or +`LlmJsonStream::from_closeable`; call `stream.close().await` on early exit. Go +`LlmStream.Close` now returns an `error`, so update call sites to handle its +result. Python keeps `await stream.aclose()`, Node.js provides +`await stream.close()`, and C callers can call `nemo_relay_stream_close` before +`nemo_relay_stream_free`. + +The [Rust Library Reference](/reference/api/rust-library-reference) documents +the current Rust API. Go and raw C FFI remain experimental and source-first. +PR [#465](https://github.com/NVIDIA/NeMo-Relay/pull/465) introduced the unified +close behavior. + + #### Reinstall Coding-Agent Integrations diff --git a/docs/instrument-applications/code-examples.mdx b/docs/instrument-applications/code-examples.mdx index 0c2fcabab..9b706998b 100644 --- a/docs/instrument-applications/code-examples.mdx +++ b/docs/instrument-applications/code-examples.mdx @@ -146,7 +146,7 @@ let response = llm_call_execute( ## Streaming LLM Execution -Use the streaming helper when subscribers need chunk collection plus one final response payload. +Use the streaming helper when subscribers need chunk collection plus one final response payload. Natural exhaustion finalizes the collected response automatically. If your application stops reading early, explicitly close the stream in cleanup; close is idempotent, waits for producer cleanup, finalizes the partial response, and emits an interrupted end event. @@ -179,6 +179,12 @@ stream = await llm_stream_execute( chunk_json_codec=DataclassCodec(Chunk), response_json_codec=DataclassCodec(FinalResponse), ) + +try: + first_chunk = await anext(stream) + # Stop early after deciding no more output is needed. +finally: + await stream.aclose() ``` @@ -213,6 +219,13 @@ const stream = await typedLlmStreamExecute( chunkCodec, responseCodec, ); + +try { + const firstChunk = await stream.next(); + // Stop early after deciding no more output is needed. +} finally { + await stream.close(); +} ``` @@ -221,19 +234,21 @@ const stream = await typedLlmStreamExecute( use nemo_relay::api::llm::{ llm_stream_call_execute, LlmAttributes, LlmRequest, LlmStreamCallExecuteParams, }; +use nemo_relay::api::runtime::LlmJsonStream; use serde_json::json; +use tokio_stream::StreamExt; let request = LlmRequest { headers: Default::default(), content: json!({"messages": [{"role": "user", "content": "hello"}]}), }; -let stream = llm_stream_call_execute( +let mut stream = llm_stream_call_execute( LlmStreamCallExecuteParams::builder() .name("demo-provider") .request(request) .func(std::sync::Arc::new(|_req| Box::pin(async move { - Ok(Box::pin(tokio_stream::iter(vec![Ok(json!({"delta": "hi"}))]))) + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({"delta": "hi"}))]))) }))) .collector(Box::new(|_chunk| Ok(()))) .finalizer(Box::new(|| json!({"text": "hi"}))) @@ -241,6 +256,10 @@ let stream = llm_stream_call_execute( .model_name("demo-model") .build(), ).await?; + +let first_chunk = stream.next().await; +// Stop early after deciding no more output is needed. +stream.close().await?; ``` diff --git a/docs/resources/glossary.mdx b/docs/resources/glossary.mdx index c91106710..3c5a17d44 100644 --- a/docs/resources/glossary.mdx +++ b/docs/resources/glossary.mdx @@ -115,7 +115,7 @@ surface are experimental unless a page says otherwise. **Finalizer** -A **finalizer** is the callback used by streaming LLM helpers when the stream ends. It turns collected stream state into the response payload that sanitize-response guardrails, subscribers, and exporters can observe. +A **finalizer** is the callback used by streaming LLM helpers when the stream ends through natural exhaustion or explicit close. It turns collected stream state into the response payload that sanitize-response guardrails, subscribers, and exporters can observe. **Global And Scope-Local Registration** @@ -151,7 +151,7 @@ An **LLM call** is an instrumented model-provider invocation. Managed LLM calls **LLM Stream** -An **LLM stream** is a streaming model response managed across multiple chunks rather than a single response object. NeMo Relay captures the originating scope stack, runs stream execution intercepts, collects chunks, and finalizes the stream into a response-side event payload. +An **LLM stream** is a streaming model response managed across multiple chunks rather than a single response object. NeMo Relay captures the originating scope stack, runs stream execution intercepts, collects chunks, and finalizes the stream into a response-side event payload. Consumers that stop early explicitly close the stream so the producer can clean up and the partial response can be finalized. **Managed Execution And Manual Lifecycle** diff --git a/go/nemo_relay/atof_test.go b/go/nemo_relay/atof_test.go index a873bfc86..5f166aa4a 100644 --- a/go/nemo_relay/atof_test.go +++ b/go/nemo_relay/atof_test.go @@ -43,6 +43,33 @@ func TestNewAtofExporterConfigDefaults(t *testing.T) { } } +func TestAtofSinkConfigConstructorsSerializeTheirDiscriminators(t *testing.T) { + file := NewAtofFileSinkConfig() + if file.Mode != AtofExporterModeAppend { + t.Fatalf("file sink mode = %q, want append", file.Mode) + } + fileJSON, err := json.Marshal(file) + if err != nil { + t.Fatalf("marshal file sink: %v", err) + } + if !strings.Contains(string(fileJSON), `"type":"file"`) { + t.Fatalf("file sink discriminator missing: %s", fileJSON) + } + + stream := NewAtofStreamSinkConfig("http://localhost:8080/events") + if stream.Transport != AtofEndpointTransportHTTPPost || stream.TimeoutMillis != 3000 || + stream.FieldNamePolicy != AtofEndpointFieldNamePolicyPreserve { + t.Fatalf("unexpected stream sink defaults: %#v", stream) + } + streamJSON, err := json.Marshal(stream) + if err != nil { + t.Fatalf("marshal stream sink: %v", err) + } + if !strings.Contains(string(streamJSON), `"type":"stream"`) { + t.Fatalf("stream sink discriminator missing: %s", streamJSON) + } +} + func TestAtofExporterLifecycleWritesRawJSONL(t *testing.T) { dir := t.TempDir() exporter, err := NewAtofExporter(AtofExporterConfig{Sink: AtofFileSinkConfig{ diff --git a/go/nemo_relay/coverage_gap_test.go b/go/nemo_relay/coverage_gap_test.go index b219d0a38..f57b74a6c 100644 --- a/go/nemo_relay/coverage_gap_test.go +++ b/go/nemo_relay/coverage_gap_test.go @@ -22,15 +22,24 @@ func TestEventBaseNilPointerFallbacks(t *testing.T) { if got := event.Kind(); got != "" { t.Fatalf("expected empty Kind, got %q", got) } + if got := event.ATOFVersion(); got != "" { + t.Fatalf("expected empty ATOFVersion, got %q", got) + } if got := event.ScopeType(); got != "" { t.Fatalf("expected empty ScopeType, got %q", got) } if got := event.Attributes(); got != 0 { t.Fatalf("expected zero Attributes, got %d", got) } + if got := event.AttributesJSON(); got != nil { + t.Fatalf("expected nil AttributesJSON, got %s", got) + } if got := event.Data(); got != nil { t.Fatalf("expected nil Data, got %s", got) } + if got := event.DataSchema(); got != nil { + t.Fatalf("expected nil DataSchema, got %s", got) + } if got := event.Metadata(); got != nil { t.Fatalf("expected nil Metadata, got %s", got) } diff --git a/go/nemo_relay/event_sanitizers_test.go b/go/nemo_relay/event_sanitizers_test.go index e2bf42abd..a1bc467af 100644 --- a/go/nemo_relay/event_sanitizers_test.go +++ b/go/nemo_relay/event_sanitizers_test.go @@ -117,6 +117,56 @@ func TestScopeLocalEventSanitizerInheritanceAndCleanup(t *testing.T) { runTestWithScopeStack(t, testScopeLocalEventSanitizerInheritanceAndCleanup) } +func TestScopeLocalEventSanitizersCanBeDeregistered(t *testing.T) { + runTestWithScopeStack(t, func(t *testing.T) { + owner, err := PushScope("deregister-owner", ScopeTypeAgent) + if err != nil { + t.Fatalf("PushScope failed: %v", err) + } + defer func() { + if err := PopScope(owner); err != nil { + t.Fatalf("PopScope failed: %v", err) + } + }() + + passThrough := func(_ Event, fields EventSanitizeFields) EventSanitizeFields { return fields } + for _, sanitizer := range []struct { + name string + register func() error + deregister func() error + }{ + { + name: "mark", + register: func() error { + return ScopeRegisterMarkSanitizeGuardrail(owner.UUID(), "deregister-mark", 0, passThrough) + }, + deregister: func() error { return ScopeDeregisterMarkSanitizeGuardrail(owner.UUID(), "deregister-mark") }, + }, + { + name: "scope start", + register: func() error { + return ScopeRegisterScopeSanitizeStartGuardrail(owner.UUID(), "deregister-start", 0, passThrough) + }, + deregister: func() error { return ScopeDeregisterScopeSanitizeStartGuardrail(owner.UUID(), "deregister-start") }, + }, + { + name: "scope end", + register: func() error { + return ScopeRegisterScopeSanitizeEndGuardrail(owner.UUID(), "deregister-end", 0, passThrough) + }, + deregister: func() error { return ScopeDeregisterScopeSanitizeEndGuardrail(owner.UUID(), "deregister-end") }, + }, + } { + if err := sanitizer.register(); err != nil { + t.Fatalf("register %s sanitizer: %v", sanitizer.name, err) + } + if err := sanitizer.deregister(); err != nil { + t.Fatalf("deregister %s sanitizer: %v", sanitizer.name, err) + } + } + }) +} + func testScopeLocalEventSanitizerInheritanceAndCleanup(t *testing.T) { var mu sync.Mutex seen := map[string]json.RawMessage{} diff --git a/go/nemo_relay/llm_test.go b/go/nemo_relay/llm_test.go index 6bf406512..9bb402b33 100644 --- a/go/nemo_relay/llm_test.go +++ b/go/nemo_relay/llm_test.go @@ -974,9 +974,15 @@ func TestLlmStreamCloseIsIdempotent(t *testing.T) { t.Fatalf(llmStreamCallExecuteFailed, err) } - stream.Close() - stream.Close() - stream.Close() + if err := stream.Close(); err != nil { + t.Fatalf("first Close failed: %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close failed: %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("third Close failed: %v", err) + } _, err = stream.Next() if err != io.EOF { @@ -984,6 +990,239 @@ func TestLlmStreamCloseIsIdempotent(t *testing.T) { } } +func TestLlmStreamConcurrentCloseIsSafe(t *testing.T) { + stream, err := LlmStreamCallExecute("concurrent_close_llm", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`"data: [DONE]\n\n"`), nil + }, + nil, nil, + ) + if err != nil { + t.Fatalf(llmStreamCallExecuteFailed, err) + } + + var wait sync.WaitGroup + errs := make(chan error, 8) + for i := 0; i < 8; i++ { + wait.Add(1) + go func() { + defer wait.Done() + errs <- stream.Close() + }() + } + wait.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("concurrent Close failed: %v", err) + } + } + if _, err := stream.Next(); err != io.EOF { + t.Fatalf("Next after concurrent Close error = %v, want io.EOF", err) + } +} + +func TestLlmStreamCollectorCanClose(t *testing.T) { + var stream *LlmStream + var closeErr error + stream, closeErr = LlmStreamCallExecute("collector_close_llm", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`"data: {\"chunk\": 1}\n\ndata: [DONE]\n\n"`), nil + }, + func(json.RawMessage) { + closeErr = stream.Close() + }, + nil, + ) + if closeErr != nil { + t.Fatalf(llmStreamCallExecuteFailed, closeErr) + } + if _, err := stream.Next(); err != nil { + t.Fatalf("Next failed: %v", err) + } + if closeErr != nil { + t.Fatalf("Close from collector failed: %v", closeErr) + } + if _, err := stream.Next(); err != io.EOF { + t.Fatalf("Next after collector Close error = %v, want io.EOF", err) + } +} + +func TestLlmStreamHelperAndReleaseCoverage(t *testing.T) { + chunk := json.RawMessage(`{"chunk": true}`) + collectorCalls := 0 + returnedChunk, err := llmStreamNextResult(1, chunk, func(json.RawMessage) { + collectorCalls++ + }, nil) + if err != nil || string(returnedChunk) != string(chunk) || collectorCalls != 1 { + t.Fatalf("chunk result = %s, %v; collector calls = %d", returnedChunk, err, collectorCalls) + } + + finalizerCalls := 0 + finalizer := FinalizerFunc(func() string { + finalizerCalls++ + return `{}` + }) + if _, err := llmStreamNextResult(0, nil, nil, &finalizer); err != io.EOF || finalizerCalls != 1 || finalizer != nil { + t.Fatalf("EOF result = %v; finalizer calls = %d", err, finalizerCalls) + } + + stream, err := LlmStreamCallExecute("release_llm", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`"data: [DONE]\n\n"`), nil + }, + nil, nil, + ) + if err != nil { + t.Fatalf(llmStreamCallExecuteFailed, err) + } + stream.release() + stream.release() + if _, err := stream.Next(); err != io.EOF { + t.Fatalf("Next after release error = %v, want io.EOF", err) + } +} + +func TestLlmStreamFinishCloseWaitsForInFlightWork(t *testing.T) { + stream, err := LlmStreamCallExecute("finish_close_wait_llm", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`"data: [DONE]\n\n"`), nil + }, + nil, nil, + ) + if err != nil { + t.Fatalf(llmStreamCallExecuteFailed, err) + } + + stream.mu.Lock() + ptr := stream.ptr + stream.inFlight = 1 + stream.idle = make(chan struct{}) + stream.closing = true + stream.closeDone = make(chan struct{}) + stream.mu.Unlock() + + finished := make(chan struct{}) + go func() { + stream.finishClose(ptr, nil) + close(finished) + }() + stream.idle <- struct{}{} + + stream.mu.Lock() + stream.inFlight = 0 + close(stream.idle) + stream.mu.Unlock() + <-finished + + if _, err := stream.Next(); err != io.EOF { + t.Fatalf("Next after finishClose error = %v, want io.EOF", err) + } +} + +func TestLlmStreamCloseWaitsForActiveCollectorBeforeFinalizing(t *testing.T) { + collectorStarted := make(chan struct{}) + collectorFinished := make(chan struct{}) + allowCollector := make(chan struct{}) + finalizerStarted := make(chan struct{}) + finalizerFinished := make(chan struct{}) + allowFinalizer := make(chan struct{}) + stream, err := LlmStreamCallExecute("collector_close_order_llm", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`"data: {\"chunk\": 1}\n\ndata: [DONE]\n\n"`), nil + }, + func(json.RawMessage) { + close(collectorStarted) + <-allowCollector + close(collectorFinished) + }, + func() string { + select { + case <-collectorFinished: + close(finalizerStarted) + default: + panic("finalizer ran before collector completed") + } + <-allowFinalizer + close(finalizerFinished) + return `{"partial": true}` + }, + ) + if err != nil { + t.Fatalf(llmStreamCallExecuteFailed, err) + } + + nextResult := make(chan error, 1) + go func() { + _, err := stream.Next() + nextResult <- err + }() + <-collectorStarted + closeResults := make(chan error, 2) + go func() { closeResults <- stream.Close() }() + go func() { closeResults <- stream.Close() }() + select { + case err := <-closeResults: + t.Fatalf("Close returned before collector completed: %v", err) + default: + } + select { + case <-finalizerStarted: + t.Fatal("finalizer ran before collector completed") + default: + } + + close(allowCollector) + if err := <-nextResult; err != nil { + t.Fatalf("Next failed: %v", err) + } + <-finalizerStarted + select { + case err := <-closeResults: + t.Fatalf("Close returned before finalizer completed: %v", err) + default: + } + close(allowFinalizer) + <-finalizerFinished + for range []struct{}{{}, {}} { + if err := <-closeResults; err != nil { + t.Fatalf("Close failed: %v", err) + } + } +} + +func TestLlmStreamCloseFinalizesPartialResponse(t *testing.T) { + request := makeRequest() + finalizerCalls := 0 + stream, err := LlmStreamCallExecute("close_partial_llm", request, + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + chunks := `data: {"chunk": 1}` + "\n\n" + + `data: {"chunk": 2}` + "\n\n" + + `data: [DONE]` + "\n\n" + return json.RawMessage(`"` + strings.ReplaceAll(chunks, `"`, `\"`) + `"`), nil + }, + nil, func() string { + finalizerCalls++ + return `{"partial": true}` + }, + ) + if err != nil { + t.Fatalf(llmStreamCallExecuteFailed, err) + } + if _, err := stream.Next(); err != nil { + t.Fatalf("first stream chunk failed: %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + if finalizerCalls != 1 { + t.Fatalf("finalizer calls = %d, want 1", finalizerCalls) + } + if _, err := stream.Next(); err != io.EOF { + t.Fatalf("Next after Close error = %v, want io.EOF", err) + } +} + func TestLlmStreamNilCollectorFinalizer(t *testing.T) { request := makeRequest() diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 44a6c14f1..b8b5280dd 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -1120,13 +1120,14 @@ func LlmCallExecute(name string, request interface{}, fn LLMExecutionFunc, opts // an [LlmStream] that yields individual SSE (Server-Sent Event) chunks. // Stream execution intercepts are applied to each chunk as it is consumed. // The caller must call [LlmStream.Next] repeatedly until [io.EOF] is -// returned, then call [LlmStream.Close]. +// returned, or call [LlmStream.Close] to stop early. Close waits for producer +// cleanup, finalizes the partial response, and returns any cleanup error. // // The optional collector callback is invoked with each intercepted chunk string, // allowing the caller to accumulate chunks for aggregation. The optional -// finalizer callback is invoked once when the stream is exhausted and must -// return a JSON string representing the aggregated response. Pass nil for -// either to use the default no-op behavior. +// finalizer callback is invoked once when the stream is exhausted or closed +// early and must return a JSON string representing the aggregated response. +// Pass nil for either to use the default no-op behavior. func LlmStreamCallExecute(name string, request interface{}, fn LLMExecutionFunc, collector CollectorFunc, finalizer FinalizerFunc, opts ...LLMCallOption) (*LlmStream, error) { o := &llmCallOptions{} for _, opt := range opts { diff --git a/go/nemo_relay/observability_plugin_test.go b/go/nemo_relay/observability_plugin_test.go index 9adcf40d9..6a0eea564 100644 --- a/go/nemo_relay/observability_plugin_test.go +++ b/go/nemo_relay/observability_plugin_test.go @@ -81,6 +81,32 @@ func TestObservabilityConfigHelpers(t *testing.T) { assertWrappedObservabilityConfig(t, wrapped) } +func TestObservabilityAtofSinkConfigConstructorsSerializeTheirDiscriminators(t *testing.T) { + file := NewObservabilityAtofFileSinkConfig() + if file.Mode != "append" { + t.Fatalf("file sink mode = %q, want append", file.Mode) + } + fileJSON, err := json.Marshal(file) + if err != nil { + t.Fatalf("marshal file sink: %v", err) + } + if !strings.Contains(string(fileJSON), `"type":"file"`) { + t.Fatalf("file sink discriminator missing: %s", fileJSON) + } + + stream := NewObservabilityAtofStreamSinkConfig("http://localhost:8080/events") + if stream.Transport != "http_post" || stream.TimeoutMillis != 3000 || stream.FieldNamePolicy != "preserve" { + t.Fatalf("unexpected stream sink defaults: %#v", stream) + } + streamJSON, err := json.Marshal(stream) + if err != nil { + t.Fatalf("marshal stream sink: %v", err) + } + if !strings.Contains(string(streamJSON), `"type":"stream"`) { + t.Fatalf("stream sink discriminator missing: %s", streamJSON) + } +} + func assertWrappedObservabilityConfig(t *testing.T, wrapped PluginComponentSpec) { t.Helper() if _, ok := wrapped.Config["atof"].(map[string]any); !ok { diff --git a/go/nemo_relay/stream.go b/go/nemo_relay/stream.go index 46643055c..c7a3462ab 100644 --- a/go/nemo_relay/stream.go +++ b/go/nemo_relay/stream.go @@ -10,6 +10,7 @@ package nemo_relay typedef struct FfiStream FfiStream; extern int32_t nemo_relay_stream_next(FfiStream* stream, char** out_chunk); +extern int32_t nemo_relay_stream_close(FfiStream* stream); extern void nemo_relay_stream_free(FfiStream* stream); extern void nemo_relay_string_free(char* ptr); */ @@ -19,6 +20,9 @@ import ( "encoding/json" "io" "runtime" + "strconv" + "strings" + "sync" ) // LlmStream wraps a streaming LLM response returned by [LlmStreamCallExecute]. @@ -44,16 +48,54 @@ import ( // fmt.Print(chunk) // } // -// The stream is not safe for concurrent use. If not closed explicitly, the +// Calls to Next and Close are synchronized. If not closed explicitly, the // underlying C resources are freed automatically by a Go runtime finalizer. // // Each stream carries its own collector and finalizer callbacks, so multiple // streams can operate concurrently without interfering with one another. type LlmStream struct { - ptr *C.FfiStream - closed bool - collector CollectorFunc - finalizer FinalizerFunc + nextMu sync.Mutex + mu sync.Mutex + ptr *C.FfiStream + closed bool + closing bool + inFlight int + idle chan struct{} + closeDone chan struct{} + closeErr error + callbackGoroutine uint64 + collector CollectorFunc + finalizer FinalizerFunc +} + +// currentGoroutineID identifies a callback-reentrant Close so that only that +// callback avoids waiting for its own completion; external callers wait for +// collector and finalizer cleanup. +func currentGoroutineID() uint64 { + var stack [64]byte + n := runtime.Stack(stack[:], false) + fields := strings.Fields(string(stack[:n])) + if len(fields) < 2 { + return 0 + } + id, _ := strconv.ParseUint(fields[1], 10, 64) + return id +} + +func (s *LlmStream) beginCallback() uint64 { + id := currentGoroutineID() + s.mu.Lock() + s.callbackGoroutine = id + s.mu.Unlock() + return id +} + +func (s *LlmStream) finishCallback(id uint64) { + s.mu.Lock() + if s.callbackGoroutine == id { + s.callbackGoroutine = 0 + } + s.mu.Unlock() } func llmStreamNextResult(rc int32, chunk json.RawMessage, collector CollectorFunc, finalizer *FinalizerFunc) (json.RawMessage, error) { @@ -83,12 +125,25 @@ func newLlmStream(ptr *C.FfiStream, collector CollectorFunc, finalizer Finalizer collector: collector, finalizer: finalizer, } - runtime.SetFinalizer(s, func(s *LlmStream) { - s.Close() - }) + runtime.SetFinalizer(s, (*LlmStream).release) return s } +// release frees the native handle without waiting for producer cleanup or +// invoking user callbacks. Explicit Close owns deterministic cleanup. +func (s *LlmStream) release() { + s.mu.Lock() + defer s.mu.Unlock() + if s.ptr == nil { + return + } + C.nemo_relay_stream_free(s.ptr) + s.ptr = nil + s.closed = true + s.collector = nil + s.finalizer = nil +} + // Next returns the next chunk from the stream as a JSON value. It returns // [io.EOF] when the stream is exhausted and all chunks have been consumed. // Any registered stream execution intercepts are applied to each chunk before @@ -100,35 +155,139 @@ func newLlmStream(ptr *C.FfiStream, collector CollectorFunc, finalizer Finalizer // // If the stream has already been closed, Next returns io.EOF. func (s *LlmStream) Next() (json.RawMessage, error) { - if s.closed || s.ptr == nil { + s.nextMu.Lock() + defer s.nextMu.Unlock() + + s.mu.Lock() + if s.closed || s.closing || s.ptr == nil { + s.mu.Unlock() return nil, io.EOF } + ptr := s.ptr + s.inFlight++ + if s.inFlight == 1 { + s.idle = make(chan struct{}) + } + s.mu.Unlock() var chunk *C.char - rc := C.nemo_relay_stream_next(s.ptr, &chunk) + rc := C.nemo_relay_stream_next(ptr, &chunk) + defer s.finishNext() if rc == 1 { // Chunk available text := C.GoString(chunk) C.nemo_relay_string_free(chunk) - return llmStreamNextResult(int32(rc), json.RawMessage(text), s.collector, &s.finalizer) + chunk := json.RawMessage(text) + s.mu.Lock() + collector := s.collector + s.mu.Unlock() + if collector != nil { + callbackID := s.beginCallback() + defer s.finishCallback(callbackID) + collector(chunk) + } + return chunk, nil } - return llmStreamNextResult(int32(rc), nil, s.collector, &s.finalizer) + if rc == 0 { + s.mu.Lock() + finalizer := s.finalizer + s.finalizer = nil + s.mu.Unlock() + if finalizer != nil { + finalizer() + } + return nil, io.EOF + } + return nil, lastError() } -// Close releases the underlying C stream resources. It is safe to call Close -// multiple times; subsequent calls are no-ops. After Close is called, any -// further calls to [LlmStream.Next] return [io.EOF]. -// -// If the stream has not been fully consumed, the finalizer (if provided) will -// NOT be called. Callers should consume the stream to completion or explicitly -// handle finalization if needed before closing early. -func (s *LlmStream) Close() { - if !s.closed && s.ptr != nil { - C.nemo_relay_stream_free(s.ptr) - s.ptr = nil - s.closed = true - s.collector = nil - s.finalizer = nil +func (s *LlmStream) finishNext() { + s.mu.Lock() + defer s.mu.Unlock() + s.inFlight-- + if s.inFlight == 0 { + close(s.idle) + } +} + +// Close stops the producer, waits for cleanup, and releases the underlying C +// stream resources. It is safe to call Close multiple times; subsequent calls +// are no-ops. After Close is called, any further calls to [LlmStream.Next] +// return [io.EOF]. The finalizer runs once with any collected partial response. +func (s *LlmStream) Close() error { + callerID := currentGoroutineID() + s.mu.Lock() + if s.closing { + if s.callbackGoroutine != 0 && s.callbackGoroutine == callerID { + err := s.closeErr + s.mu.Unlock() + return err + } + done := s.closeDone + s.mu.Unlock() + <-done + s.mu.Lock() + err := s.closeErr + s.mu.Unlock() + return err + } + if s.closed || s.ptr == nil { + err := s.closeErr + s.mu.Unlock() + return err + } + s.closing = true + s.closeDone = make(chan struct{}) + ptr := s.ptr + runtime.SetFinalizer(s, nil) + s.mu.Unlock() + + status := C.nemo_relay_stream_close(ptr) + var err error + if status != 0 { + err = lastError() + } + + s.mu.Lock() + if s.inFlight > 0 && s.callbackGoroutine != 0 && s.callbackGoroutine == callerID { + s.mu.Unlock() + // A collector can close its own stream. Finish cleanup after the callback + // returns instead of waiting here and deadlocking that callback. + go s.finishClose(ptr, err) + return err + } + s.mu.Unlock() + s.finishClose(ptr, err) + return err +} + +func (s *LlmStream) finishClose(ptr *C.FfiStream, err error) { + s.mu.Lock() + for s.inFlight > 0 { + idle := s.idle + s.mu.Unlock() + <-idle + s.mu.Lock() } + s.closeErr = err + s.ptr = nil + s.closed = true + s.collector = nil + finalizer := s.finalizer + s.finalizer = nil + s.mu.Unlock() + + C.nemo_relay_stream_free(ptr) + + if finalizer != nil { + callbackID := s.beginCallback() + finalizer() + s.finishCallback(callbackID) + } + + s.mu.Lock() + s.closing = false + close(s.closeDone) + s.mu.Unlock() } diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index d540da1cf..fac3ed23d 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -933,6 +933,9 @@ class LlmStream: async def __anext__(self) -> _Json: """Return the next JSON chunk or raise ``StopAsyncIteration``.""" ... + async def aclose(self) -> None: + """Cancel the producer and release native stream resources.""" + ... class OpenTelemetryConfig: """Mutable configuration for ``OpenTelemetrySubscriber``. diff --git a/python/nemo_relay/llm.py b/python/nemo_relay/llm.py index dd2bd66e0..a2890de01 100644 --- a/python/nemo_relay/llm.py +++ b/python/nemo_relay/llm.py @@ -317,8 +317,11 @@ def stream_execute( Notes: ``collector`` observes the post-intercept chunk values. ``finalizer`` - runs once at stream completion and should return a representation of - the full response, not the final chunk. + runs once at natural stream completion or explicit close and should + return a representation of the full response, not the final chunk. If + the caller stops consuming early, call ``await stream.aclose()`` to + cancel the producer, finalize the partial response, and release native + stream resources. Example:: diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py index 5e9c738ef..f2575c9aa 100644 --- a/python/tests/test_llm.py +++ b/python/tests/test_llm.py @@ -3,6 +3,7 @@ """Tests for NeMo Relay LLM lifecycle, guardrails, intercepts, and streaming.""" +import asyncio from typing import NoReturn, cast import pytest @@ -512,6 +513,32 @@ def finalizer(): # Collector should have received all chunks assert len(collected) == len(chunks) + async def test_stream_execute_aclose_stops_partially_consumed_stream(self): + producer_closed = asyncio.Event() + wait_for_more_chunks = asyncio.Event() + + async def stream_func(request): + try: + yield {"token": "first"} + await wait_for_more_chunks.wait() + finally: + producer_closed.set() + + stream = await llm.stream_execute( + "stream_aclose_llm", + make_request(), + stream_func, + lambda chunk: None, + lambda: {}, + ) + assert await anext(stream) == {"token": "first"} + + await asyncio.wait_for(stream.aclose(), timeout=1) + assert producer_closed.is_set() + await asyncio.wait_for(stream.aclose(), timeout=1) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(stream), timeout=1) + async def test_stream_execute_propagates_generator_error(self): def stream_func(request): async def gen():