diff --git a/Cargo.lock b/Cargo.lock index 1ad534707..2a1123651 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1714,6 +1714,7 @@ dependencies = [ name = "nemo-relay" version = "0.8.0" dependencies = [ + "async-stream", "bitflags", "chrono", "futures", @@ -1886,6 +1887,7 @@ dependencies = [ name = "nemo-relay-plugin" version = "0.8.0" dependencies = [ + "futures", "nemo-relay-types", "serde", "serde_json", diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index b4ec5452d..fb5d1afb4 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -63,6 +63,7 @@ tokio = { version = "1", default-features = false, features = ["rt", "rt-multi-t tokio-stream = { version = "0.1", default-features = false, features = ["sync"] } typed-builder = "0.23.2" futures-util = "0.3" +async-stream = "0.3" opentelemetry = { workspace = true, features = ["trace"] } opentelemetry-semantic-conventions.workspace = true opentelemetry_sdk = { workspace = true, features = ["trace", "internal-logs"] } diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index f398ebe93..611e1aeca 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -28,7 +28,8 @@ use crate::api::runtime::subscriber_dispatcher::{ use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, - MiddlewareContinuationContext, with_active_event_uuid, + MiddlewareContinuationContext, targeted_llm_execution, targeted_llm_stream_execution, + with_active_event_uuid, }; use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; @@ -1494,7 +1495,11 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.llm_build_execution_chain(&execution_name, func, &scope_locals) + state.llm_build_execution_chain( + &execution_name, + targeted_llm_execution(func), + &scope_locals, + ) }; execution(intercepted_request).await }), @@ -1703,7 +1708,11 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.llm_stream_build_execution_chain(&execution_name, func, &scope_locals) + state.llm_stream_build_execution_chain( + &execution_name, + targeted_llm_stream_execution(func), + &scope_locals, + ) }; let execution_context = MiddlewareContinuationContext::capture(); execution(intercepted_request) diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 12d6612c3..52e065df1 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -6,6 +6,7 @@ pub mod callbacks; mod continuation_context; pub mod global; +mod llm_dispatch_context; pub mod scope_stack; pub mod state; pub mod subscriber_dispatcher; @@ -23,6 +24,11 @@ pub use continuation_context::MiddlewareContinuationContext; #[cfg(test)] pub(crate) use continuation_context::MiddlewareContinuationLease; pub use global::global_context; +#[cfg(test)] +pub(crate) use llm_dispatch_context::current_llm_dispatch_target; +pub(crate) use llm_dispatch_context::{ + LlmDispatchTargetContext, targeted_llm_execution, targeted_llm_stream_execution, +}; pub use scope_stack::{ PropagationContext, ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding, capture_propagation_context, capture_propagation_context_with_root, capture_thread_scope_stack, diff --git a/crates/core/src/api/runtime/continuation_context.rs b/crates/core/src/api/runtime/continuation_context.rs index 7ce3a1c0c..5257bac53 100644 --- a/crates/core/src/api/runtime/continuation_context.rs +++ b/crates/core/src/api/runtime/continuation_context.rs @@ -8,6 +8,9 @@ use std::future::Future; use crate::api::optimization::{ LlmOptimizationRecorder, current_llm_optimization_recorder, scope_llm_optimization_recorder, }; +use crate::api::runtime::llm_dispatch_context::{ + LlmDispatchTargetContext, current_llm_dispatch_target, scope_llm_dispatch_target, +}; use crate::api::runtime::scope_stack::{ ScopeStackHandle, TASK_SCOPE_STACK, active_event_uuid, current_context_scope_stack, current_scope_stack, scope_stack_active, snapshot_scope_stack, with_active_event_uuid, @@ -31,6 +34,7 @@ pub struct MiddlewareContinuationContext { publication_context: Option, publication_buffer: Option, optimization_recorder: Option, + llm_dispatch_target: Option, } impl MiddlewareContinuationContext { @@ -44,6 +48,7 @@ impl MiddlewareContinuationContext { publication_context: capture_publication_context(), publication_buffer: capture_nested_publication_buffer(), optimization_recorder: current_llm_optimization_recorder(), + llm_dispatch_target: current_llm_dispatch_target(), } } @@ -69,6 +74,7 @@ impl MiddlewareContinuationContext { publication_context: self.publication_context.clone(), publication_buffer: self.publication_buffer.clone(), optimization_recorder: self.optimization_recorder.clone(), + llm_dispatch_target: self.llm_dispatch_target.clone(), }) } @@ -97,12 +103,36 @@ impl MiddlewareContinuationContext { None => published.await, } }; - match &self.optimization_recorder { - Some(recorder) => scope_llm_optimization_recorder(recorder.clone(), active).await, - None => active.await, + let optimized = async { + match &self.optimization_recorder { + Some(recorder) => scope_llm_optimization_recorder(recorder.clone(), active).await, + None => active.await, + } + }; + match &self.llm_dispatch_target { + Some(target) => { + scope_llm_dispatch_target(self.active_event_uuid, target.clone(), optimized).await + } + None => optimized.await, } } + /// Invoke a callback and poll its future with the captured Relay context and typed LLM target. + #[doc(hidden)] + pub(crate) async fn invoke_with_llm_dispatch_target( + &self, + target: LlmDispatchTargetContext, + callback: C, + ) -> F::Output + where + C: FnOnce() -> F, + F: Future, + { + let mut context = self.clone(); + context.llm_dispatch_target = Some(target); + context.invoke(callback).await + } + /// Invoke a callback and poll its future with the captured Relay context. /// /// The callback itself can inspect Relay task state before constructing its diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs new file mode 100644 index 000000000..decdc75d4 --- /dev/null +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -0,0 +1,422 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Invocation-scoped target and core HTTP transport for managed LLM continuations. + +use std::collections::BTreeMap; +use std::fmt; +use std::future::Future; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use async_stream::stream; +use futures_util::StreamExt; +use reqwest::header::{self, HeaderMap, HeaderName, HeaderValue}; +use reqwest::{Client, StatusCode, Url}; + +use crate::api::llm::LlmRequest; +use crate::api::runtime::scope_stack::active_event_uuid; +use crate::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn}; +use crate::codec::streaming::SseEventDecoder; +use crate::error::{ + FlowError, MAX_UPSTREAM_FAILURE_BODY_BYTES, Result, UpstreamFailure, UpstreamFailureClass, + bounded_utf8, sanitize_upstream_failure_headers, +}; +use crate::json::Json; + +const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(30); +const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(300); +const HTTP_READ_TIMEOUT: Duration = Duration::from_secs(300); +const MAX_BUFFERED_SUCCESS_BODY_BYTES: usize = 16 * 1024 * 1024; +tokio::task_local! { + static TASK_LLM_DISPATCH_TARGET: LlmDispatchTargetBinding; +} + +#[derive(Clone)] +struct LlmDispatchTargetBinding { + // The active LLM event identifies the exact managed continuation chain. + // Nested managed calls install their own event UUID and cannot consume it. + active_event_uuid: Option, + target: LlmDispatchTargetContext, +} + +/// Validated provider transport target bound to one LLM continuation invocation. +/// +/// The target stays outside [`crate::api::llm::LlmRequest`] so credentials and +/// transport routing cannot leak into provider JSON or observability payloads. +#[doc(hidden)] +#[derive(Clone)] +#[cfg_attr(test, derive(PartialEq, Eq))] +pub(crate) struct LlmDispatchTargetContext { + url: Url, + headers: HeaderMap, +} + +impl LlmDispatchTargetContext { + /// Validate and construct a target for one continuation invocation. + pub(crate) fn try_new(url: String, headers: BTreeMap) -> Result { + let url = Url::parse(&url).map_err(|_| invalid_target_url())?; + if !matches!(url.scheme(), "http" | "https") + || !url.has_host() + || !url.username().is_empty() + || url.password().is_some() + { + return Err(invalid_target_url()); + } + let mut validated_headers = HeaderMap::new(); + for (name, value) in headers { + let name = HeaderName::from_bytes(name.as_bytes()).map_err(|_| { + FlowError::InvalidArgument( + "LLM continuation contained an invalid target header name".into(), + ) + })?; + if prohibited_target_header(&name) { + return Err(FlowError::InvalidArgument(format!( + "LLM continuation target header {name} is host-owned or prohibited" + ))); + } + if validated_headers.contains_key(&name) { + return Err(FlowError::InvalidArgument(format!( + "LLM continuation target header {name} was specified more than once" + ))); + } + let value = HeaderValue::from_str(&value).map_err(|_| { + FlowError::InvalidArgument(format!( + "LLM continuation target header {name} had an invalid value" + )) + })?; + validated_headers.insert(name, value); + } + validated_headers + .entry(header::CONTENT_TYPE) + .or_insert(HeaderValue::from_static("application/json")); + Ok(Self { + url, + headers: validated_headers, + }) + } +} + +impl fmt::Debug for LlmDispatchTargetContext { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut redacted_url = self.url.clone(); + redacted_url.set_query(None); + redacted_url.set_fragment(None); + formatter + .debug_struct("LlmDispatchTargetContext") + .field("url", &redacted_url) + .field( + "header_names", + &self + .headers + .keys() + .map(HeaderName::as_str) + .collect::>(), + ) + .finish() + } +} + +fn invalid_target_url() -> FlowError { + FlowError::InvalidArgument( + "LLM continuation target must be an absolute HTTP(S) URL without user info".into(), + ) +} + +fn prohibited_target_header(name: &HeaderName) -> bool { + let name = name.as_str(); + name.starts_with("x-nemo-relay-internal-") + || matches!( + name, + "host" + | "content-length" + | "connection" + | "transfer-encoding" + | "upgrade" + | "proxy-connection" + | "keep-alive" + | "trailer" + | "te" + ) +} + +pub(crate) fn current_llm_dispatch_target() -> Option { + TASK_LLM_DISPATCH_TARGET + .try_with(|binding| { + (binding.active_event_uuid == active_event_uuid()).then(|| binding.target.clone()) + }) + .ok() + .flatten() +} + +/// Poll a future with one typed target bound to its continuation invocation. +pub(crate) async fn scope_llm_dispatch_target( + event_uuid: Option, + target: LlmDispatchTargetContext, + future: F, +) -> F::Output { + TASK_LLM_DISPATCH_TARGET + .scope( + LlmDispatchTargetBinding { + active_event_uuid: event_uuid, + target, + }, + future, + ) + .await +} + +/// Wrap a host callback with core-owned targeted dispatch at the terminal step. +pub(crate) fn targeted_llm_execution(fallback: LlmExecutionNextFn) -> LlmExecutionNextFn { + Arc::new(move |request| { + let fallback = fallback.clone(); + Box::pin(async move { + match current_llm_dispatch_target() { + Some(target) => dispatch_buffered(&target, request).await, + None => fallback(request).await, + } + }) + }) +} + +/// Wrap a streaming host callback with core-owned targeted dispatch at the terminal step. +pub(crate) fn targeted_llm_stream_execution( + fallback: LlmStreamExecutionNextFn, +) -> LlmStreamExecutionNextFn { + Arc::new(move |request| { + let fallback = fallback.clone(); + Box::pin(async move { + match current_llm_dispatch_target() { + Some(target) => dispatch_stream(&target, request).await, + None => fallback(request).await, + } + }) + }) +} + +async fn dispatch_buffered(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result { + let response = send(target, request, Some(HTTP_REQUEST_TIMEOUT)).await?; + let status = response.status(); + if !status.is_success() { + let headers = safe_failure_headers(response.headers()); + let bytes = bounded_response_body(target, response).await?; + return Err(http_error(status, headers, &bytes)); + } + let bytes = bounded_success_body(target, response, MAX_BUFFERED_SUCCESS_BODY_BYTES).await?; + serde_json::from_slice(&bytes).map_err(|_| { + FlowError::Internal("targeted LLM provider returned malformed response JSON".into()) + }) +} + +async fn dispatch_stream( + target: &LlmDispatchTargetContext, + request: LlmRequest, +) -> Result { + let response = send(target, request, None).await?; + let status = response.status(); + if !status.is_success() { + let headers = safe_failure_headers(response.headers()); + let body = bounded_response_body(target, response).await?; + return Err(http_error(status, headers, &body)); + } + if !is_event_stream_content_type(response.headers()) { + let content_type = response + .headers() + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_owned(); + let body = bounded_response_body(target, response).await?; + return Err(unexpected_stream_content_type_error(&content_type, &body)); + } + + let target = target.clone(); + let mut decoder = SseEventDecoder::new(); + let mut bytes = response.bytes_stream(); + Ok(LlmJsonStream::new(stream! { + while let Some(chunk) = bytes.next().await { + match chunk { + Ok(buffer) => { + for result in decoder.push_bytes_results(&buffer) { + match result { + Ok(event) => yield Ok(event.data), + Err(error) => { + yield Err(error); + return; + } + } + } + } + Err(error) => { + yield Err(transport_error(&target, error)); + return; + } + } + } + match decoder.finish() { + Ok(Some(event)) => yield Ok(event.data), + Ok(None) => {} + Err(error) => yield Err(error), + } + })) +} + +fn is_event_stream_content_type(headers: &HeaderMap) -> bool { + headers + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .is_some_and(|media_type| media_type.trim().eq_ignore_ascii_case("text/event-stream")) +} + +fn unexpected_stream_content_type_error(content_type: &str, body: &[u8]) -> FlowError { + let detail = String::from_utf8_lossy(body); + let diagnostic = if detail.trim().is_empty() { + format!( + "targeted LLM provider expected Content-Type text/event-stream, received {content_type}" + ) + } else { + format!( + "targeted LLM provider expected Content-Type text/event-stream, received {content_type}; provider response body: {detail}" + ) + }; + FlowError::Internal(bounded_utf8(diagnostic, MAX_UPSTREAM_FAILURE_BODY_BYTES)) +} + +async fn send( + target: &LlmDispatchTargetContext, + request: LlmRequest, + timeout: Option, +) -> Result { + let body = serde_json::to_vec(&request.content) + .map_err(|error| FlowError::InvalidArgument(error.to_string()))?; + let mut outbound = targeted_http_client().post(target.url.clone()).body(body); + for (name, value) in &target.headers { + outbound = outbound.header(name, value); + } + if let Some(timeout) = timeout { + outbound = outbound.timeout(timeout); + } + outbound + .send() + .await + .map_err(|error| transport_error(target, error)) +} + +fn targeted_http_client() -> &'static Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + Client::builder() + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .read_timeout(HTTP_READ_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("core targeted LLM HTTP client configuration is valid") + }) +} + +async fn bounded_success_body( + target: &LlmDispatchTargetContext, + response: reqwest::Response, + max_bytes: usize, +) -> Result> { + if response + .content_length() + .is_some_and(|length| length > max_bytes as u64) + { + return Err(success_body_too_large(target, max_bytes)); + } + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|error| transport_error(target, error))?; + if body.len().saturating_add(chunk.len()) > max_bytes { + return Err(success_body_too_large(target, max_bytes)); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + +fn success_body_too_large(target: &LlmDispatchTargetContext, max_bytes: usize) -> FlowError { + log::warn!( + target: "nemo_relay.runtime", + event = "targeted_llm_response_too_large", + provider_host = target.url.host_str().unwrap_or(""), + max_bytes; + "Targeted LLM provider response exceeded the buffered body limit" + ); + FlowError::Internal(format!( + "targeted LLM provider response exceeded the {max_bytes}-byte buffered body limit" + )) +} + +async fn bounded_response_body( + target: &LlmDispatchTargetContext, + response: reqwest::Response, +) -> Result> { + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while body.len() < MAX_UPSTREAM_FAILURE_BODY_BYTES { + let Some(chunk) = stream.next().await else { + break; + }; + let chunk = chunk.map_err(|error| transport_error(target, error))?; + let remaining = MAX_UPSTREAM_FAILURE_BODY_BYTES - body.len(); + body.extend_from_slice(&chunk[..chunk.len().min(remaining)]); + } + Ok(body) +} + +fn transport_error(target: &LlmDispatchTargetContext, error: reqwest::Error) -> FlowError { + let timeout = error.is_timeout(); + let diagnostic = error.without_url(); + log::warn!( + target: "nemo_relay.runtime", + event = "targeted_llm_transport_failed", + provider_host = target.url.host_str().unwrap_or(""), + failure_kind = if timeout { "timeout" } else { "transport" }; + "Targeted LLM provider request failed: {diagnostic}" + ); + FlowError::Upstream(UpstreamFailure { + status: None, + body: if timeout { + "provider request timed out".into() + } else { + "provider transport failed".into() + }, + headers: BTreeMap::new(), + class: if timeout { + UpstreamFailureClass::Timeout + } else { + UpstreamFailureClass::Connection + }, + }) +} + +fn http_error(status: StatusCode, headers: BTreeMap, body: &[u8]) -> FlowError { + let body = String::from_utf8_lossy(&body[..body.len().min(MAX_UPSTREAM_FAILURE_BODY_BYTES)]); + FlowError::Upstream(UpstreamFailure { + status: Some(status.as_u16()), + body: body.into_owned(), + headers, + class: if matches!(status.as_u16(), 408 | 425 | 429 | 500 | 502 | 503 | 504) { + UpstreamFailureClass::RetryableStatus + } else { + UpstreamFailureClass::Other + }, + }) +} + +fn safe_failure_headers(headers: &HeaderMap) -> BTreeMap { + sanitize_upstream_failure_headers(headers.iter().map(|(name, value)| { + ( + name.as_str().to_owned(), + String::from_utf8_lossy(value.as_bytes()).into_owned(), + ) + })) +} + +#[cfg(test)] +#[path = "../../../tests/unit/llm_dispatch_context_tests.rs"] +mod tests; diff --git a/crates/core/src/codec/streaming.rs b/crates/core/src/codec/streaming.rs index e1bbbd777..8bffbf1f3 100644 --- a/crates/core/src/codec/streaming.rs +++ b/crates/core/src/codec/streaming.rs @@ -27,6 +27,11 @@ use crate::error::{FlowError, Result}; use crate::json::Json; use serde::{Deserialize, Serialize}; +// Bound one provider event independently of Relay's bounded event queue. Keep the limit generous +// for large tool or multimodal deltas; without a terminator there is no event to enqueue and +// backpressure cannot apply. +const MAX_SSE_FRAME_BYTES: usize = 8 * 1024 * 1024; + /// Provider-neutral incremental stream item used by cross-protocol transcoders. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -98,7 +103,7 @@ pub trait StreamingCodec: Send + Sync { /// terminator are retained for the next call. #[derive(Default)] pub struct SseEventDecoder { - buffer: String, + buffer: Vec, } /// One decoded SSE frame, paired with the parsed `data:` payload. @@ -118,15 +123,13 @@ impl SseEventDecoder { /// Appends `bytes` to the internal buffer and returns every now-complete SSE event. /// - /// Bytes are interpreted as UTF-8 with replacement characters for invalid sequences; provider - /// SSE streams are well-formed UTF-8 in practice, but lossy decoding keeps the decoder honest - /// rather than failing on a single corrupt chunk. + /// UTF-8 is validated only after a complete frame arrives, so a multibyte code point may span + /// arbitrary transport chunks without being replaced or corrupted. /// /// Returns `Ok(events)` containing zero or more events whose `data:` payloads parsed - /// successfully. Frames whose `data:` line is non-empty but does not parse as JSON are - /// surfaced as [`FlowError::Internal`] so the caller can decide whether to abort the stream - /// or skip the frame; frames with no `data:` line at all (e.g. SSE heartbeats) are silently - /// dropped. + /// successfully. Invalid UTF-8, oversized frames, and non-empty `data:` payloads that do not + /// parse as JSON are surfaced as [`FlowError::Internal`] so the caller can abort the stream; + /// frames with no `data:` line at all (e.g. SSE heartbeats) are silently dropped. pub fn push_bytes(&mut self, bytes: &[u8]) -> Result> { self.push_bytes_results(bytes).into_iter().collect() } @@ -141,25 +144,66 @@ impl SseEventDecoder { // providers emit mixed line endings on the wire; normalizing once here keeps the inner // loop cheap. If CRLF is split across chunks, retain the trailing CR until the next append // and remove it only when the next byte completes the sequence. - if self.buffer.ends_with('\r') && bytes.first() == Some(&b'\n') { + // The previous call drained every complete frame, so a new delimiter can begin only at + // the former trailing byte or inside this chunk. Avoid rescanning a growing partial frame + // from its beginning on every network read. + let mut scan_from = self.buffer.len().saturating_sub(1); + let mut offset = 0; + if self.buffer.last() == Some(&b'\r') && bytes.first() == Some(&b'\n') { self.buffer.pop(); + scan_from = self.buffer.len().saturating_sub(1); + self.buffer.push(b'\n'); + offset = 1; + } + while offset < bytes.len() { + if bytes[offset] == b'\r' && bytes.get(offset + 1) == Some(&b'\n') { + self.buffer.push(b'\n'); + offset += 2; + } else { + self.buffer.push(bytes[offset]); + offset += 1; + } } - let chunk = String::from_utf8_lossy(bytes).replace("\r\n", "\n"); - self.buffer.push_str(&chunk); + let mut results = Vec::new(); - while let Some(cut) = self.buffer.find("\n\n") { - let frame: String = self.buffer.drain(..cut).collect(); + while let Some(relative_cut) = self.buffer[scan_from..] + .windows(2) + .position(|pair| pair == b"\n\n") + { + let cut = scan_from + relative_cut; + if cut > MAX_SSE_FRAME_BYTES { + self.buffer.clear(); + results.push(Err(oversized_sse_frame_error())); + return results; + } + let frame: Vec = self.buffer.drain(..cut).collect(); // Drop the `\n\n` terminator itself. self.buffer.drain(..2); - match parse_sse_frame(&frame) { + scan_from = 0; + let frame = match std::str::from_utf8(&frame) { + Ok(frame) => frame, + Err(error) => { + self.buffer.clear(); + results.push(Err(FlowError::Internal(format!( + "streaming codec received invalid UTF-8 SSE frame: {error}" + )))); + return results; + } + }; + match parse_sse_frame(frame) { Ok(Some(event)) => results.push(Ok(event)), Ok(None) => {} Err(error) => { + self.buffer.clear(); results.push(Err(error)); - break; + return results; } } } + if self.buffer.len() > MAX_SSE_FRAME_BYTES { + self.buffer.clear(); + results.push(Err(oversized_sse_frame_error())); + } results } @@ -170,14 +214,25 @@ impl SseEventDecoder { /// captures the last bytes the upstream sent before disconnect. pub fn finish(mut self) -> Result> { let trailing = std::mem::take(&mut self.buffer); + let trailing = std::str::from_utf8(&trailing).map_err(|error| { + FlowError::Internal(format!( + "streaming codec received incomplete or invalid UTF-8 at end of SSE stream: {error}" + )) + })?; if trailing.trim().is_empty() { Ok(None) } else { - parse_sse_frame(&trailing) + parse_sse_frame(trailing) } } } +fn oversized_sse_frame_error() -> FlowError { + FlowError::Internal(format!( + "streaming codec SSE frame exceeded the {MAX_SSE_FRAME_BYTES}-byte limit" + )) +} + // Parses a single SSE frame. Returns `None` for frames without a `data:` line, `Some(event)` for // frames whose `data:` JSON parsed successfully. fn parse_sse_frame(frame: &str) -> Result> { diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 270f9737a..29585b2ec 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -12,7 +12,57 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use thiserror::Error; -/// Stable classification for a failure from an upstream provider attempt. +pub(crate) const MAX_UPSTREAM_FAILURE_BODY_BYTES: usize = 16 * 1024; +pub(crate) const MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES: usize = 1024; + +pub(crate) fn bounded_utf8(value: String, max_bytes: usize) -> String { + if value.len() <= max_bytes { + return value; + } + let mut boundary = max_bytes; + while !value.is_char_boundary(boundary) { + boundary -= 1; + } + value[..boundary].to_owned() +} + +pub(crate) fn sanitize_upstream_failure_headers( + headers: impl IntoIterator, +) -> BTreeMap { + headers + .into_iter() + .filter_map(|(name, value)| { + let normalized = name.to_ascii_lowercase(); + matches!( + normalized.as_str(), + "retry-after" + | "request-id" + | "traceparent" + | "x-request-id" + | "x-ratelimit-limit" + | "x-ratelimit-remaining" + | "x-ratelimit-reset" + | "ratelimit-limit" + | "ratelimit-remaining" + | "ratelimit-reset" + ) + .then(|| { + ( + normalized, + bounded_utf8(value, MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES), + ) + }) + }) + .collect() +} + +pub(crate) fn sanitize_upstream_failure(mut failure: UpstreamFailure) -> UpstreamFailure { + failure.body = bounded_utf8(failure.body, MAX_UPSTREAM_FAILURE_BODY_BYTES); + failure.headers = sanitize_upstream_failure_headers(failure.headers); + failure +} + +/// Stable classification for an upstream provider failure captured by managed dispatch. #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub enum UpstreamFailureClass { @@ -34,7 +84,7 @@ pub enum UpstreamFailureClass { Other, } -/// Structured failure returned by one upstream provider attempt. +/// Structured provider failure surfaced by targeted or explicitly retry-aware dispatch. #[derive(Clone, Debug, Deserialize, Serialize)] pub struct UpstreamFailure { /// HTTP status when a provider response was received. @@ -48,7 +98,7 @@ pub struct UpstreamFailure { } impl UpstreamFailure { - /// Whether Switchyard may be consulted for another bounded provider attempt. + /// Whether a provider-neutral routing policy may make another bounded attempt. pub fn is_retryable(&self) -> bool { matches!( self.class, @@ -122,7 +172,7 @@ pub enum FlowError { #[error("guardrail rejected: {0}")] GuardrailRejected(String), - /// Structured upstream provider failure from retry-aware gateway dispatch. + /// Structured upstream provider failure from managed provider dispatch. #[error("{0}")] Upstream(UpstreamFailure), diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 3b0124b5b..a474d6d47 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -16,7 +16,7 @@ use std::path::{Path, PathBuf}; use std::pin::Pin; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::task::{Context, Poll}; use futures_util::FutureExt; @@ -24,15 +24,17 @@ use futures_util::FutureExt; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionFn, - LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, - LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, + LlmDispatchTargetContext, LlmExecutionFn, LlmExecutionNextFn, LlmJsonStream, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, MiddlewareContinuationContext, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::{ ScopeStackHandle, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack, - restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack, with_scope_stack, + current_scope_stack, restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack, + with_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, @@ -41,7 +43,9 @@ use crate::api::scope::{event as emit_scope_mark, get_handle, pop_scope, push_sc use crate::api::tool::ToolExecutionInterceptOutcome; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; -use crate::error::{FlowError, Result as FlowResult}; +use crate::error::{ + FlowError, Result as FlowResult, UpstreamFailureClass, bounded_utf8, sanitize_upstream_failure, +}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, deregister_plugin_registration_checked, register_plugin_tracked, @@ -49,22 +53,28 @@ use crate::plugin::{ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ + LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmNonHttpFailureKindV2, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, - NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, - NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, + NemoRelayNativeAsyncLlmStreamForwardCbV2, NemoRelayNativeAsyncLlmStreamNextCbV2, + NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeAsyncTaskPollCbV2, + NemoRelayNativeAsyncTaskV2, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, + NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, - NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, - NemoRelayNativePluginEntry, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, - NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, - NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, - NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, + NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativeLlmStreamV2, + NemoRelayNativePluginContext, NemoRelayNativePluginEntry, NemoRelayNativePluginV1, + NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, + NemoRelayNativeScopeType, NemoRelayNativeString, NemoRelayNativeToolConditionalCb, + NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, + NemoRelayStatus, }; use semver::{Version, VersionReq}; use serde_json::{Map, Value as Json}; @@ -90,8 +100,9 @@ pub struct NativePluginLoadSpec { /// Owns native dynamic libraries registered into the plugin registry. /// /// Dropping this value deregisters the native plugin kinds before unloading -/// their libraries. Clear active plugin configuration before dropping it so -/// runtime callbacks cannot outlive their code. +/// their libraries. Clear active plugin configuration before dropping it. If +/// an opaque native handle still owns plugin code after deregistration, Relay +/// conservatively keeps that library mapped until process exit. pub struct NativePluginActivation { plugins: Vec>, plugin_registrations: Vec<(String, u64)>, @@ -107,7 +118,24 @@ impl NativePluginActivation { pub fn clear(self) {} pub(crate) fn deregister_plugin_kinds_checked(&mut self) -> DynamicPluginTeardownOutcome { - deregister_tracked_registrations_checked(&mut self.plugin_registrations, "native") + let outcome = + deregister_tracked_registrations_checked(&mut self.plugin_registrations, "native"); + self.retain_libraries_for_outstanding_references(); + outcome + } + + fn retain_libraries_for_outstanding_references(&self) { + for plugin in &self.plugins { + // Deregistration removes the registry adapter's Arc. The + // activation owns the one remaining expected reference, so any + // additional reference belongs to an escaped callback, task, + // continuation, or provider stream. Its final host release may + // run from plugin code; unloading there would unmap the caller + // before the FFI function can return. + if Arc::strong_count(plugin) > 1 { + plugin.retain_library_on_drop.store(true, Ordering::Release); + } + } } #[cfg(test)] @@ -124,6 +152,7 @@ impl Drop for NativePluginActivation { for (plugin_kind, registration_id) in self.plugin_registrations.iter().rev() { let _ = deregister_plugin_registration_checked(plugin_kind, *registration_id); } + self.retain_libraries_for_outstanding_references(); } } @@ -279,7 +308,8 @@ struct NativePluginInstance { relay_compat: String, allows_multiple_components: bool, plugin: Mutex, - _library: Library, + library: Option, + retain_library_on_drop: AtomicBool, } unsafe impl Send for NativePluginInstance {} @@ -290,6 +320,13 @@ impl Drop for NativePluginInstance { if let Ok(mut plugin) = self.plugin.lock() { drop_native_plugin_descriptor(&mut plugin); } + if self.retain_library_on_drop.load(Ordering::Acquire) + && let Some(library) = self.library.take() + { + // Only the OS library handle is retained. The plugin descriptor + // and every Relay-owned allocation above are still released. + std::mem::forget(library); + } } } @@ -327,11 +364,12 @@ fn load_one_native_plugin( .as_deref() .expect("validated native manifest must declare compat.relay") .to_string(); - if manifest.compat.native_api.as_deref().map(str::trim) != Some("1") { + let native_api = manifest.compat.native_api.as_deref().map(str::trim); + if !matches!(native_api, Some("1" | "2")) { return Err(PluginError::InvalidConfig(format!( - "dynamic plugin '{}' declares unsupported compat.native_api '{}'; expected 1", + "dynamic plugin '{}' declares unsupported compat.native_api '{}'; expected 1 or 2", spec.plugin_id, - manifest.compat.native_api.as_deref().unwrap_or("") + native_api.unwrap_or("") ))); } let DynamicPluginManifestLoad::RustDynamic(load) = &manifest.load else { @@ -380,13 +418,22 @@ fn load_one_native_plugin( library_path.display() )) })?; - let mut status = entry(native_host_api(), &mut plugin); - // SDKs compiled against ABI v2 correctly reject a v3 table. Retry - // their entry point with the frozen v2 prefix instead of making a - // runtime upgrade a breaking change for installed native plugins. - if status == NemoRelayStatus::InvalidArg { + let host_apis: &[*const NemoRelayNativeHostApiV1] = match native_api { + Some("1") => &[native_host_api_v3(), native_host_api_v2()], + Some("2") => &[native_host_api()], + _ => unreachable!("native API version was validated"), + }; + let mut status = NemoRelayStatus::InvalidArg; + for (index, host_api) in host_apis.iter().enumerate() { + status = entry(*host_api, &mut plugin); + if status == NemoRelayStatus::Ok { + break; + } + if status != NemoRelayStatus::InvalidArg || index + 1 == host_apis.len() { + break; + } drop_native_plugin_descriptor(&mut plugin); - status = entry(native_host_api_legacy(), &mut plugin); + plugin = NemoRelayNativePluginV1::default(); } if status != NemoRelayStatus::Ok { drop_native_plugin_descriptor(&mut plugin); @@ -419,7 +466,8 @@ fn load_one_native_plugin( relay_compat, allows_multiple_components: plugin.allows_multiple_components, plugin: Mutex::new(plugin), - _library: library, + library: Some(library), + retain_library_on_drop: AtomicBool::new(false), })) } @@ -797,16 +845,21 @@ unsafe extern "C" fn native_llm_response_codec_decode( } fn native_host_api() -> *const NemoRelayNativeHostApiV1 { + static HOST_API: OnceLock = OnceLock::new(); + &HOST_API.get_or_init(build_native_host_api_v4).v3.v1 as *const NemoRelayNativeHostApiV1 +} + +fn native_host_api_v3() -> *const NemoRelayNativeHostApiV1 { static HOST_API: OnceLock = OnceLock::new(); &HOST_API.get_or_init(build_native_host_api_v3).v1 as *const NemoRelayNativeHostApiV1 } -fn native_host_api_legacy() -> *const NemoRelayNativeHostApiV1 { +fn native_host_api_v2() -> *const NemoRelayNativeHostApiV1 { static HOST_API: OnceLock = OnceLock::new(); - HOST_API.get_or_init(build_native_host_api_legacy) as *const _ + HOST_API.get_or_init(build_native_host_api_v2) as *const _ } -fn build_native_host_api_legacy() -> NemoRelayNativeHostApiV1 { +fn build_native_host_api_v2() -> NemoRelayNativeHostApiV1 { static RELAY_VERSION: &[u8] = concat!(env!("CARGO_PKG_VERSION"), "\0").as_bytes(); NemoRelayNativeHostApiV1 { abi_version: NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, @@ -867,7 +920,7 @@ fn build_native_host_api_legacy() -> NemoRelayNativeHostApiV1 { } fn build_native_host_api_v3() -> NemoRelayNativeHostApiV3 { - let mut v1 = build_native_host_api_legacy(); + let mut v1 = build_native_host_api_v2(); v1.abi_version = NEMO_RELAY_NATIVE_ABI_VERSION; v1.struct_size = std::mem::size_of::(); NemoRelayNativeHostApiV3 { @@ -891,6 +944,27 @@ fn build_native_host_api_v3() -> NemoRelayNativeHostApiV3 { } } +fn build_native_host_api_v4() -> NemoRelayNativeHostApiV4 { + let mut v3 = build_native_host_api_v3(); + v3.v1.abi_version = NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS; + v3.v1.struct_size = std::mem::size_of::(); + v3.async_stream_push_json = native_async_stream_push_json_v2; + v3.async_stream_reject = native_async_stream_reject_v2; + NemoRelayNativeHostApiV4 { + v3, + async_llm_next_invoke_result_v2: native_async_llm_next_invoke_result_v2, + async_llm_next_open_stream_v2: native_async_llm_next_open_stream_v2, + async_llm_stream_next_v2: native_async_llm_stream_next_v2, + async_llm_stream_release_v2: native_async_llm_stream_release_v2, + async_completion_spawn_task_v2: native_async_completion_spawn_task_v2, + async_stream_spawn_task_v2: native_async_stream_spawn_task_v2, + async_task_retain_v2: native_async_task_retain_v2, + async_task_wake_v2: native_async_task_wake_v2, + async_task_release_v2: native_async_task_release_v2, + async_llm_next_forward_stream_v2: native_async_llm_next_forward_stream_v2, + } +} + fn read_native_string(value: *const NemoRelayNativeString) -> crate::plugin::Result { if value.is_null() { return Ok(String::new()); @@ -1368,6 +1442,56 @@ impl Drop for NativeCallbackUserDataGuard { } } +struct NativeInvocationStringGuard(usize); + +impl Drop for NativeInvocationStringGuard { + fn drop(&mut self) { + unsafe { native_string_free(self.0 as *mut NemoRelayNativeString) }; + } +} + +struct NativeCompletionHandoff { + raw: usize, + armed: bool, +} + +impl NativeCompletionHandoff { + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for NativeCompletionHandoff { + fn drop(&mut self) { + if self.armed { + unsafe { + native_async_completion_release(self.raw as *const NemoRelayNativeAsyncCompletion) + }; + } + } +} + +struct NativeNextHandoff { + raw: Option, + armed: bool, +} + +impl NativeNextHandoff { + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for NativeNextHandoff { + fn drop(&mut self) { + if self.armed + && let Some(raw) = self.raw + { + unsafe { native_async_next_release(raw as *const NemoRelayNativeAsyncNext) }; + } + } +} + unsafe impl Send for NativeCallbackUserData {} unsafe impl Sync for NativeCallbackUserData {} @@ -1391,6 +1515,8 @@ fn make_user_data( }) } +// Incremental plugin output is bounded for backpressure. Keep the established +// native API v1 capacity for every native plugin API version. const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY: usize = 64; struct NativeAsyncCompletion { @@ -1398,6 +1524,9 @@ struct NativeAsyncCompletion { cancelled: AtomicBool, next_invoked: AtomicBool, next_abort: Mutex>, + runtime: tokio::runtime::Handle, + context: MiddlewareContinuationContext, + task: Mutex>>, #[cfg(test)] before_settlement_lock: Option>, // A pending native callback can continue running after its completion @@ -1436,6 +1565,8 @@ impl Drop for NativeAsyncWait { if let Some(abort) = next_abort.take() { abort.abort(); } + drop(next_abort); + cancel_native_async_task_slot(&self.completion.task); } } @@ -1449,6 +1580,7 @@ struct NativeAsyncNext { inner: NativeAsyncNextInner, runtime: tokio::runtime::Handle, context: MiddlewareContinuationContext, + in_flight_aborts: Arc>>, // The native callback owns this handle independently of its completion. // Retaining the library here prevents an unload while it still uses `next`. _callback_user_data: Option>, @@ -1464,22 +1596,274 @@ impl NativeAsyncNext { inner, runtime, context: MiddlewareContinuationContext::capture(), + in_flight_aborts: Arc::new(Mutex::new(HashMap::new())), _callback_user_data: callback_user_data, } } } +impl Drop for NativeAsyncNext { + fn drop(&mut self) { + let mut in_flight = self + .in_flight_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + for (_, abort) in in_flight.drain() { + abort.abort(); + } + } +} + struct NativeAsyncStream { sender: Mutex>>>, cancelled: AtomicBool, settled: AtomicBool, downstream_aborts: Mutex>, settlement: Mutex<()>, + runtime: tokio::runtime::Handle, + context: MiddlewareContinuationContext, + task: Mutex>>, #[cfg(test)] before_settlement_lock: Option>, _callback_user_data: Option>, } +enum NativeAsyncTaskOwnerV2 { + Completion(Weak), + Stream(Weak), +} + +impl NativeAsyncTaskOwnerV2 { + fn is_terminal(&self) -> bool { + match self { + Self::Completion(completion) => completion.upgrade().is_none_or(|completion| { + completion.cancelled.load(Ordering::Acquire) + || completion + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + }), + Self::Stream(stream) => stream.upgrade().is_none_or(|stream| { + stream.cancelled.load(Ordering::Acquire) || stream.settled.load(Ordering::Acquire) + }), + } + } + + fn detach(&self, task: &Arc) { + match self { + Self::Completion(completion) => { + if let Some(owner) = completion.upgrade() { + clear_native_async_task_owner_slot(&owner.task, task); + } + } + Self::Stream(stream) => { + if let Some(owner) = stream.upgrade() { + clear_native_async_task_owner_slot(&owner.task, task); + } + } + } + } +} + +fn clear_native_async_task_owner_slot( + slot: &Mutex>>, + task: &Arc, +) { + let mut slot = slot.lock().unwrap_or_else(|error| error.into_inner()); + if slot + .as_ref() + .is_some_and(|current| current.as_ptr() == Arc::as_ptr(task)) + { + slot.take(); + } +} + +struct NativeAsyncTaskPluginStateV2 { + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} + +impl Drop for NativeAsyncTaskPluginStateV2 { + fn drop(&mut self) { + if let Some(free_fn) = self.free_fn { + unsafe { free_fn(self.user_data as *mut c_void) }; + } + } +} + +struct NativeAsyncTaskStateV2 { + polling: bool, + cancel_requested: bool, + complete: bool, + plugin: Option, +} + +struct NativeAsyncTaskV2 { + context: MiddlewareContinuationContext, + owner: NativeAsyncTaskOwnerV2, + wake: tokio::sync::Notify, + state: Mutex, + // Stale Rust wakers retain task references whose vtables live in the + // plugin. Keep the dynamic library loaded until the final task reference + // is released, even after per-invocation state has completed. + _library_guard: Option>, +} + +impl NativeAsyncTaskV2 { + fn wake(task: &Arc) { + if !task + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .complete + { + task.wake.notify_one(); + } + } + + async fn run(self: Arc) { + loop { + self.wake.notified().await; + if self + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .complete + { + break; + } + self.clone().poll_once().await; + if self + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .complete + { + break; + } + } + } + + async fn poll_once(self: Arc) { + let callback = { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.complete { + return; + } + state.polling = true; + state + .plugin + .as_ref() + .map(|plugin| (plugin.cb, plugin.user_data)) + }; + let Some((cb, user_data)) = callback else { + self.finish(); + return; + }; + + if self.owner.is_terminal() { + self.finish(); + return; + } + + let task = Arc::clone(&self); + let callback_result = self + .context + .run(async move { + catch_unwind(AssertUnwindSafe(|| unsafe { + cb( + user_data as *mut c_void, + Arc::as_ptr(&task) as *const NemoRelayNativeAsyncTaskV2, + ) + })) + }) + .await; + + let cancel_requested = { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.polling = false; + state.cancel_requested + }; + + let callback_state = callback_result + .ok() + .and_then(|state| NemoRelayNativeAsyncCallbackState::try_from(state).ok()); + match callback_state { + _ if cancel_requested => self.finish(), + Some(NemoRelayNativeAsyncCallbackState::Pending) if !self.owner.is_terminal() => {} + Some(NemoRelayNativeAsyncCallbackState::Pending) => self.finish(), + Some(NemoRelayNativeAsyncCallbackState::Complete) if self.owner.is_terminal() => { + self.finish(); + } + Some(NemoRelayNativeAsyncCallbackState::Complete) => { + self.settle_internal( + "native async task returned Complete without settling its owner", + ) + .await; + self.finish(); + } + None => { + self.settle_internal( + "native async task panicked or returned an invalid callback state", + ) + .await; + self.finish(); + } + } + } + + async fn settle_internal(&self, message: &str) { + match &self.owner { + NativeAsyncTaskOwnerV2::Completion(completion) => { + if let Some(completion) = completion.upgrade() { + settle_native_async_completion_error(&completion, message.to_owned()); + } + } + NativeAsyncTaskOwnerV2::Stream(stream) => { + if let Some(stream) = stream.upgrade() { + settle_native_async_stream_error(stream, message.to_owned()).await; + } + } + } + } + + fn finish(self: &Arc) { + let plugin = { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.complete { + return; + } + state.complete = true; + state.plugin.take() + }; + self.owner.detach(self); + drop(plugin); + } + + fn cancel(task: &Arc) { + let plugin = { + let mut state = task.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.complete { + return; + } + state.cancel_requested = true; + if state.polling { + None + } else { + state.complete = true; + state.plugin.take() + } + }; + if let Some(plugin) = plugin { + task.owner.detach(task); + drop(plugin); + } + task.wake.notify_one(); + } +} + struct NativeAsyncStreamReceiver { receiver: tokio::sync::mpsc::Receiver>, stream: Arc, @@ -1493,71 +1877,497 @@ struct NativeAsyncStreamCallbackGuard { active: bool, } -impl NativeAsyncStreamCallbackGuard { - fn finish(&mut self) { - self.active = false; +struct NativeLlmProviderStreamV2 { + stream: tokio::sync::Mutex>, + runtime: tokio::runtime::Handle, + context: MiddlewareContinuationContext, + target: LlmDispatchTargetContext, + output: Arc, + lifecycle: Mutex, +} + +enum NativeLlmProviderStreamLifecycleV2 { + Idle, + Pulling { task_id: tokio::task::Id }, + Terminal, + Cancelled, +} + +impl NativeLlmProviderStreamV2 { + fn cancel(&self) { + let pulling = { + let mut lifecycle = self + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + match std::mem::replace( + &mut *lifecycle, + NativeLlmProviderStreamLifecycleV2::Cancelled, + ) { + NativeLlmProviderStreamLifecycleV2::Pulling { task_id } => Some(task_id), + NativeLlmProviderStreamLifecycleV2::Terminal => { + *lifecycle = NativeLlmProviderStreamLifecycleV2::Terminal; + None + } + NativeLlmProviderStreamLifecycleV2::Idle + | NativeLlmProviderStreamLifecycleV2::Cancelled => None, + } + }; + if let Some(task_id) = pulling + && let Some(abort) = self + .output + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&task_id) + { + abort.abort(); + } } +} - fn fail(&mut self, error: &str) { - if !self.active { - return; +impl Drop for NativeLlmProviderStreamV2 { + fn drop(&mut self) { + self.cancel(); + } +} + +struct NativeLlmProviderNextCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, + user_data: usize, + provider: Arc, + active: bool, +} + +impl NativeLlmProviderNextCallbackGuardV2 { + fn complete_pull(&self, terminal: bool) -> bool { + let output_live = !self.provider.output.cancelled.load(Ordering::Acquire) + && !self.provider.output.settled.load(Ordering::Acquire); + let task_id = { + let mut lifecycle = self + .provider + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + match &*lifecycle { + NativeLlmProviderStreamLifecycleV2::Pulling { task_id } => { + let task_id = *task_id; + *lifecycle = if output_live { + if terminal { + NativeLlmProviderStreamLifecycleV2::Terminal + } else { + NativeLlmProviderStreamLifecycleV2::Idle + } + } else { + NativeLlmProviderStreamLifecycleV2::Cancelled + }; + Some(task_id) + } + _ => None, + } + }; + if let Some(task_id) = task_id { + self.provider + .output + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&task_id); } - // Cancellation owns terminal delivery. Leave the guard active so its - // Drop implementation can notify the plugin and release callback data. - if self.stream.cancelled.load(Ordering::Acquire) { - return; + task_id.is_some() && output_live + } + + fn cancel_pull(&self) { + let task_id = { + let mut lifecycle = self + .provider + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + match &*lifecycle { + NativeLlmProviderStreamLifecycleV2::Pulling { task_id } => { + let task_id = *task_id; + *lifecycle = NativeLlmProviderStreamLifecycleV2::Cancelled; + Some(task_id) + } + _ => None, + } + }; + if let Some(task_id) = task_id { + self.provider + .output + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&task_id); } - if let Some(message) = native_string_from_str(error) { - unsafe { - let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); - native_string_free(message); + } + + fn invoke_failure(&self, error: &LlmContinuationFailureV2) { + let error = native_string_from_json( + &serde_json::to_value(error) + .expect("native API v2 LLM failures contain serializable Relay DTOs"), + ); + unsafe { + (self.cb)( + self.user_data as *mut c_void, + ptr::null(), + error.unwrap_or(ptr::null_mut()), + error.is_some(), + ); + if let Some(error) = error { + native_string_free(error); } - self.active = false; } } -} -impl Drop for NativeAsyncStreamCallbackGuard { - fn drop(&mut self) { + fn complete(&mut self, outcome: std::result::Result, LlmContinuationFailureV2>) { if !self.active { return; } - if self.stream.cancelled.load(Ordering::Acquire) { - if let Some(message) = - native_string_from_str("native async stream continuation was cancelled") - { - unsafe { - let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); - native_string_free(message); + self.active = false; + match outcome { + Ok(Some(chunk)) => { + let chunk = native_string_from_json(&chunk); + let allowed = self.complete_pull(chunk.is_none()); + if !allowed { + self.invoke_failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM provider stream was cancelled".into(), + )); + } else { + unsafe { + (self.cb)( + self.user_data as *mut c_void, + chunk.unwrap_or(ptr::null_mut()), + ptr::null(), + false, + ); + if let Some(chunk) = chunk { + native_string_free(chunk); + } + } } } - } else if self.stream.settled.load(Ordering::Acquire) { - if let Some(message) = - native_string_from_str("native async stream continuation output settled") - { - unsafe { - let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); - native_string_free(message); + Ok(None) => { + if self.complete_pull(true) { + unsafe { + (self.cb)( + self.user_data as *mut c_void, + ptr::null(), + ptr::null(), + true, + ); + } + } else { + self.invoke_failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM provider stream was cancelled".into(), + )); } } - } else { - unsafe { - let _ = (self.cb)( - self.user_data as *mut c_void, - ptr::null(), - ptr::null(), - true, - ); + Err(error) => { + if self.complete_pull(true) { + self.invoke_failure(&error); + } else { + self.invoke_failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM provider stream was cancelled".into(), + )); + } } } } } -impl Stream for NativeAsyncStreamReceiver { - type Item = FlowResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.receiver.poll_recv(cx) +impl Drop for NativeLlmProviderNextCallbackGuardV2 { + fn drop(&mut self) { + if !self.active { + return; + } + self.active = false; + self.cancel_pull(); + self.invoke_failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM provider stream pull was cancelled".into(), + )); + } +} + +struct NativeLlmStreamOpenCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, + user_data: usize, + active: bool, + _library_guard: Option>, +} + +struct NativeLlmResultCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: usize, + active: bool, + _library_guard: Option>, +} + +struct NativeAsyncResultCallbackGuard { + cb: NemoRelayNativeAsyncNextResultCb, + user_data: usize, + active: bool, + _library_guard: Option>, +} + +impl NativeAsyncResultCallbackGuard { + fn complete(&mut self, result: FlowResult) { + if !self.active { + return; + } + match result { + Ok(value) => { + let value = native_string_from_json(&value); + unsafe { + (self.cb)( + self.user_data as *mut c_void, + value.unwrap_or(ptr::null_mut()), + ptr::null(), + ); + if let Some(value) = value { + native_string_free(value); + } + } + } + Err(error) => { + let error = native_string_from_str(&error.to_string()); + unsafe { + (self.cb)( + self.user_data as *mut c_void, + ptr::null(), + error.unwrap_or(ptr::null_mut()), + ); + if let Some(error) = error { + native_string_free(error); + } + } + } + } + self.active = false; + } +} + +impl Drop for NativeAsyncResultCallbackGuard { + fn drop(&mut self) { + if self.active { + self.complete(Err(FlowError::Internal( + "native continuation was cancelled".into(), + ))); + } + } +} + +impl NativeLlmResultCallbackGuardV2 { + fn complete(&mut self, result: std::result::Result) { + if !self.active { + return; + } + let (response, error) = match result { + Ok(response) => (native_string_from_json(&response), None), + Err(error) => ( + None, + native_string_from_json( + &serde_json::to_value(error) + .expect("native API v2 LLM failures contain serializable Relay DTOs"), + ), + ), + }; + unsafe { + (self.cb)( + self.user_data as *mut c_void, + response.unwrap_or(ptr::null_mut()), + error.unwrap_or(ptr::null_mut()), + ); + if let Some(response) = response { + native_string_free(response); + } + if let Some(error) = error { + native_string_free(error); + } + } + self.active = false; + } +} + +impl Drop for NativeLlmResultCallbackGuardV2 { + fn drop(&mut self) { + if self.active { + self.complete(Err(non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM continuation was cancelled".into(), + ))); + } + } +} + +impl NativeLlmStreamOpenCallbackGuardV2 { + fn success(&mut self, stream: Arc) { + if !self.active { + return; + } + let stream = Arc::into_raw(stream) as *const NemoRelayNativeLlmStreamV2; + unsafe { (self.cb)(self.user_data as *mut c_void, stream, ptr::null()) }; + self.active = false; + } + + fn failure(&mut self, error: &LlmContinuationFailureV2) { + if !self.active { + return; + } + let error = native_string_from_json( + &serde_json::to_value(error) + .expect("native API v2 LLM failures contain serializable Relay DTOs"), + ); + unsafe { + (self.cb)( + self.user_data as *mut c_void, + ptr::null(), + error.unwrap_or(ptr::null_mut()), + ); + if let Some(error) = error { + native_string_free(error); + } + } + self.active = false; + } +} + +impl Drop for NativeLlmStreamOpenCallbackGuardV2 { + fn drop(&mut self) { + if self.active { + self.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM stream setup was cancelled".into(), + )); + } + } +} + +struct NativeLlmStreamForwardCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: usize, + stream: Arc, + active: bool, +} + +impl NativeLlmStreamForwardCallbackGuardV2 { + fn settle(&mut self) { + if !self.active { + return; + } + self.active = false; + unsafe { (self.cb)(self.user_data as *mut c_void) }; + } + + fn cancel_unsettled_output(&self) { + let sender = { + let _settlement = self + .stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if self.stream.cancelled.load(Ordering::Acquire) + || self.stream.settled.load(Ordering::Acquire) + { + None + } else { + self.stream.cancelled.store(true, Ordering::Release); + self.stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + } + }; + drop(sender); + abort_native_stream_downstream_tasks(&self.stream); + } +} + +impl Drop for NativeLlmStreamForwardCallbackGuardV2 { + fn drop(&mut self) { + if !self.active { + return; + } + self.cancel_unsettled_output(); + self.settle(); + } +} + +impl NativeAsyncStreamCallbackGuard { + fn finish(&mut self) { + self.active = false; + } + + fn fail(&mut self, error: &str) { + if !self.active { + return; + } + // Cancellation owns terminal delivery. Leave the guard active so its + // Drop implementation can notify the plugin and release callback data. + if self.stream.cancelled.load(Ordering::Acquire) { + return; + } + if let Some(message) = native_string_from_str(error) { + unsafe { + let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); + native_string_free(message); + } + self.active = false; + } + } +} + +impl Drop for NativeAsyncStreamCallbackGuard { + fn drop(&mut self) { + if !self.active { + return; + } + if self.stream.cancelled.load(Ordering::Acquire) { + if let Some(message) = + native_string_from_str("native async stream continuation was cancelled") + { + unsafe { + let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); + native_string_free(message); + } + } + } else if self.stream.settled.load(Ordering::Acquire) { + if let Some(message) = + native_string_from_str("native async stream continuation output settled") + { + unsafe { + let _ = (self.cb)(self.user_data as *mut c_void, ptr::null(), message, false); + native_string_free(message); + } + } + } else { + unsafe { + let _ = (self.cb)( + self.user_data as *mut c_void, + ptr::null(), + ptr::null(), + true, + ); + } + } + } +} + +impl Stream for NativeAsyncStreamReceiver { + type Item = FlowResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let result = self.receiver.poll_recv(cx); + if result.is_ready() { + wake_native_async_stream_task(&self.stream); + } + result } } @@ -1583,6 +2393,20 @@ impl Drop for NativeAsyncStreamReceiver { .lock() .unwrap_or_else(|error| error.into_inner()) .take(); + drop(_settlement); + cancel_native_async_task_slot(&self.stream.task); + } +} + +fn wake_native_async_stream_task(stream: &NativeAsyncStream) { + let task = stream + .task + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref() + .and_then(|task| task.upgrade()); + if let Some(task) = task { + NativeAsyncTaskV2::wake(&task); } } @@ -1592,15 +2416,12 @@ async fn invoke_native_async_callback( invocation: Json, next: Option, ) -> FlowResult { - let runtime = if next.is_some() { - Some(tokio::runtime::Handle::try_current().map_err(|error| { - FlowError::Internal(format!( - "native async intercept requires a Tokio runtime: {error}" - )) - })?) - } else { - None - }; + let runtime = tokio::runtime::Handle::try_current().map_err(|error| { + FlowError::Internal(format!( + "native async middleware requires a Tokio runtime: {error}" + )) + })?; + let context = MiddlewareContinuationContext::capture(); let invocation = native_string_from_json(&invocation) .ok_or_else(|| FlowError::Internal("failed to allocate native async invocation".into()))? as usize; @@ -1610,6 +2431,9 @@ async fn invoke_native_async_callback( cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), next_abort: Mutex::new(None), + runtime: runtime.clone(), + context, + task: Mutex::new(None), #[cfg(test)] before_settlement_lock: None, _callback_user_data: Some(user_data.clone()), @@ -1620,66 +2444,73 @@ async fn invoke_native_async_callback( completed: false, }; let completion_ref = Arc::into_raw(completion.clone()) as usize; - let next_ref = match (next, runtime) { - (Some(inner), Some(runtime)) => Some(Arc::into_raw(Arc::new(NativeAsyncNext::new( + let next_ref = next.map(|inner| { + Arc::into_raw(Arc::new(NativeAsyncNext::new( inner, runtime, Some(user_data.clone()), - ))) as usize), - (None, None) => None, - _ => unreachable!("runtime is present exactly for native async intercepts"), + ))) as usize + }); + let callback_user_data = user_data.ptr as usize; + let callback_scope_stack = current_scope_stack(); + let invocation_guard = NativeInvocationStringGuard(invocation); + let completion_handoff = NativeCompletionHandoff { + raw: completion_ref, + armed: true, }; - let state = match catch_unwind(AssertUnwindSafe(|| unsafe { - cb( - user_data.ptr, - invocation as *const NemoRelayNativeString, - next_ref - .map(|next| next as *const NemoRelayNativeAsyncNext) - .unwrap_or(ptr::null()), - completion_ref as *const NemoRelayNativeAsyncCompletion, - ) - })) { - Ok(state) => state, - Err(_) => { - unsafe { - drop(Arc::from_raw( - completion_ref as *const NativeAsyncCompletion, - )); - native_string_free(invocation as *mut NemoRelayNativeString); - } - return Err(FlowError::Internal("native async callback panicked".into())); - } + let next_handoff = NativeNextHandoff { + raw: next_ref, + armed: true, }; - unsafe { native_string_free(invocation as *mut NemoRelayNativeString) }; - let state = match NemoRelayNativeAsyncCallbackState::try_from(state) { - Ok(state) => state, - Err(()) => { - unsafe { - drop(Arc::from_raw( - completion_ref as *const NativeAsyncCompletion, - )); - } - return Err(FlowError::Internal( - "native async callback returned an invalid state".into(), - )); + let invoke = move || { + let _invocation = invocation_guard; + let mut completion = completion_handoff; + let mut next = next_handoff; + let result = catch_unwind(AssertUnwindSafe(|| { + // Ownership transfers at callback entry. From this point the + // plugin must release `next`, even if its callback panics. + next.disarm(); + with_scope_stack(callback_scope_stack, || unsafe { + cb( + callback_user_data as *mut c_void, + invocation as *const NemoRelayNativeString, + next_ref + .map(|next| next as *const NemoRelayNativeAsyncNext) + .unwrap_or(ptr::null()), + completion_ref as *const NemoRelayNativeAsyncCompletion, + ) + }) + })); + if result + .as_ref() + .ok() + .and_then(|state| NemoRelayNativeAsyncCallbackState::try_from(*state).ok()) + == Some(NemoRelayNativeAsyncCallbackState::Pending) + { + // A pending plugin callback retains the completion and assumes + // responsibility for settling and releasing it. + completion.disarm(); } + result }; - if state == NemoRelayNativeAsyncCallbackState::Complete { - unsafe { - drop(Arc::from_raw( - completion_ref as *const NativeAsyncCompletion, - )); - } - if completion + let callback_result = invoke(); + let state = match callback_result { + Ok(state) => NemoRelayNativeAsyncCallbackState::try_from(state).map_err(|()| { + FlowError::Internal("native async callback returned an invalid state".into()) + }), + Err(_) => Err(FlowError::Internal("native async callback panicked".into())), + }?; + if state == NemoRelayNativeAsyncCallbackState::Complete + && completion .sender .lock() .unwrap_or_else(|error| error.into_inner()) .is_some() - { - return Err(FlowError::Internal( - "native async callback returned Complete without settling".into(), - )); - } + && !completion.cancelled.load(Ordering::Acquire) + { + return Err(FlowError::Internal( + "native async callback returned Complete without settling".into(), + )); } wait.receive().await } @@ -1719,6 +2550,8 @@ unsafe extern "C" fn native_async_completion_resolve_json( return NemoRelayStatus::InvalidArg; }; let _ = sender.send(Ok(value)); + drop(next_abort); + cancel_native_async_task_slot(&completion.task); NemoRelayStatus::Ok } @@ -1764,6 +2597,8 @@ unsafe extern "C" fn native_async_completion_reject( return NemoRelayStatus::InvalidArg; }; let _ = sender.send(Err(FlowError::Internal(message))); + drop(next_abort); + cancel_native_async_task_slot(&completion.task); NemoRelayStatus::Ok } @@ -1782,16 +2617,196 @@ unsafe extern "C" fn native_async_completion_release( } } -unsafe extern "C" fn native_async_next_release(next: *const NemoRelayNativeAsyncNext) { - if !next.is_null() { - unsafe { drop(Arc::from_raw(next as *const NativeAsyncNext)) }; +fn settle_native_async_completion_error(completion: &NativeAsyncCompletion, message: String) { + let mut next_abort = completion + .next_abort + .lock() + .unwrap_or_else(|error| error.into_inner()); + if completion.cancelled.load(Ordering::Acquire) { + return; } -} - -unsafe extern "C" fn native_async_stream_push_json( - stream: *const NemoRelayNativeAsyncStream, - chunk_json: *const NemoRelayNativeString, -) -> NemoRelayStatus { + if let Some(abort) = next_abort.take() { + abort.abort(); + } + let sender = completion + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if let Some(sender) = sender { + let _ = sender.send(Err(FlowError::Internal(message))); + } +} + +fn cancel_native_async_task_slot(slot: &Mutex>>) { + let task = { + let mut slot = slot.lock().unwrap_or_else(|error| error.into_inner()); + let task = slot.as_ref().and_then(Weak::upgrade); + if task.is_none() { + slot.take(); + } + task + }; + if let Some(task) = task { + NativeAsyncTaskV2::cancel(&task); + } +} + +unsafe fn weak_from_arc_raw(raw: *const T) -> Weak { + unsafe { Arc::increment_strong_count(raw) }; + let owner = unsafe { Arc::from_raw(raw) }; + Arc::downgrade(&owner) +} + +unsafe extern "C" fn native_async_completion_spawn_task_v2( + completion: *const NemoRelayNativeAsyncCompletion, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + clear_native_last_error(); + let Some(completion) = (unsafe { (completion as *const NativeAsyncCompletion).as_ref() }) + else { + return NemoRelayStatus::NullPointer; + }; + if completion.cancelled.load(Ordering::Acquire) + || completion + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + { + set_native_last_error("cannot spawn a task for a settled async completion"); + return NemoRelayStatus::InvalidArg; + } + spawn_native_async_task_v2( + completion.runtime.clone(), + completion.context.clone(), + NativeAsyncTaskOwnerV2::Completion(unsafe { + weak_from_arc_raw(completion as *const NativeAsyncCompletion) + }), + &completion.task, + completion._callback_user_data.clone(), + "completion", + (cb, user_data as usize, free_fn), + ) +} + +unsafe extern "C" fn native_async_stream_spawn_task_v2( + stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + clear_native_last_error(); + let Some(stream) = (unsafe { (stream as *const NativeAsyncStream).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + if stream.cancelled.load(Ordering::Acquire) || stream.settled.load(Ordering::Acquire) { + set_native_last_error("cannot spawn a task for a settled async stream"); + return NemoRelayStatus::InvalidArg; + } + spawn_native_async_task_v2( + stream.runtime.clone(), + stream.context.clone(), + NativeAsyncTaskOwnerV2::Stream(unsafe { + weak_from_arc_raw(stream as *const NativeAsyncStream) + }), + &stream.task, + stream._callback_user_data.clone(), + "stream", + (cb, user_data as usize, free_fn), + ) +} + +fn spawn_native_async_task_v2( + runtime: tokio::runtime::Handle, + context: MiddlewareContinuationContext, + owner: NativeAsyncTaskOwnerV2, + slot: &Mutex>>, + library_guard: Option>, + owner_label: &str, + callback: ( + NemoRelayNativeAsyncTaskPollCbV2, + usize, + NemoRelayNativeFreeFn, + ), +) -> NemoRelayStatus { + let mut slot = slot.lock().unwrap_or_else(|error| error.into_inner()); + if slot.as_ref().and_then(Weak::upgrade).is_some() { + set_native_last_error(format!("an async {owner_label} task is already active")); + return NemoRelayStatus::InvalidArg; + } + let (cb, user_data, free_fn) = callback; + let task = Arc::new(NativeAsyncTaskV2 { + context, + owner, + wake: tokio::sync::Notify::new(), + state: Mutex::new(NativeAsyncTaskStateV2 { + polling: false, + cancel_requested: false, + complete: false, + plugin: Some(NativeAsyncTaskPluginStateV2 { + cb, + user_data, + free_fn, + }), + }), + _library_guard: library_guard, + }); + *slot = Some(Arc::downgrade(&task)); + drop(slot); + runtime.spawn(Arc::clone(&task).run()); + NativeAsyncTaskV2::wake(&task); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn native_async_task_retain_v2(task: *const NemoRelayNativeAsyncTaskV2) { + if !task.is_null() { + unsafe { Arc::increment_strong_count(task as *const NativeAsyncTaskV2) }; + } +} + +unsafe extern "C" fn native_async_task_wake_v2(task: *const NemoRelayNativeAsyncTaskV2) { + if task.is_null() { + return; + } + unsafe { Arc::increment_strong_count(task as *const NativeAsyncTaskV2) }; + let task = unsafe { Arc::from_raw(task as *const NativeAsyncTaskV2) }; + NativeAsyncTaskV2::wake(&task); +} + +unsafe extern "C" fn native_async_task_release_v2(task: *const NemoRelayNativeAsyncTaskV2) { + if !task.is_null() { + unsafe { drop(Arc::from_raw(task as *const NativeAsyncTaskV2)) }; + } +} + +unsafe extern "C" fn native_async_next_release(next: *const NemoRelayNativeAsyncNext) { + if !next.is_null() { + unsafe { drop(Arc::from_raw(next as *const NativeAsyncNext)) }; + } +} + +unsafe extern "C" fn native_async_stream_push_json( + stream: *const NemoRelayNativeAsyncStream, + chunk_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { native_async_stream_push_json_impl(stream, chunk_json, NemoRelayStatus::Internal) } +} + +unsafe extern "C" fn native_async_stream_push_json_v2( + stream: *const NemoRelayNativeAsyncStream, + chunk_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { native_async_stream_push_json_impl(stream, chunk_json, NemoRelayStatus::WouldBlock) } +} + +unsafe fn native_async_stream_push_json_impl( + stream: *const NemoRelayNativeAsyncStream, + chunk_json: *const NemoRelayNativeString, + backpressure_status: NemoRelayStatus, +) -> NemoRelayStatus { clear_native_last_error(); let Some(stream) = (unsafe { (stream as *const NativeAsyncStream).as_ref() }) else { return NemoRelayStatus::NullPointer; @@ -1831,7 +2846,7 @@ unsafe extern "C" fn native_async_stream_push_json( set_native_last_error( "native async stream is backpressured; retry the chunk after the consumer advances", ); - NemoRelayStatus::Internal + backpressure_status } Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => NemoRelayStatus::InvalidArg, } @@ -1872,6 +2887,9 @@ unsafe extern "C" fn native_async_stream_finish( for (_, abort) in downstream_aborts.drain() { abort.abort(); } + drop(downstream_aborts); + drop(_settlement); + cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } else { NemoRelayStatus::InvalidArg @@ -1881,6 +2899,21 @@ unsafe extern "C" fn native_async_stream_finish( unsafe extern "C" fn native_async_stream_reject( stream: *const NemoRelayNativeAsyncStream, message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { native_async_stream_reject_impl(stream, message, NemoRelayStatus::Internal) } +} + +unsafe extern "C" fn native_async_stream_reject_v2( + stream: *const NemoRelayNativeAsyncStream, + message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { native_async_stream_reject_impl(stream, message, NemoRelayStatus::WouldBlock) } +} + +unsafe fn native_async_stream_reject_impl( + stream: *const NemoRelayNativeAsyncStream, + message: *const NemoRelayNativeString, + backpressure_status: NemoRelayStatus, ) -> NemoRelayStatus { clear_native_last_error(); let Some(stream) = (unsafe { (stream as *const NativeAsyncStream).as_ref() }) else { @@ -1920,13 +2953,17 @@ unsafe extern "C" fn native_async_stream_reject( for (_, abort) in downstream_aborts.drain() { abort.abort(); } + drop(downstream_aborts); + drop(sender_guard); + drop(_settlement); + cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { set_native_last_error( "native async stream is backpressured; retry rejection after the consumer advances", ); - NemoRelayStatus::Internal + backpressure_status } Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => NemoRelayStatus::InvalidArg, } @@ -2078,84 +3115,548 @@ unsafe extern "C" fn native_async_next_invoke_result( let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { return NemoRelayStatus::NullPointer; }; - let invocation = match parse_json_arg(invocation_json, "native async next invocation") { - Ok(value) => value, + let invocation = match parse_json_arg(invocation_json, "native async next invocation") { + Ok(value) => value, + Err(status) => return status, + }; + let future: Pin> + Send>> = match &next.inner { + NativeAsyncNextInner::Tool(next_fn) => { + let next_fn = next_fn.clone(); + Box::pin(async move { next_fn(invocation).await }) + } + NativeAsyncNextInner::Llm(next_fn) => { + let request = match serde_json::from_value(invocation) { + Ok(request) => request, + Err(error) => { + set_native_last_error(error.to_string()); + return NemoRelayStatus::InvalidJson; + } + }; + let next_fn = next_fn.clone(); + Box::pin(async move { next_fn(request).await }) + } + NativeAsyncNextInner::LlmStream(_) => { + set_native_last_error( + "stream continuations require async_next_invoke_stream; unary result callbacks cannot buffer a stream", + ); + return NemoRelayStatus::InvalidArg; + } + }; + let continuation_context = match next.context.isolated_for_current_invocation() { + Ok(context) => context, + Err(error) => return status_from_flow_error(error), + }; + let mut in_flight = next + .in_flight_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + let callbacks = Arc::clone(&next.in_flight_aborts); + let mut callback_guard = NativeAsyncResultCallbackGuard { + cb, + user_data: user_data as usize, + active: true, + _library_guard: next._callback_user_data.clone(), + }; + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let task = next.runtime.spawn(async move { + if start_rx.await.is_err() { + return; + } + let result = AssertUnwindSafe(continuation_context.run(future)) + .catch_unwind() + .await + .unwrap_or_else(|payload| { + Err(FlowError::Internal(format!( + "native async next continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ))) + }); + callback_guard.complete(result); + callbacks + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&tokio::task::id()); + }); + let abort = task.abort_handle(); + in_flight.insert(task.id(), abort); + let _ = start_tx.send(()); + NemoRelayStatus::Ok +} + +const NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES: usize = 4 * 1024; + +fn prepare_llm_continuation_invocation( + invocation: LlmContinuationInvocationV2, +) -> std::result::Result<(LlmRequest, LlmDispatchTargetContext), NemoRelayStatus> { + let target = + LlmDispatchTargetContext::try_new(invocation.target.url, invocation.target.headers) + .map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidArg + })?; + Ok((invocation.request, target)) +} + +fn non_http_llm_failure( + kind: LlmNonHttpFailureKindV2, + message: String, +) -> LlmContinuationFailureV2 { + LlmContinuationFailureV2::NonHttp { + kind, + message: bounded_utf8(message, NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES), + } +} + +fn typed_llm_failure(error: FlowError) -> LlmContinuationFailureV2 { + match error { + FlowError::Upstream(failure) => { + let failure = sanitize_upstream_failure(failure); + match failure.status { + Some(status) => LlmContinuationFailureV2::Http { + status, + body: failure.body, + headers: failure.headers, + }, + None => { + let kind = if failure.class == UpstreamFailureClass::Timeout { + LlmNonHttpFailureKindV2::Timeout + } else { + LlmNonHttpFailureKindV2::Transport + }; + non_http_llm_failure(kind, failure.body) + } + } + } + FlowError::GuardrailRejected(message) => { + non_http_llm_failure(LlmNonHttpFailureKindV2::Guardrail, message) + } + FlowError::InvalidArgument(message) => { + non_http_llm_failure(LlmNonHttpFailureKindV2::InvalidRequest, message) + } + other => non_http_llm_failure(LlmNonHttpFailureKindV2::Internal, other.to_string()), + } +} + +/// Invokes a unary LLM continuation through native API v2. +unsafe extern "C" fn native_async_llm_next_invoke_result_v2( + next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + let NativeAsyncNextInner::Llm(next_fn) = &next.inner else { + set_native_last_error("targeted LLM continuation requires an LLM execution continuation"); + return NemoRelayStatus::InvalidArg; + }; + let invocation = + match parse_json_arg(invocation_json, "targeted LLM continuation").and_then(|value| { + serde_json::from_value(value).map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidJson + }) + }) { + Ok(invocation) => invocation, + Err(status) => return status, + }; + let (request, target) = match prepare_llm_continuation_invocation(invocation) { + Ok(prepared) => prepared, + Err(status) => return status, + }; + let continuation_context = match next.context.isolated_for_current_invocation() { + Ok(context) => context, + Err(error) => return status_from_flow_error(error), + }; + let next_fn = next_fn.clone(); + let mut in_flight = next + .in_flight_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + let callbacks = Arc::clone(&next.in_flight_aborts); + let mut callback_guard = NativeLlmResultCallbackGuardV2 { + cb, + user_data: user_data as usize, + active: true, + _library_guard: next._callback_user_data.clone(), + }; + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let task = next.runtime.spawn(async move { + if start_rx.await.is_err() { + return; + } + let result = AssertUnwindSafe( + continuation_context.invoke_with_llm_dispatch_target(target, move || next_fn(request)), + ) + .catch_unwind() + .await + .unwrap_or_else(|payload| { + Err(FlowError::Internal(format!( + "typed native LLM continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ))) + }); + callback_guard.complete(result.map_err(typed_llm_failure)); + callbacks + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&tokio::task::id()); + }); + let abort = task.abort_handle(); + in_flight.insert(task.id(), abort); + let _ = start_tx.send(()); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn native_async_next_invoke_stream( + next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncNextStreamCb, + user_data: *mut c_void, +) -> NemoRelayStatus { + let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + if output_stream.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { Arc::increment_strong_count(output_stream as *const NativeAsyncStream) }; + let output_stream = unsafe { Arc::from_raw(output_stream as *const NativeAsyncStream) }; + let NativeAsyncNextInner::LlmStream(next_fn) = &next.inner else { + return NemoRelayStatus::InvalidArg; + }; + let request = match parse_json_arg(invocation_json, "native async stream next invocation") + .and_then(|value| { + serde_json::from_value(value).map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidJson + }) + }) { + Ok(request) => request, + Err(status) => return status, + }; + let continuation_context = match next.context.isolated_for_current_invocation() { + Ok(context) => context, + Err(error) => return status_from_flow_error(error), + }; + let _settlement = output_stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if output_stream.cancelled.load(Ordering::Acquire) + || output_stream.settled.load(Ordering::Acquire) + || output_stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + { + return NemoRelayStatus::InvalidArg; + } + let mut downstream_aborts = output_stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + let next_fn = next_fn.clone(); + let user_data = user_data as usize; + let output_stream_for_task = Arc::clone(&output_stream); + let output_stream_for_cleanup = Arc::clone(&output_stream); + let callback_guard = NativeAsyncStreamCallbackGuard { + cb, + user_data, + stream: output_stream_for_task, + _library_guard: next._callback_user_data.clone(), + active: true, + }; + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let task = next.runtime.spawn(async move { + if start_rx.await.is_err() { + return; + } + continuation_context + .run(deliver_native_async_next_stream( + next_fn, + request, + callback_guard, + )) + .await; + output_stream_for_cleanup + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&tokio::task::id()); + }); + let abort = task.abort_handle(); + downstream_aborts.insert(task.id(), abort); + let _ = start_tx.send(()); + NemoRelayStatus::Ok +} + +async fn deliver_native_async_next_stream( + next_fn: LlmStreamExecutionNextFn, + request: LlmRequest, + mut callback_guard: NativeAsyncStreamCallbackGuard, +) { + let result = AssertUnwindSafe(async { + match next_fn(request).await { + Ok(stream) => forward_native_async_next_stream(stream, &mut callback_guard).await, + Err(error) => callback_guard.fail(&error.to_string()), + } + }) + .catch_unwind() + .await; + if let Err(payload) = result { + callback_guard.fail(&format!( + "native async stream continuation panicked: {}", + panic_payload_message(payload.as_ref()) + )); + } +} + +async fn forward_native_async_next_stream( + stream: LlmJsonStream, + callback_guard: &mut NativeAsyncStreamCallbackGuard, +) { + forward_native_async_next_stream_with(stream, callback_guard, native_string_from_json).await; +} + +async fn forward_native_async_next_stream_with( + mut stream: LlmJsonStream, + callback_guard: &mut NativeAsyncStreamCallbackGuard, + to_native_string: impl Fn(&Json) -> Option<*mut NemoRelayNativeString>, +) { + while let Some(item) = stream.next().await { + match item { + Ok(chunk) => { + let Some(chunk) = to_native_string(&chunk) else { + callback_guard.fail( + "failed to serialize or allocate native async stream continuation chunk", + ); + return; + }; + let keep_going = unsafe { + (callback_guard.cb)( + callback_guard.user_data as *mut c_void, + chunk, + ptr::null(), + false, + ) + }; + unsafe { native_string_free(chunk) }; + if !keep_going { + callback_guard.finish(); + return; + } + } + Err(error) => { + callback_guard.fail(&error.to_string()); + return; + } + } + } + unsafe { + let _ = (callback_guard.cb)( + callback_guard.user_data as *mut c_void, + ptr::null(), + ptr::null(), + true, + ); + } + callback_guard.finish(); +} + +fn remove_current_native_stream_task(stream: &NativeAsyncStream) { + stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&tokio::task::id()); +} + +fn abort_native_stream_downstream_tasks(stream: &NativeAsyncStream) { + let mut downstream_aborts = stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + for (_, abort) in downstream_aborts.drain() { + abort.abort(); + } +} + +async fn push_forwarded_native_stream_chunk( + stream: &NativeAsyncStream, + chunk: Json, +) -> FlowResult<()> { + if stream.cancelled.load(Ordering::Acquire) || stream.settled.load(Ordering::Acquire) { + return Err(FlowError::Internal( + "native LLM pass-through output was cancelled".into(), + )); + } + let sender = stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + .ok_or_else(|| FlowError::Internal("native LLM pass-through output was settled".into()))?; + sender + .send(Ok(chunk)) + .await + .map_err(|_| FlowError::Internal("native LLM pass-through output was cancelled".into())) +} + +fn finish_forwarded_native_stream(stream: &NativeAsyncStream) { + let sender = { + let _settlement = stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if stream.cancelled.load(Ordering::Acquire) || stream.settled.load(Ordering::Acquire) { + return; + } + let sender = stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if sender.is_none() { + return; + } + stream.settled.store(true, Ordering::Release); + sender + }; + abort_native_stream_downstream_tasks(stream); + drop(sender); +} + +async fn reject_forwarded_native_stream(stream: &NativeAsyncStream, error: FlowError) { + let sender = { + let _settlement = stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if stream.cancelled.load(Ordering::Acquire) || stream.settled.load(Ordering::Acquire) { + return; + } + let sender = stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + let Some(sender) = sender else { + return; + }; + stream.settled.store(true, Ordering::Release); + sender + }; + abort_native_stream_downstream_tasks(stream); + let _ = sender.send(Err(error)).await; +} + +/// Forwards the ordinary downstream LLM continuation into the host-owned +/// output stream without crossing the plugin boundary for each event. +unsafe extern "C" fn native_async_llm_next_forward_stream_v2( + next: *const NemoRelayNativeAsyncNext, + request_json: *const NemoRelayNativeString, + output_stream: *const NemoRelayNativeAsyncStream, + terminal_callback: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + if output_stream.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { Arc::increment_strong_count(output_stream as *const NativeAsyncStream) }; + let output_stream = unsafe { Arc::from_raw(output_stream as *const NativeAsyncStream) }; + let NativeAsyncNextInner::LlmStream(next_fn) = &next.inner else { + set_native_last_error( + "native LLM stream pass-through requires an LLM stream execution continuation", + ); + return NemoRelayStatus::InvalidArg; + }; + let request = match parse_llm_request_arg(request_json, "native LLM stream pass-through") { + Ok(request) => request, Err(status) => return status, }; - let future: Pin> + Send>> = match &next.inner { - NativeAsyncNextInner::Tool(next_fn) => { - let next_fn = next_fn.clone(); - Box::pin(async move { next_fn(invocation).await }) - } - NativeAsyncNextInner::Llm(next_fn) => { - let request = match serde_json::from_value(invocation) { - Ok(request) => request, - Err(error) => { - set_native_last_error(error.to_string()); - return NemoRelayStatus::InvalidJson; - } - }; - let next_fn = next_fn.clone(); - Box::pin(async move { next_fn(request).await }) - } - NativeAsyncNextInner::LlmStream(_) => { - set_native_last_error( - "stream continuations require async_next_invoke_stream; unary result callbacks cannot buffer a stream", - ); - return NemoRelayStatus::InvalidArg; - } - }; let continuation_context = match next.context.isolated_for_current_invocation() { Ok(context) => context, Err(error) => return status_from_flow_error(error), }; - let user_data = user_data as usize; - let _library_guard = next._callback_user_data.clone(); - next.runtime.spawn(async move { - let result = AssertUnwindSafe(continuation_context.run(future)) - .catch_unwind() - .await - .unwrap_or_else(|payload| { - Err(FlowError::Internal(format!( - "native async next continuation panicked: {}", - panic_payload_message(payload.as_ref()) - ))) - }); - match result { - Ok(value) => { - if let Some(value) = native_string_from_json(&value) { - unsafe { - cb(user_data as *mut c_void, value, ptr::null()); - native_string_free(value); - } - } else if let Some(error) = - native_string_from_str("failed to allocate native async next result") - { - unsafe { - cb(user_data as *mut c_void, ptr::null(), error); - native_string_free(error); - } - } - } - Err(error) => { - if let Some(error) = native_string_from_str(&error.to_string()) { - unsafe { - cb(user_data as *mut c_void, ptr::null(), error); - native_string_free(error); + let settlement = output_stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if output_stream.cancelled.load(Ordering::Acquire) + || output_stream.settled.load(Ordering::Acquire) + || output_stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + { + return NemoRelayStatus::InvalidArg; + } + let mut downstream_aborts = output_stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + let next_fn = next_fn.clone(); + let stream_for_pump = Arc::clone(&output_stream); + let callback_guard = NativeLlmStreamForwardCallbackGuardV2 { + cb: terminal_callback, + user_data: user_data as usize, + stream: Arc::clone(&output_stream), + active: true, + }; + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let task = next.runtime.spawn(async move { + if start_rx.await.is_err() { + return; + } + let mut callback_guard = callback_guard; + let result = continuation_context + .run(async move { + AssertUnwindSafe(async move { + let mut downstream = next_fn(request).await?; + while let Some(item) = downstream.next().await { + push_forwarded_native_stream_chunk(&stream_for_pump, item?).await?; } - } - } + Ok(()) + }) + .catch_unwind() + .await + .unwrap_or_else(|payload| { + Err(FlowError::Internal(format!( + "native LLM stream pass-through panicked: {}", + panic_payload_message(payload.as_ref()) + ))) + }) + }) + .await; + remove_current_native_stream_task(&callback_guard.stream); + match result { + Ok(()) => finish_forwarded_native_stream(&callback_guard.stream), + Err(error) => reject_forwarded_native_stream(&callback_guard.stream, error).await, } - drop(_library_guard); + callback_guard.settle(); }); + let abort = task.abort_handle(); + downstream_aborts.insert(task.id(), abort); + drop(downstream_aborts); + drop(settlement); + let _ = start_tx.send(()); NemoRelayStatus::Ok } -unsafe extern "C" fn native_async_next_invoke_stream( +/// Opens a streaming LLM continuation through native API v2. +unsafe extern "C" fn native_async_llm_next_open_stream_v2( next: *const NemoRelayNativeAsyncNext, invocation_json: *const NemoRelayNativeString, output_stream: *const NemoRelayNativeAsyncStream, - cb: NemoRelayNativeAsyncNextStreamCb, + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, user_data: *mut c_void, ) -> NemoRelayStatus { let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { @@ -2167,23 +3668,30 @@ unsafe extern "C" fn native_async_next_invoke_stream( unsafe { Arc::increment_strong_count(output_stream as *const NativeAsyncStream) }; let output_stream = unsafe { Arc::from_raw(output_stream as *const NativeAsyncStream) }; let NativeAsyncNextInner::LlmStream(next_fn) = &next.inner else { + set_native_last_error( + "targeted LLM continuation requires an LLM stream execution continuation", + ); return NemoRelayStatus::InvalidArg; }; - let request = match parse_json_arg(invocation_json, "native async stream next invocation") + let invocation = match parse_json_arg(invocation_json, "targeted streaming LLM continuation") .and_then(|value| { serde_json::from_value(value).map_err(|error| { set_native_last_error(error.to_string()); NemoRelayStatus::InvalidJson }) }) { - Ok(request) => request, + Ok(invocation) => invocation, + Err(status) => return status, + }; + let (request, target) = match prepare_llm_continuation_invocation(invocation) { + Ok(prepared) => prepared, Err(status) => return status, }; let continuation_context = match next.context.isolated_for_current_invocation() { Ok(context) => context, Err(error) => return status_from_flow_error(error), }; - let _settlement = output_stream + let settlement = output_stream .settlement .lock() .unwrap_or_else(|error| error.into_inner()); @@ -2202,13 +3710,14 @@ unsafe extern "C" fn native_async_next_invoke_stream( .lock() .unwrap_or_else(|error| error.into_inner()); let next_fn = next_fn.clone(); - let user_data = user_data as usize; + let provider_runtime = next.runtime.clone(); + let provider_context = continuation_context.clone(); + let provider_target = target.clone(); let output_stream_for_task = Arc::clone(&output_stream); let output_stream_for_cleanup = Arc::clone(&output_stream); - let callback_guard = NativeAsyncStreamCallbackGuard { + let callback_guard = NativeLlmStreamOpenCallbackGuardV2 { cb, - user_data, - stream: output_stream_for_task, + user_data: user_data as usize, _library_guard: next._callback_user_data.clone(), active: true, }; @@ -2217,13 +3726,46 @@ unsafe extern "C" fn native_async_next_invoke_stream( if start_rx.await.is_err() { return; } - continuation_context - .run(deliver_native_async_next_stream( - next_fn, - request, - callback_guard, - )) + let mut callback_guard = callback_guard; + let result = continuation_context + .invoke_with_llm_dispatch_target(target, move || { + AssertUnwindSafe(next_fn(request)).catch_unwind() + }) .await; + match result { + Ok(Ok(provider_stream)) => { + let provider = Arc::new(NativeLlmProviderStreamV2 { + stream: tokio::sync::Mutex::new(Some(provider_stream)), + runtime: provider_runtime, + context: provider_context, + target: provider_target, + output: Arc::clone(&output_stream_for_task), + lifecycle: Mutex::new(NativeLlmProviderStreamLifecycleV2::Idle), + }); + let _settlement = output_stream_for_task + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if output_stream_for_task.cancelled.load(Ordering::Acquire) + || output_stream_for_task.settled.load(Ordering::Acquire) + { + callback_guard.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM stream output settled during setup".into(), + )); + } else { + callback_guard.success(provider); + } + } + Ok(Err(error)) => callback_guard.failure(&typed_llm_failure(error)), + Err(payload) => callback_guard.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!( + "typed native LLM stream continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ), + )), + } output_stream_for_cleanup .downstream_aborts .lock() @@ -2232,81 +3774,126 @@ unsafe extern "C" fn native_async_next_invoke_stream( }); let abort = task.abort_handle(); downstream_aborts.insert(task.id(), abort); + drop(settlement); let _ = start_tx.send(()); NemoRelayStatus::Ok } -async fn deliver_native_async_next_stream( - next_fn: LlmStreamExecutionNextFn, - request: LlmRequest, - mut callback_guard: NativeAsyncStreamCallbackGuard, -) { - let result = AssertUnwindSafe(async { - match next_fn(request).await { - Ok(stream) => forward_native_async_next_stream(stream, &mut callback_guard).await, - Err(error) => callback_guard.fail(&error.to_string()), +/// Requests one event from a native API v2 provider stream. +unsafe extern "C" fn native_async_llm_stream_next_v2( + stream: *const NemoRelayNativeLlmStreamV2, + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let Some(stream) = (unsafe { (stream as *const NativeLlmProviderStreamV2).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + let settlement = stream + .output + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if stream.output.cancelled.load(Ordering::Acquire) + || stream.output.settled.load(Ordering::Acquire) + { + set_native_last_error("native API v2 provider stream output is settled"); + return NemoRelayStatus::InvalidArg; + } + let mut lifecycle = stream + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + match &*lifecycle { + NativeLlmProviderStreamLifecycleV2::Idle => {} + NativeLlmProviderStreamLifecycleV2::Pulling { .. } => { + set_native_last_error( + "native API v2 provider stream already has a pending next operation", + ); + return NemoRelayStatus::InvalidArg; + } + NativeLlmProviderStreamLifecycleV2::Terminal => { + set_native_last_error("native API v2 provider stream is already terminal"); + return NemoRelayStatus::InvalidArg; + } + NativeLlmProviderStreamLifecycleV2::Cancelled => { + set_native_last_error("native API v2 provider stream is cancelled"); + return NemoRelayStatus::InvalidArg; } - }) - .catch_unwind() - .await; - if let Err(payload) = result { - callback_guard.fail(&format!( - "native async stream continuation panicked: {}", - panic_payload_message(payload.as_ref()) - )); } -} - -async fn forward_native_async_next_stream( - stream: LlmJsonStream, - callback_guard: &mut NativeAsyncStreamCallbackGuard, -) { - forward_native_async_next_stream_with(stream, callback_guard, native_string_from_json).await; -} - -async fn forward_native_async_next_stream_with( - mut stream: LlmJsonStream, - callback_guard: &mut NativeAsyncStreamCallbackGuard, - to_native_string: impl Fn(&Json) -> Option<*mut NemoRelayNativeString>, -) { - while let Some(item) = stream.next().await { - match item { - Ok(chunk) => { - let Some(chunk) = to_native_string(&chunk) else { - callback_guard.fail( - "failed to serialize or allocate native async stream continuation chunk", - ); - return; - }; - let keep_going = unsafe { - (callback_guard.cb)( - callback_guard.user_data as *mut c_void, - chunk, - ptr::null(), - false, - ) - }; - unsafe { native_string_free(chunk) }; - if !keep_going { - callback_guard.finish(); - return; + unsafe { Arc::increment_strong_count(stream as *const NativeLlmProviderStreamV2) }; + let provider = unsafe { Arc::from_raw(stream as *const NativeLlmProviderStreamV2) }; + let callback_guard = NativeLlmProviderNextCallbackGuardV2 { + cb, + user_data: user_data as usize, + provider: Arc::clone(&provider), + active: true, + }; + let provider_for_task = Arc::clone(&provider); + let context = provider.context.clone(); + let target = provider.target.clone(); + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let task = provider.runtime.spawn(async move { + let mut callback_guard = callback_guard; + if start_rx.await.is_err() { + return; + } + let result = AssertUnwindSafe(context.invoke_with_llm_dispatch_target(target, || async { + let mut stream = provider_for_task.stream.lock().await; + let Some(provider_stream) = stream.as_mut() else { + return Ok(None); + }; + match provider_stream.next().await { + Some(Ok(chunk)) => Ok(Some(chunk)), + Some(Err(error)) => { + stream.take(); + Err(typed_llm_failure(error)) + } + None => { + stream.take(); + Ok(None) } } - Err(error) => { - callback_guard.fail(&error.to_string()); - return; + })) + .catch_unwind() + .await; + let outcome = match result { + Ok(outcome) => outcome, + Err(payload) => { + provider_for_task.stream.lock().await.take(); + Err(non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!( + "typed native LLM provider stream panicked: {}", + panic_payload_message(payload.as_ref()) + ), + )) } - } - } - unsafe { - let _ = (callback_guard.cb)( - callback_guard.user_data as *mut c_void, - ptr::null(), - ptr::null(), - true, - ); + }; + callback_guard.complete(outcome); + }); + let task_id_value = task.id(); + let abort = task.abort_handle(); + *lifecycle = NativeLlmProviderStreamLifecycleV2::Pulling { + task_id: task_id_value, + }; + provider + .output + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert(task_id_value, abort); + drop(lifecycle); + drop(settlement); + let _ = start_tx.send(()); + NemoRelayStatus::Ok +} + +/// Releases a native API v2 provider-stream reference. +unsafe extern "C" fn native_async_llm_stream_release_v2(stream: *const NemoRelayNativeLlmStreamV2) { + if !stream.is_null() { + let stream = unsafe { Arc::from_raw(stream as *const NativeLlmProviderStreamV2) }; + stream.cancel(); } - callback_guard.finish(); } fn wrap_native_async_tool_json( @@ -2548,6 +4135,13 @@ fn wrap_native_async_llm_execution( free_fn: NemoRelayNativeFreeFn, ) -> LlmExecutionFn { let user_data = make_user_data(instance, user_data, free_fn); + wrap_native_async_llm_execution_with_user_data(cb, user_data) +} + +fn wrap_native_async_llm_execution_with_user_data( + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: Arc, +) -> LlmExecutionFn { Arc::new(move |name, request, next| { let user_data = user_data.clone(); let name = name.to_owned(); @@ -2570,25 +4164,50 @@ fn wrap_native_incremental_llm_stream_execution( free_fn: NemoRelayNativeFreeFn, ) -> LlmStreamExecutionFn { let user_data = make_user_data(instance, user_data, free_fn); - wrap_native_incremental_llm_stream_execution_with_user_data(cb, user_data) + wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( + cb, + user_data, + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY, + ) } +#[cfg(test)] fn wrap_native_incremental_llm_stream_execution_with_user_data( cb: NemoRelayNativeAsyncStreamMiddlewareCb, user_data: Arc, +) -> LlmStreamExecutionFn { + wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( + cb, + user_data, + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY, + ) +} + +fn wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: Arc, + channel_capacity: usize, ) -> LlmStreamExecutionFn { Arc::new(move |name, request, next| { let user_data = user_data.clone(); let name = name.to_owned(); Box::pin(async move { - let (sender, receiver) = - tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY); + let runtime = tokio::runtime::Handle::try_current().map_err(|error| { + FlowError::Internal(format!( + "native async stream intercept requires a Tokio runtime: {error}" + )) + })?; + let context = MiddlewareContinuationContext::capture(); + let (sender, receiver) = tokio::sync::mpsc::channel(channel_capacity); let stream = Arc::new(NativeAsyncStream { sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: runtime.clone(), + context, + task: Mutex::new(None), #[cfg(test)] before_settlement_lock: None, _callback_user_data: Some(user_data.clone()), @@ -2605,11 +4224,6 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data( "failed to allocate native async stream invocation".into(), ) })?; - let runtime = tokio::runtime::Handle::try_current().map_err(|error| { - FlowError::Internal(format!( - "native async stream intercept requires a Tokio runtime: {error}" - )) - })?; let next_ref = Arc::into_raw(Arc::new(NativeAsyncNext::new( NativeAsyncNextInner::LlmStream(next), runtime, @@ -2651,6 +4265,36 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data( }) } +async fn settle_native_async_stream_error(stream: Arc, message: String) { + let sender = { + let _settlement = stream + .settlement + .lock() + .unwrap_or_else(|error| error.into_inner()); + if stream.cancelled.load(Ordering::Acquire) || stream.settled.swap(true, Ordering::AcqRel) { + return; + } + stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + }; + if let Some(sender) = sender { + let _ = sender.send(Err(FlowError::Internal(message))).await; + } + let mut downstream_aborts = stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()); + for (_, abort) in downstream_aborts.drain() { + abort.abort(); + } + drop(downstream_aborts); + wake_native_async_stream_task(&stream); + cancel_native_async_task_slot(&stream.task); +} + unsafe extern "C" fn native_plugin_context_register_async_stream_middleware( ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, diff --git a/crates/core/tests/coverage/error_tests.rs b/crates/core/tests/coverage/error_tests.rs index 55e823a72..5909473b5 100644 --- a/crates/core/tests/coverage/error_tests.rs +++ b/crates/core/tests/coverage/error_tests.rs @@ -128,3 +128,36 @@ fn upstream_failures_classify_retryability_and_render_status() { assert!(!rejected.is_retryable()); assert!(rejected.to_string().contains("transport failure")); } + +#[test] +fn upstream_failure_sanitization_bounds_data_and_keeps_only_safe_headers() { + use std::collections::BTreeMap; + + let failure = sanitize_upstream_failure(UpstreamFailure { + status: Some(429), + body: "é".repeat(MAX_UPSTREAM_FAILURE_BODY_BYTES), + headers: BTreeMap::from([ + ("Retry-After".into(), "2".into()), + ("Authorization".into(), "Bearer secret".into()), + ( + "X-Request-Id".into(), + format!( + "{}é", + "x".repeat(MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) + ), + ), + ]), + class: UpstreamFailureClass::RetryableStatus, + }); + + assert_eq!(failure.body.len(), MAX_UPSTREAM_FAILURE_BODY_BYTES); + assert_eq!( + failure.headers.get("retry-after").map(String::as_str), + Some("2") + ); + assert_eq!( + failure.headers.get("x-request-id").map(String::len), + Some(MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) + ); + assert!(!failure.headers.contains_key("authorization")); +} diff --git a/crates/core/tests/fixtures/native_plugin/Cargo.toml b/crates/core/tests/fixtures/native_plugin/Cargo.toml index 5b8356fcd..66a90de40 100644 --- a/crates/core/tests/fixtures/native_plugin/Cargo.toml +++ b/crates/core/tests/fixtures/native_plugin/Cargo.toml @@ -14,5 +14,6 @@ publish = false crate-type = ["cdylib"] [dependencies] +futures = "0.3" nemo-relay-plugin = { path = "../../../../plugin" } serde_json = "1" diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index a8c47a1bc..b44d7aa72 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -3,17 +3,20 @@ use std::ffi::c_void; use std::ptr; +use std::sync::Mutex; use std::sync::atomic::{AtomicBool, Ordering}; +use futures::StreamExt; use nemo_relay_plugin::{ CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, - Json, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareCb, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, - NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, - NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, - PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, - ToolExecutionInterceptOutcome, + Json, LlmContinuationInvocationV2, LlmContinuationTargetV2, LlmContinuationV2, + LlmJsonAsyncStreamV2, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, + LlmStreamExecutionOutcomeV2, NativePlugin, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, + NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, + ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, NEMO_RELAY_NATIVE_ABI_VERSION, }; use serde_json::{Map, json}; @@ -297,6 +300,176 @@ fn mark_json(mut value: Json, key: &str) -> Json { } nemo_relay_plugin::nemo_relay_plugin!(nemo_relay_fixture_native_plugin, || FixtureNativePlugin); +nemo_relay_plugin::nemo_relay_plugin_v2!(nemo_relay_fixture_native_api_v2_plugin, || { + FixtureNativePlugin +}); + +struct TargetedFixturePlugin; + +static ESCAPED_V2_CONTINUATION: Mutex> = Mutex::new(None); +static TARGETED_V2_PLUGIN_DROPPED: AtomicBool = AtomicBool::new(false); + +impl Drop for TargetedFixturePlugin { + fn drop(&mut self) { + TARGETED_V2_PLUGIN_DROPPED.store(true, Ordering::Release); + } +} + +#[unsafe(no_mangle)] +pub extern "C" fn nemo_relay_fixture_release_escaped_v2_continuation() -> bool { + ESCAPED_V2_CONTINUATION + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .is_some() +} + +#[unsafe(no_mangle)] +pub extern "C" fn nemo_relay_fixture_v2_library_probe() -> u64 { + 0x4e52_5632_u64 +} + +#[unsafe(no_mangle)] +pub extern "C" fn nemo_relay_fixture_targeted_v2_plugin_dropped() -> bool { + TARGETED_V2_PLUGIN_DROPPED.load(Ordering::Acquire) +} + +impl NativePlugin for TargetedFixturePlugin { + fn plugin_kind(&self) -> &str { + "fixture_native" + } + + fn register( + &mut self, + plugin_config: &Map, + ctx: &mut PluginContext<'_>, + ) -> nemo_relay_plugin::Result<()> { + let runtime = ctx.runtime(); + let url = plugin_config + .get("target_url") + .and_then(Json::as_str) + .ok_or_else(|| "targeted fixture requires target_url".to_string())? + .to_owned(); + let buffered_url = url.clone(); + let buffered_runtime = runtime.clone(); + let escape_continuation = plugin_config + .get("escape_continuation") + .and_then(Json::as_bool) + .unwrap_or(false); + ctx.register_async_llm_execution_v2( + "fixture_targeted_llm", + 0, + move |name, mut request, continuation| { + let url = buffered_url.clone(); + let runtime = buffered_runtime.clone(); + async move { + if escape_continuation { + *ESCAPED_V2_CONTINUATION + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some(continuation.clone()); + } + if name == "fixture_passthrough_llm" { + return continuation.call_passthrough(request).await; + } + cooperative_yield_once().await; + runtime.emit_mark( + "fixture.native.v2.after_await", + Some(&json!({"phase": "resumed"})), + None, + )?; + request.content["fixture_targeted"] = json!(true); + continuation + .call(targeted_fixture_invocation(url, request)) + .await + .map_err(|error| format!("targeted fixture dispatch failed: {error:?}")) + } + }, + )?; + + ctx.register_async_llm_stream_execution_v2( + "fixture_targeted_llm_stream", + 0, + move |name, mut request, continuation| { + let url = url.clone(); + async move { + if name == "fixture_passthrough_llm_stream" { + return Ok(LlmStreamExecutionOutcomeV2::Passthrough(request)); + } + request.content["fixture_targeted_stream"] = json!(true); + let stream = continuation + .open_stream(targeted_fixture_invocation(url, request)) + .await + .map_err(|error| { + format!("targeted fixture stream dispatch failed: {error:?}") + })?; + let stream: LlmJsonAsyncStreamV2 = Box::pin(stream.map(|item| { + item.map_err(|error| { + format!("targeted fixture provider stream failed: {error:?}") + }) + })); + Ok(LlmStreamExecutionOutcomeV2::Stream(stream)) + } + }, + ) + } +} + +async fn cooperative_yield_once() { + let mut yielded = false; + futures::future::poll_fn(move |cx| { + if yielded { + std::task::Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; +} + +fn targeted_fixture_invocation(url: String, request: LlmRequest) -> LlmContinuationInvocationV2 { + LlmContinuationInvocationV2 { + request, + target: LlmContinuationTargetV2 { + url, + headers: std::collections::BTreeMap::from([( + "authorization".into(), + "Bearer fixture-target".into(), + )]), + }, + } +} + +nemo_relay_plugin::nemo_relay_plugin_v2!(nemo_relay_fixture_targeted_v2_plugin, || { + TargetedFixturePlugin +}); + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_fixture_native_api_v1_plugin( + host: *const NemoRelayNativeHostApiV1, + out: *mut NemoRelayNativePluginV1, +) -> NemoRelayStatus { + let Some(host_ref) = (unsafe { host.as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + if host_ref.abi_version != NEMO_RELAY_NATIVE_ABI_VERSION + || host_ref.struct_size != std::mem::size_of::() + { + return NemoRelayStatus::InvalidArg; + } + unsafe { + write_raw_descriptor( + host, + out, + "fixture_native", + None, + None, + Some(raw_noop_register), + ) + } +} #[unsafe(no_mangle)] pub unsafe extern "C" fn nemo_relay_fixture_async_entry( diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 8d572f247..230b8843b 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -5,7 +5,7 @@ use std::path::{Path, PathBuf}; use std::process::Command; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use nemo_relay::api::event::{Event, ScopeCategory}; @@ -42,6 +42,9 @@ use uuid::Uuid; static NATIVE_PLUGIN_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); const PLUGIN_DISCOVERY_TEST_CHILD: &str = "NEMO_RELAY_PLUGIN_DISCOVERY_TEST_CHILD"; +const NATIVE_V2_UNLOAD_TEST_CHILD: &str = "NEMO_RELAY_NATIVE_V2_UNLOAD_TEST_CHILD"; +const NATIVE_V2_UNLOAD_TEST_MANIFEST: &str = "NEMO_RELAY_NATIVE_V2_UNLOAD_TEST_MANIFEST"; +const NATIVE_V2_UNLOAD_TEST_LIBRARY: &str = "NEMO_RELAY_NATIVE_V2_UNLOAD_TEST_LIBRARY"; struct ReplacementRegistryPlugin; @@ -1198,7 +1201,7 @@ fn native_loader_rejects_manifest_contract_errors_before_loading_library() { &native_manifest_text( "fixture_native", &format!("={}", env!("CARGO_PKG_VERSION")), - "2", + "3", "libdoes-not-need-to-exist.so", "nemo_relay_fixture_native_plugin", ), @@ -1246,6 +1249,553 @@ entrypoint = "fixture.worker:create_plugin" assert!(error.contains("only supports rust_dynamic"), "{error}"); } +#[test] +fn native_api_v1_plugin_loads_unchanged_beside_native_api_v2() { + let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); + let fixture = build_fixture_plugin(); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "1", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_native_api_v1_plugin", + ), + ); + + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("native API v1 plugin should load against the preserved host table"); + activation.clear(); +} + +#[test] +fn native_api_v2_plugin_requires_the_v2_manifest_contract() { + let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); + let fixture = build_fixture_plugin(); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "2", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_native_api_v2_plugin", + ), + ); + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("native API v2 plugin should receive the typed host table"); + activation.clear(); + + let v1_manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "1", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_native_api_v2_plugin", + ), + ); + let error = expect_native_load_error( + load_spec("fixture_native", &v1_manifest_ref), + "native API v2-only plugin must reject native API v1 negotiation", + ); + assert!(error.contains("entry symbol"), "{error}"); +} + +#[test] +fn native_api_v2_fixture_dispatches_cooperatively_without_a_cli_gateway() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .max_blocking_threads(1) + .enable_all() + .build() + .expect("fixture runtime should build"); + runtime.block_on(async { + let _guard = NATIVE_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_plugin(); + let provider = EmbeddedFakeProvider::spawn( + br#"{"id":"embedded-target-response"}"#, + "application/json", + ); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "2", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_targeted_v2_plugin", + ), + ); + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("targeted native API v2 fixture should load"); + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_native".into(), + enabled: true, + config: Map::from_iter([("target_url".into(), json!(provider.url))]), + }); + initialize_plugins_exact(plugin_config) + .await + .expect("targeted native API v2 fixture should initialize"); + let mut cleanup = NativePluginTestCleanup::new(); + cleanup.mark_plugin_configuration_active(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured_events = events.clone(); + register_subscriber( + "native_plugin_v2_cooperative_events", + Arc::new(move |event| captured_events.lock().unwrap().push(event.clone())), + ) + .expect("cooperative v2 event subscriber should register"); + cleanup.mark_subscriber_registered("native_plugin_v2_cooperative_events"); + + let (blocking_entered_tx, blocking_entered_rx) = tokio::sync::oneshot::channel(); + let (blocking_release_tx, blocking_release_rx) = std::sync::mpsc::channel(); + let blocking_worker = tokio::task::spawn_blocking(move || { + let _ = blocking_entered_tx.send(()); + let _ = blocking_release_rx.recv(); + }); + blocking_entered_rx + .await + .expect("the only blocking worker should be occupied"); + + let original_provider_called = Arc::new(AtomicBool::new(false)); + let original_provider_called_for_fn = original_provider_called.clone(); + let response = tokio::time::timeout( + std::time::Duration::from_secs(3), + llm_call_execute( + LlmCallExecuteParams::builder() + .name("embedded-targeted-native-v2") + .request(LlmRequest { + headers: Map::new(), + content: json!({"model": "caller-model", "prompt": "hello"}), + }) + .func(Arc::new(move |_| { + original_provider_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(json!({"id": "wrong-provider"})) }) + })) + .build(), + ), + ) + .await; + let _ = blocking_release_tx.send(()); + blocking_worker + .await + .expect("occupied blocking worker should exit cleanly"); + let response = response + .expect("safe v2 callback must not wait for the occupied blocking worker") + .expect("embedded targeted dispatch should succeed"); + + assert_eq!(response, json!({"id": "embedded-target-response"})); + assert!(!original_provider_called.load(Ordering::SeqCst)); + let captured = + String::from_utf8(provider.request()).expect("captured request should be UTF-8"); + assert!(captured.contains("authorization: Bearer fixture-target\r\n")); + assert!(captured.contains("\"fixture_targeted\":true")); + + flush_subscribers().expect("cooperative v2 events should flush"); + let events = events.lock().unwrap(); + let llm_start = find_event( + &events, + "embedded-targeted-native-v2", + Some(ScopeCategory::Start), + ); + let resumed_mark = find_event(&events, "fixture.native.v2.after_await", None); + assert_eq!( + resumed_mark.parent_uuid(), + llm_start.parent_uuid(), + "a resumed safe callback must retain the LLM invocation's active scope" + ); + assert_eq!(resumed_mark.data().unwrap()["phase"], "resumed"); + drop(events); + + drop(cleanup); + activation.clear(); + }); +} + +#[test] +fn native_api_v2_escaped_continuation_release_keeps_library_mapped() { + if std::env::var_os(NATIVE_V2_UNLOAD_TEST_CHILD).is_none() { + let fixture = build_fixture_plugin(); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "2", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_targeted_v2_plugin", + ), + ); + let output = Command::new(std::env::current_exe().expect("test executable should resolve")) + .args([ + "--exact", + "native_api_v2_escaped_continuation_release_keeps_library_mapped", + "--nocapture", + ]) + .env(NATIVE_V2_UNLOAD_TEST_CHILD, "1") + .env(NATIVE_V2_UNLOAD_TEST_MANIFEST, &manifest_ref) + .env(NATIVE_V2_UNLOAD_TEST_LIBRARY, &fixture.library_path) + .output() + .expect("native unload-safety child process should run"); + assert!( + output.status.success(), + "native unload-safety child process failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("fixture runtime should build"); + runtime.block_on(async { + let _guard = NATIVE_PLUGIN_TEST_LOCK.lock().await; + let manifest_ref = std::env::var(NATIVE_V2_UNLOAD_TEST_MANIFEST) + .expect("child manifest path should be provided"); + let library_path = std::env::var(NATIVE_V2_UNLOAD_TEST_LIBRARY) + .expect("child library path should be provided"); + let activation = load_native_plugins([NativePluginLoadSpec { + plugin_id: "fixture_native".into(), + manifest_ref, + }]) + .expect("targeted native API v2 fixture should load"); + + type ReleaseContinuation = unsafe extern "C" fn() -> bool; + type LibraryProbe = unsafe extern "C" fn() -> u64; + type PluginDropped = unsafe extern "C" fn() -> bool; + let fixture_library = unsafe { libloading::Library::new(library_path) } + .expect("targeted fixture should open for lifecycle probes"); + let release_continuation = unsafe { + *fixture_library + .get::(b"nemo_relay_fixture_release_escaped_v2_continuation\0") + .expect("fixture should export its continuation release probe") + }; + let library_probe = unsafe { + *fixture_library + .get::(b"nemo_relay_fixture_v2_library_probe\0") + .expect("fixture should export its library probe") + }; + let plugin_dropped = unsafe { + *fixture_library + .get::(b"nemo_relay_fixture_targeted_v2_plugin_dropped\0") + .expect("fixture should export its descriptor-drop probe") + }; + assert!(!unsafe { release_continuation() }); + assert!(!unsafe { plugin_dropped() }); + // Do not retain an independent libloading handle across activation + // teardown; that would mask the self-unload bug under test. + drop(fixture_library); + + let provider = + EmbeddedFakeProvider::spawn(br#"{"id":"escaped-handle-response"}"#, "application/json"); + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_native".into(), + enabled: true, + config: Map::from_iter([ + ("target_url".into(), json!(provider.url)), + ("escape_continuation".into(), json!(true)), + ]), + }); + initialize_plugins_exact(plugin_config) + .await + .expect("targeted native API v2 fixture should initialize"); + + let response = llm_call_execute( + LlmCallExecuteParams::builder() + .name("escaped-v2-continuation") + .request(LlmRequest { + headers: Map::new(), + content: json!({"model": "caller-model", "prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { Ok(json!({"id": "wrong-provider"})) }) + })) + .build(), + ) + .await + .expect("targeted dispatch should succeed before teardown"); + assert_eq!(response, json!({"id": "escaped-handle-response"})); + let _ = provider.request(); + + clear_plugin_configuration().expect("plugin callbacks should clear before activation"); + activation.clear(); + + // Dropping the escaped safe wrapper calls back into Relay. Before the + // unload-safety fix, that host release could drop the final library + // guard and unmap this function before it returned. + assert!(unsafe { release_continuation() }); + assert!(unsafe { plugin_dropped() }); + assert_eq!(unsafe { library_probe() }, 0x4e52_5632_u64); + }); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn native_api_v2_safe_stream_fixture_dispatches_through_embedded_core() { + let _guard = NATIVE_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_plugin(); + let provider = EmbeddedFakeProvider::spawn( + b"data: {\"delta\":\"embedded-stream-response\"}\n\n", + "text/event-stream", + ); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "2", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_targeted_v2_plugin", + ), + ); + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("safe streaming native API v2 fixture should load"); + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_native".into(), + enabled: true, + config: Map::from_iter([("target_url".into(), json!(provider.url))]), + }); + initialize_plugins_exact(plugin_config) + .await + .expect("safe streaming native API v2 fixture should initialize"); + let mut cleanup = NativePluginTestCleanup::new(); + cleanup.mark_plugin_configuration_active(); + + let original_provider_called = Arc::new(AtomicBool::new(false)); + let original_provider_called_for_fn = original_provider_called.clone(); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("embedded-targeted-native-v2-stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({"model": "caller-model", "prompt": "hello"}), + }) + .func(Arc::new(move |_| { + original_provider_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .expect("embedded targeted stream dispatch should open"); + + assert_eq!( + stream + .next() + .await + .expect("targeted provider should emit one event") + .expect("targeted provider event should succeed"), + json!({"delta": "embedded-stream-response"}) + ); + assert!(stream.next().await.is_none()); + assert!(!original_provider_called.load(Ordering::SeqCst)); + let captured = String::from_utf8(provider.request()).expect("captured request should be UTF-8"); + assert!(captured.contains("authorization: Bearer fixture-target\r\n")); + assert!(captured.contains("\"fixture_targeted_stream\":true")); + + drop(cleanup); + activation.clear(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn native_api_v2_safe_fixture_forwards_passthrough_inside_embedded_core() { + let _guard = NATIVE_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_plugin(); + let manifest_ref = write_raw_manifest( + fixture.manifest_dir.path(), + &native_manifest_text( + "fixture_native", + &format!("={}", env!("CARGO_PKG_VERSION")), + "2", + &fixture.library_path.to_string_lossy(), + "nemo_relay_fixture_targeted_v2_plugin", + ), + ); + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("safe native API v2 fixture should load"); + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_native".into(), + enabled: true, + config: Map::from_iter([("target_url".into(), json!("http://127.0.0.1:1/unused"))]), + }); + initialize_plugins_exact(plugin_config) + .await + .expect("safe native API v2 fixture should initialize"); + let mut cleanup = NativePluginTestCleanup::new(); + cleanup.mark_plugin_configuration_active(); + + let response = llm_call_execute( + LlmCallExecuteParams::builder() + .name("fixture_passthrough_llm") + .request(LlmRequest { + headers: Map::new(), + content: json!({"prompt": "buffered"}), + }) + .func(Arc::new(|request| { + Box::pin(async move { Ok(json!({"ordinary": request.content})) }) + })) + .build(), + ) + .await + .expect("safe buffered pass-through should use the ordinary continuation"); + assert_eq!(response, json!({"ordinary": {"prompt": "buffered"}})); + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("fixture_passthrough_llm_stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({"prompt": "stream"}), + }) + .func(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter( + (0..40).map(|index| Ok(json!({"index": index}))), + ))) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .expect("safe streaming pass-through should open"); + for expected in 0..40 { + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + assert_eq!( + stream + .next() + .await + .expect("all pass-through events should arrive") + .expect("pass-through event should succeed"), + json!({"index": expected}) + ); + } + assert!(stream.next().await.is_none()); + + let mut failed = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("fixture_passthrough_llm_stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({"prompt": "late failure"}), + }) + .func(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(json!({"first": true})), + Err(nemo_relay::error::FlowError::Internal( + "fixture late failure".into(), + )), + ]))) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .expect("late-failing pass-through stream should open"); + assert_eq!( + failed.next().await.unwrap().unwrap(), + json!({"first": true}) + ); + assert!( + failed + .next() + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("fixture late failure") + ); + + let provider_dropped = Arc::new(AtomicBool::new(false)); + let provider_dropped_for_fn = provider_dropped.clone(); + let provider_started = Arc::new(AtomicBool::new(false)); + let provider_started_for_fn = provider_started.clone(); + let cancelled = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("fixture_passthrough_llm_stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({"prompt": "cancel"}), + }) + .func(Arc::new(move |_| { + let dropped = provider_dropped_for_fn.clone(); + let started = provider_started_for_fn.clone(); + Box::pin(async move { + let stream = LlmJsonStream::new(PendingDropStream { dropped }); + started.store(true, Ordering::SeqCst); + Ok(stream) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .expect("pending pass-through stream should open"); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !provider_started.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("downstream provider stream should start before cancellation"); + drop(cancelled); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !provider_dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("consumer cancellation should drop downstream production"); + + let open_passthrough = |id: usize| { + llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("fixture_passthrough_llm_stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({"id": id}), + }) + .func(Arc::new(move |_| { + Box::pin(async move { + Ok(LlmJsonStream::new(tokio_stream::iter([Ok( + json!({"id": id}), + )]))) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + }; + let (left, right) = tokio::join!(open_passthrough(1), open_passthrough(2)); + let mut left = left.expect("first concurrent pass-through should open"); + let mut right = right.expect("second concurrent pass-through should open"); + assert_eq!(left.next().await.unwrap().unwrap(), json!({"id": 1})); + assert_eq!(right.next().await.unwrap().unwrap(), json!({"id": 2})); + assert!(left.next().await.is_none()); + assert!(right.next().await.is_none()); + + drop(cleanup); + activation.clear(); +} + #[test] fn native_manifest_writer_escapes_toml_strings() { let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); @@ -2040,6 +2590,148 @@ struct BuiltFixture { library_path: PathBuf, } +struct PendingDropStream { + dropped: Arc, +} + +impl tokio_stream::Stream for PendingDropStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } +} + +impl Drop for PendingDropStream { + fn drop(&mut self) { + self.dropped.store(true, Ordering::SeqCst); + } +} + +struct EmbeddedFakeProvider { + url: String, + request: std::sync::mpsc::Receiver>, + thread: Option>>, +} + +impl EmbeddedFakeProvider { + fn spawn(body: &'static [u8], content_type: &'static str) -> Self { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::time::Duration; + + let listener = TcpListener::bind("127.0.0.1:0").expect("embedded provider should bind"); + listener + .set_nonblocking(true) + .expect("embedded provider listener should be nonblocking"); + let address = listener.local_addr().expect("embedded provider address"); + let (request_tx, request) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || -> Result<(), String> { + let deadline = std::time::Instant::now() + Duration::from_secs(30); + let (mut socket, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) + if error.kind() == std::io::ErrorKind::WouldBlock + && std::time::Instant::now() < deadline => + { + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + return Err(format!( + "embedded provider timed out waiting for a request: {error}" + )); + } + Err(error) => { + return Err(format!("embedded provider failed to accept: {error}")); + } + } + }; + socket.set_nonblocking(false).map_err(|error| { + format!("embedded provider failed to set blocking mode: {error}") + })?; + socket + .set_read_timeout(Some(Duration::from_secs(5))) + .map_err(|error| { + format!("embedded provider failed to set read timeout: {error}") + })?; + let mut request_bytes = Vec::new(); + let mut buffer = [0_u8; 4096]; + loop { + let read = socket.read(&mut buffer).map_err(|error| { + format!("embedded provider failed to read request: {error}") + })?; + if read == 0 { + break; + } + request_bytes.extend_from_slice(&buffer[..read]); + let Some(header_end) = request_bytes + .windows(4) + .position(|window| window == b"\r\n\r\n") + else { + continue; + }; + let content_length = String::from_utf8_lossy(&request_bytes[..header_end]) + .lines() + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + }) + .unwrap_or(0); + if request_bytes.len() >= header_end + 4 + content_length { + break; + } + } + request_tx + .send(request_bytes) + .map_err(|_| "embedded provider request receiver was dropped".to_owned())?; + let headers = format!( + "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + socket + .write_all(headers.as_bytes()) + .map_err(|error| format!("embedded provider failed to write headers: {error}"))?; + socket + .write_all(body) + .map_err(|error| format!("embedded provider failed to write body: {error}"))?; + Ok(()) + }); + Self { + url: format!("http://{address}/v1/chat/completions"), + request, + thread: Some(thread), + } + } + + fn request(&self) -> Vec { + self.request + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("embedded provider request should arrive") + } +} + +impl Drop for EmbeddedFakeProvider { + fn drop(&mut self) { + if let Some(thread) = self.thread.take() { + let result = thread.join(); + if !std::thread::panicking() { + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => panic!("embedded provider thread failed: {error}"), + Err(panic) => std::panic::resume_unwind(panic), + } + } + } + } +} + fn build_fixture_plugin() -> BuiltFixture { let _ = spdlog::init_log_crate_proxy(); log::set_max_level(log::LevelFilter::Info); diff --git a/crates/core/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 3a9856429..ad9b52b4d 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -31,6 +31,73 @@ fn buffers_partial_frames_across_pushes() { assert_eq!(events[0].data, json!({"a": 1})); } +#[test] +fn preserves_utf8_code_points_split_across_transport_chunks() { + let mut decoder = SseEventDecoder::new(); + let frame = "data: {\"text\":\"hello 🦀\"}\n\n".as_bytes(); + let split = frame.iter().position(|byte| *byte == 0xf0).unwrap() + 2; + + assert!(decoder.push_bytes(&frame[..split]).unwrap().is_empty()); + let events = decoder.push_bytes(&frame[split..]).unwrap(); + + assert_eq!(events.len(), 1); + assert_eq!(events[0].data, json!({"text": "hello 🦀"})); +} + +#[test] +fn rejects_an_unterminated_frame_above_the_retention_limit() { + let mut decoder = SseEventDecoder::new(); + let mut frame = b"data: \"".to_vec(); + frame.resize(MAX_SSE_FRAME_BYTES + 1, b'x'); + + let error = decoder.push_bytes(&frame).unwrap_err().to_string(); + + assert!(error.contains("SSE frame exceeded"), "{error}"); + assert!(error.contains(&MAX_SSE_FRAME_BYTES.to_string()), "{error}"); +} + +#[test] +fn preserves_completed_events_before_an_oversized_partial_frame() { + let mut decoder = SseEventDecoder::new(); + let mut bytes = b"data: {\"chunk\":\"first\"}\n\ndata: \"".to_vec(); + bytes.resize(MAX_SSE_FRAME_BYTES + 32, b'x'); + + let mut results = decoder.push_bytes_results(&bytes).into_iter(); + + assert_eq!( + results.next().unwrap().unwrap().data, + json!({"chunk": "first"}) + ); + let error = results.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("SSE frame exceeded"), "{error}"); + assert!(results.next().is_none()); +} + +#[test] +fn rejects_an_oversized_terminated_frame() { + let mut decoder = SseEventDecoder::new(); + let mut frame = b"data: \"".to_vec(); + frame.resize(MAX_SSE_FRAME_BYTES + 1, b'x'); + frame.extend_from_slice(b"\n\n"); + + let error = decoder.push_bytes(&frame).unwrap_err().to_string(); + + assert!(error.contains("SSE frame exceeded"), "{error}"); + assert!(error.contains(&MAX_SSE_FRAME_BYTES.to_string()), "{error}"); +} + +#[test] +fn rejects_invalid_utf8_in_a_terminated_frame() { + let mut decoder = SseEventDecoder::new(); + + let error = decoder + .push_bytes(b"data: \xff\n\n") + .unwrap_err() + .to_string(); + + assert!(error.contains("invalid UTF-8 SSE frame"), "{error}"); +} + #[test] fn normalizes_crlf_terminator_split_across_pushes() { let mut decoder = SseEventDecoder::new(); @@ -69,6 +136,19 @@ fn surfaces_final_partial_frame_on_finish() { assert_eq!(trailing.data, json!({"end": true})); } +#[test] +fn rejects_invalid_utf8_in_a_final_partial_frame() { + let mut decoder = SseEventDecoder::new(); + assert!(decoder.push_bytes(b"data: \xff").unwrap().is_empty()); + + let error = decoder.finish().unwrap_err().to_string(); + + assert!( + error.contains("incomplete or invalid UTF-8 at end of SSE stream"), + "{error}" + ); +} + #[test] fn drops_openai_chat_done_sentinel() { let mut decoder = SseEventDecoder::new(); diff --git a/crates/core/tests/unit/continuation_context_tests.rs b/crates/core/tests/unit/continuation_context_tests.rs index 9853f32a3..b37e976df 100644 --- a/crates/core/tests/unit/continuation_context_tests.rs +++ b/crates/core/tests/unit/continuation_context_tests.rs @@ -5,6 +5,9 @@ use super::*; use crate::api::optimization::{ LlmOptimizationRecorder, record_llm_optimization_contribution, scope_llm_optimization_recorder, }; +use crate::api::runtime::llm_dispatch_context::{ + LlmDispatchTargetContext, current_llm_dispatch_target, +}; use crate::api::runtime::scope_stack::{ TASK_SCOPE_STACK, active_event_uuid, create_scope_stack, current_scope_stack, with_active_event_uuid, @@ -64,6 +67,51 @@ fn continuation_context_restores_all_managed_execution_state() { }); } +#[test] +fn continuation_context_restores_target_across_task_and_thread_hops() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let event_uuid = uuid::Uuid::now_v7(); + let target = LlmDispatchTargetContext::try_new( + "https://provider.example/v1/chat/completions".into(), + std::collections::BTreeMap::from([( + "authorization".into(), + "Bearer target-secret".into(), + )]), + ) + .unwrap(); + let context = with_active_event_uuid(event_uuid, async { + MiddlewareContinuationContext::capture() + }) + .await; + let captured = context + .invoke_with_llm_dispatch_target(target.clone(), || async { + assert_eq!(current_llm_dispatch_target(), Some(target.clone())); + MiddlewareContinuationContext::capture() + }) + .await; + + let task_context = captured.clone(); + let task_target = tokio::spawn(async move { + task_context + .run(async { current_llm_dispatch_target() }) + .await + }) + .await + .unwrap(); + assert_eq!(task_target, Some(target.clone())); + + let thread_target = std::thread::spawn(move || { + tokio::runtime::Runtime::new() + .unwrap() + .block_on(captured.run(async { current_llm_dispatch_target() })) + }) + .join() + .unwrap(); + assert_eq!(thread_target, Some(target)); + }); +} + #[test] fn continuation_context_isolates_each_scope_stack_snapshot() { let runtime = tokio::runtime::Runtime::new().unwrap(); diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs new file mode 100644 index 000000000..ed769ce22 --- /dev/null +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -0,0 +1,631 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::future::Future; +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, mpsc}; +use std::time::Duration; + +use futures_util::StreamExt; +use serde_json::{Map, json}; + +use crate::api::llm::{ + LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, +}; +use crate::api::runtime::{MiddlewareContinuationContext, NemoRelayContextState, global_context}; +use crate::error::{FlowError, MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES, bounded_utf8}; + +use super::*; + +struct FakeProvider { + url: String, + request: mpsc::Receiver>, + thread: Option>, +} + +impl FakeProvider { + fn spawn(response: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").expect("fake provider should bind"); + listener + .set_nonblocking(true) + .expect("fake provider listener should be nonblocking"); + let address = listener.local_addr().expect("fake provider address"); + let (request_tx, request) = mpsc::channel(); + let thread = std::thread::spawn(move || { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + let (mut socket, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) + if error.kind() == std::io::ErrorKind::WouldBlock + && std::time::Instant::now() < deadline => + { + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("fake provider should accept: {error}"), + } + }; + socket + .set_nonblocking(false) + .expect("fake provider socket should be blocking"); + socket + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout should configure"); + let request_bytes = read_http_request(&mut socket); + request_tx + .send(request_bytes) + .expect("test should receive provider request"); + if let Err(error) = socket.write_all(&response) { + assert!( + matches!( + error.kind(), + std::io::ErrorKind::BrokenPipe + | std::io::ErrorKind::ConnectionAborted + | std::io::ErrorKind::ConnectionReset + ), + "fake provider should write response: {error}" + ); + } + }); + Self { + url: format!("http://{address}/v1/messages"), + request, + thread: Some(thread), + } + } + + fn request(&self) -> Vec { + self.request + .recv_timeout(Duration::from_secs(5)) + .expect("fake provider request should arrive") + } +} + +impl Drop for FakeProvider { + fn drop(&mut self) { + if let Some(thread) = self.thread.take() { + thread.join().expect("fake provider thread should finish"); + } + } +} + +fn read_http_request(socket: &mut std::net::TcpStream) -> Vec { + let mut request = Vec::new(); + let mut buffer = [0_u8; 4096]; + loop { + let read = socket + .read(&mut buffer) + .expect("request read should succeed"); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n") else { + continue; + }; + let content_length = String::from_utf8_lossy(&request[..header_end]) + .lines() + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + }) + .unwrap_or(0); + if request.len() >= header_end + 4 + content_length { + break; + } + } + request +} + +fn response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { + let mut response = format!("HTTP/1.1 {status}\r\nConnection: close\r\n"); + for (name, value) in headers { + response.push_str(name); + response.push_str(": "); + response.push_str(value); + response.push_str("\r\n"); + } + response.push_str(&format!("Content-Length: {}\r\n\r\n", body.len())); + let mut response = response.into_bytes(); + response.extend_from_slice(body); + response +} + +fn target(url: String, headers: BTreeMap) -> LlmDispatchTargetContext { + LlmDispatchTargetContext::try_new(url, headers).expect("test target should be valid") +} + +fn request() -> LlmRequest { + LlmRequest { + headers: Map::from_iter([ + ("authorization".into(), json!("Bearer request-secret")), + ( + "x-nemo-relay-internal-dispatch-url".into(), + json!("http://attacker.invalid"), + ), + ]), + content: json!({"model": "selected", "prompt": "hello"}), + } +} + +async fn scope_test_target(target: LlmDispatchTargetContext, future: F) -> F::Output { + let event_uuid = uuid::Uuid::now_v7(); + crate::api::runtime::with_active_event_uuid( + event_uuid, + scope_llm_dispatch_target(Some(event_uuid), target, future), + ) + .await +} + +#[tokio::test] +async fn buffered_target_runs_after_downstream_middleware_and_ignores_host_callback() { + let provider = FakeProvider::spawn(response( + "200 OK", + &[("Content-Type", "application/json")], + br#"{"id":"selected-response"}"#, + )); + let target = target( + format!("{}?api_key=url-secret", provider.url), + BTreeMap::from([ + ("authorization".into(), "Bearer target-secret".into()), + ("x-target".into(), "selected".into()), + ]), + ); + let debug = format!("{target:?}"); + assert!(!debug.contains("target-secret")); + assert!(!debug.contains("url-secret")); + + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_for_fn = fallback_called.clone(); + let terminal = targeted_llm_execution(Arc::new(move |_| { + fallback_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(json!({"id": "wrong-provider"})) }) + })); + let middleware_ran = Arc::new(AtomicBool::new(false)); + let middleware_ran_for_fn = middleware_ran.clone(); + let downstream = Arc::new(move |mut request: LlmRequest| { + middleware_ran_for_fn.store(true, Ordering::SeqCst); + request.content["middleware"] = json!(true); + terminal(request) + }); + + let result = scope_test_target(target, downstream(request())) + .await + .expect("targeted request should succeed"); + + assert_eq!(result, json!({"id": "selected-response"})); + assert!(middleware_ran.load(Ordering::SeqCst)); + assert!(!fallback_called.load(Ordering::SeqCst)); + let captured = String::from_utf8(provider.request()).expect("request should be UTF-8"); + assert!(captured.starts_with("POST /v1/messages?api_key=url-secret HTTP/1.1\r\n")); + assert!(captured.contains("authorization: Bearer target-secret\r\n")); + assert!(captured.contains("x-target: selected\r\n")); + assert!(captured.contains(r#"{"middleware":true,"model":"selected","prompt":"hello"}"#)); + assert!(!captured.contains("request-secret")); + assert!(!captured.contains("attacker.invalid")); +} + +#[tokio::test] +async fn malformed_success_json_is_an_internal_provider_failure() { + let provider = FakeProvider::spawn(response( + "200 OK", + &[("Content-Type", "application/json")], + b"not-json", + )); + + let error = dispatch_buffered(&target(provider.url.clone(), BTreeMap::new()), request()) + .await + .expect_err("malformed successful response should fail"); + + assert!(matches!( + error, + FlowError::Internal(message) + if message == "targeted LLM provider returned malformed response JSON" + )); + let _ = provider.request(); +} + +#[tokio::test] +async fn buffered_success_body_is_bounded_with_or_without_content_length() { + for response in [ + response( + "200 OK", + &[("Content-Type", "application/json")], + b"123456789", + ), + b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Type: application/json\r\nTransfer-Encoding: chunked\r\n\r\n9\r\n123456789\r\n0\r\n\r\n".to_vec(), + ] { + let provider = FakeProvider::spawn(response); + let target = target(provider.url.clone(), BTreeMap::new()); + let response = send(&target, request(), Some(HTTP_REQUEST_TIMEOUT)) + .await + .expect("provider should return a successful HTTP response"); + let error = bounded_success_body(&target, response, 8) + .await + .expect_err("oversized successful response should fail"); + assert!(matches!( + error, + FlowError::Internal(message) + if message == "targeted LLM provider response exceeded the 8-byte buffered body limit" + )); + let _ = provider.request(); + } +} + +#[tokio::test] +async fn buffered_http_failure_is_bounded_and_filters_headers() { + let body = vec![b'x'; MAX_UPSTREAM_FAILURE_BODY_BYTES + 1024]; + let provider = FakeProvider::spawn(response( + "429 Too Many Requests", + &[ + ("Retry-After", "2"), + ("Set-Cookie", "secret=true"), + ("Authorization", "Bearer response-secret"), + ], + &body, + )); + + let error = dispatch_buffered(&target(provider.url.clone(), BTreeMap::new()), request()) + .await + .expect_err("429 should fail"); + let FlowError::Upstream(failure) = error else { + panic!("expected structured upstream failure"); + }; + assert_eq!(failure.status, Some(429)); + assert_eq!(failure.class, UpstreamFailureClass::RetryableStatus); + assert_eq!(failure.body.len(), MAX_UPSTREAM_FAILURE_BODY_BYTES); + assert_eq!( + failure.headers.get("retry-after").map(String::as_str), + Some("2") + ); + assert!(!failure.headers.contains_key("set-cookie")); + assert!(!failure.headers.contains_key("authorization")); + let _ = provider.request(); +} + +#[tokio::test] +async fn streaming_http_failure_is_structured_and_filters_headers() { + let provider = FakeProvider::spawn(response( + "503 Service Unavailable", + &[("Retry-After", "3"), ("Set-Cookie", "secret=true")], + b"provider unavailable", + )); + + let error = + match dispatch_stream(&target(provider.url.clone(), BTreeMap::new()), request()).await { + Ok(_) => panic!("503 should fail before opening a stream"), + Err(error) => error, + }; + let FlowError::Upstream(failure) = error else { + panic!("expected structured upstream failure"); + }; + assert_eq!(failure.status, Some(503)); + assert_eq!(failure.class, UpstreamFailureClass::RetryableStatus); + assert_eq!(failure.body, "provider unavailable"); + assert_eq!( + failure.headers.get("retry-after").map(String::as_str), + Some("3") + ); + assert!(!failure.headers.contains_key("set-cookie")); + let _ = provider.request(); +} + +#[tokio::test] +async fn redirects_are_returned_without_following() { + let provider = FakeProvider::spawn(response( + "302 Found", + &[("Location", "http://127.0.0.1:9/should-not-run")], + b"redirect", + )); + + let error = dispatch_buffered(&target(provider.url.clone(), BTreeMap::new()), request()) + .await + .expect_err("redirect should fail"); + let FlowError::Upstream(failure) = error else { + panic!("expected HTTP failure"); + }; + assert_eq!(failure.status, Some(302)); + assert_eq!(failure.class, UpstreamFailureClass::Other); + let _ = provider.request(); +} + +#[tokio::test] +async fn streaming_target_decodes_events_empty_streams_and_late_errors() { + let provider = FakeProvider::spawn(response( + "200 OK", + &[("Content-Type", "Text/Event-Stream; Charset=UTF-8")], + b"data: {\"delta\":\"hello\"}\n\ndata: not-json\n\n", + )); + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_for_fn = fallback_called.clone(); + let terminal = targeted_llm_stream_execution(Arc::new(move |_| { + fallback_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(LlmJsonStream::new(futures_util::stream::empty())) }) + })); + let stream_target = target(provider.url.clone(), BTreeMap::new()); + let mut stream = scope_test_target(stream_target, terminal(request())) + .await + .expect("stream should open"); + assert_eq!( + stream.next().await.unwrap().unwrap(), + json!({"delta": "hello"}) + ); + assert!(stream.next().await.unwrap().is_err()); + assert!(!fallback_called.load(Ordering::SeqCst)); + let _ = provider.request(); + + let empty_provider = FakeProvider::spawn(response( + "200 OK", + &[("Content-Type", "text/event-stream")], + b"", + )); + let mut empty = dispatch_stream( + &target(empty_provider.url.clone(), BTreeMap::new()), + request(), + ) + .await + .expect("empty stream should open"); + assert!(empty.next().await.is_none()); + let _ = empty_provider.request(); + + let cancelled_provider = FakeProvider::spawn(response( + "200 OK", + &[("Content-Type", "text/event-stream")], + b"data: {\"delta\":\"first\"}\n\ndata: {\"delta\":\"second\"}\n\n", + )); + let mut cancelled = dispatch_stream( + &target(cancelled_provider.url.clone(), BTreeMap::new()), + request(), + ) + .await + .expect("cancellable stream should open"); + assert_eq!( + cancelled.next().await.unwrap().unwrap(), + json!({"delta": "first"}) + ); + cancelled + .close() + .await + .expect("stream should close cleanly"); + assert!(cancelled.next().await.is_none()); + let _ = cancelled_provider.request(); +} + +#[tokio::test] +async fn streaming_target_rejects_missing_or_non_sse_content_type() { + for (headers, expected_type) in [ + (vec![], ""), + ( + vec![("Content-Type", "application/json")], + "application/json", + ), + ] { + let provider = FakeProvider::spawn(response( + "200 OK", + &headers, + br#"{"error":"not an event stream"}"#, + )); + let error = match dispatch_stream(&target(provider.url.clone(), BTreeMap::new()), request()) + .await + { + Ok(_) => panic!("a successful non-SSE response should not open a stream"), + Err(error) => error, + }; + let FlowError::Internal(message) = error else { + panic!("expected a non-HTTP stream setup failure"); + }; + assert!(message.contains("expected Content-Type text/event-stream")); + assert!(message.contains(expected_type)); + assert!(message.contains("not an event stream")); + let _ = provider.request(); + } +} + +#[test] +fn target_validation_rejects_unsafe_transport_inputs() { + for (url, headers) in [ + ("ftp://provider.example/v1", BTreeMap::new()), + ("https://user:secret@provider.example/v1", BTreeMap::new()), + ( + "https://provider.example/v1", + BTreeMap::from([("host".into(), "attacker.invalid".into())]), + ), + ( + "https://provider.example/v1", + BTreeMap::from([( + "x-nemo-relay-internal-dispatch-url".into(), + "http://attacker.invalid".into(), + )]), + ), + ( + "https://provider.example/v1", + BTreeMap::from([ + ("Authorization".into(), "Bearer first".into()), + ("authorization".into(), "Bearer second".into()), + ]), + ), + ] { + assert!(LlmDispatchTargetContext::try_new(url.into(), headers).is_err()); + } +} + +#[test] +fn target_validation_reports_malformed_headers() { + let error = LlmDispatchTargetContext::try_new( + "https://provider.example/v1".into(), + BTreeMap::from([("bad header".into(), "value".into())]), + ) + .expect_err("header name containing a space should be rejected"); + let FlowError::InvalidArgument(message) = error else { + panic!("expected invalid header name argument"); + }; + assert_eq!( + message, + "LLM continuation contained an invalid target header name" + ); + + let error = LlmDispatchTargetContext::try_new( + "https://provider.example/v1".into(), + BTreeMap::from([("x-target".into(), "line one\nline two".into())]), + ) + .expect_err("header value containing a newline should be rejected"); + let FlowError::InvalidArgument(message) = error else { + panic!("expected invalid header value argument"); + }; + assert_eq!( + message, + "LLM continuation target header x-target had an invalid value" + ); +} + +#[test] +fn bounded_utf8_truncates_long_multibyte_value_at_character_boundary() { + let expected = "a".repeat(MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1); + let value = format!("{expected}\u{e9}"); + assert_eq!(value.len(), MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES + 1); + + let bounded = bounded_utf8(value, MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES); + + assert_eq!(bounded, expected); + assert_eq!(bounded.len(), MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1); +} + +#[tokio::test] +async fn transport_failures_do_not_fall_back_to_the_host_callback() { + let listener = TcpListener::bind("127.0.0.1:0").expect("temporary listener should bind"); + let address = listener.local_addr().unwrap(); + drop(listener); + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_for_fn = fallback_called.clone(); + let terminal = targeted_llm_execution(Arc::new(move |_| { + fallback_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(json!({"wrong": true})) }) + })); + let target = target( + format!("http://{address}/v1/chat/completions?api_key=transport-secret"), + BTreeMap::new(), + ); + + let error = scope_test_target(target, terminal(request())) + .await + .expect_err("connection should fail"); + let FlowError::Upstream(failure) = error else { + panic!("expected transport failure"); + }; + assert_eq!(failure.status, None); + assert_eq!(failure.class, UpstreamFailureClass::Connection); + assert!(!failure.body.contains("transport-secret")); + assert!(!fallback_called.load(Ordering::SeqCst)); +} + +#[test] +fn nested_buffered_managed_call_does_not_inherit_outer_target() { + let _guard = crate::shared_runtime::runtime_owner_test_mutex() + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::shared_runtime::reset_runtime_owner_for_tests(); + *global_context() + .write() + .unwrap_or_else(|error| error.into_inner()) = NemoRelayContextState::new(); + + tokio::runtime::Runtime::new().unwrap().block_on(async { + let target = target("http://127.0.0.1:9/v1/messages".into(), BTreeMap::new()); + let context = crate::api::runtime::with_active_event_uuid(uuid::Uuid::now_v7(), async { + MiddlewareContinuationContext::capture() + }) + .await; + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_for_fn = Arc::clone(&fallback_called); + + let response = context + .invoke_with_llm_dispatch_target(target, || async move { + assert!(current_llm_dispatch_target().is_some()); + let response = Box::pin(llm_call_execute( + LlmCallExecuteParams::builder() + .name("nested-buffered") + .request(request()) + .func(Arc::new(move |_| { + assert!(current_llm_dispatch_target().is_none()); + fallback_called_for_fn.store(true, Ordering::SeqCst); + Box::pin(async { Ok(json!({"provider": "ordinary"})) }) + })) + .build(), + )) + .await?; + assert!(current_llm_dispatch_target().is_some()); + Ok::<_, FlowError>(response) + }) + .await + .expect("nested ordinary call should use its own provider callback"); + + assert_eq!(response, json!({"provider": "ordinary"})); + assert!(fallback_called.load(Ordering::SeqCst)); + }); +} + +#[test] +fn nested_streaming_managed_call_isolated_during_lazy_polling() { + let _guard = crate::shared_runtime::runtime_owner_test_mutex() + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::shared_runtime::reset_runtime_owner_for_tests(); + *global_context() + .write() + .unwrap_or_else(|error| error.into_inner()) = NemoRelayContextState::new(); + + tokio::runtime::Runtime::new().unwrap().block_on(async { + let target = target("http://127.0.0.1:9/v1/messages".into(), BTreeMap::new()); + let context = crate::api::runtime::with_active_event_uuid(uuid::Uuid::now_v7(), async { + MiddlewareContinuationContext::capture() + }) + .await; + let provider_opened = Arc::new(AtomicBool::new(false)); + let provider_opened_for_fn = Arc::clone(&provider_opened); + let provider_polled = Arc::new(AtomicBool::new(false)); + let provider_polled_for_fn = Arc::clone(&provider_polled); + + let chunk = context + .invoke_with_llm_dispatch_target(target, || async move { + assert!(current_llm_dispatch_target().is_some()); + let mut stream = Box::pin(llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("nested-streaming") + .request(request()) + .func(Arc::new(move |_| { + assert!(current_llm_dispatch_target().is_none()); + provider_opened_for_fn.store(true, Ordering::SeqCst); + let provider_polled = Arc::clone(&provider_polled_for_fn); + Box::pin(async move { + Ok(LlmJsonStream::new(futures_util::stream::once(async move { + assert!(current_llm_dispatch_target().is_none()); + provider_polled.store(true, Ordering::SeqCst); + Ok(json!({"delta": "ordinary"})) + }))) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .build(), + )) + .await?; + assert!(current_llm_dispatch_target().is_some()); + let chunk = stream.next().await.expect("nested stream should emit")?; + assert!(current_llm_dispatch_target().is_some()); + assert!(stream.next().await.is_none()); + assert!(current_llm_dispatch_target().is_some()); + Ok::<_, FlowError>(chunk) + }) + .await + .expect("nested ordinary stream should use its own provider callback"); + + assert_eq!(chunk, json!({"delta": "ordinary"})); + assert!(provider_opened.load(Ordering::SeqCst)); + assert!(provider_polled.load(Ordering::SeqCst)); + }); +} diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index f0978fc6d..ced07d06a 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -5,7 +5,7 @@ use super::*; -use std::collections::VecDeque; +use std::collections::{BTreeMap, VecDeque}; use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; @@ -28,12 +28,25 @@ use crate::api::runtime::subscriber_dispatcher::{ }; use crate::api::runtime::{ BuiltinLlmCodec, LlmSanitizeRequestContext, LlmSanitizeResponseContext, - MiddlewareContinuationLease, NemoRelayContextState, TASK_SCOPE_STACK, current_scope_stack, - global_context, with_active_event_uuid, + MiddlewareContinuationLease, NemoRelayContextState, TASK_SCOPE_STACK, + current_llm_dispatch_target, current_scope_stack, global_context, with_active_event_uuid, }; use crate::codec::openai_chat::OpenAIChatCodec; use crate::codec::response::AnnotatedLlmResponse; +#[derive(Clone, Debug, PartialEq)] +enum LlmContinuationOutcomeV2 { + Success { response: Json }, + Failure { error: LlmContinuationFailureV2 }, +} + +#[derive(Clone, Debug, PartialEq)] +enum LlmContinuationStreamEventV2 { + Chunk { chunk: Json }, + Failure { error: LlmContinuationFailureV2 }, + Done, +} + struct ThreadScopeStackRestore(Option); impl ThreadScopeStackRestore { @@ -111,6 +124,447 @@ unsafe extern "C" fn complete_native_next_result( let _ = sender.send(result); } +unsafe extern "C" fn complete_typed_llm_result( + user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, +) { + let sender = unsafe { + Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) + }; + let outcome = typed_llm_outcome(response_json, error_json); + let _ = sender.send(outcome); +} + +fn typed_llm_outcome( + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, +) -> LlmContinuationOutcomeV2 { + if !response_json.is_null() { + LlmContinuationOutcomeV2::Success { + response: parse_json_arg(response_json, "typed LLM response") + .expect("host emitted valid provider JSON"), + } + } else if !error_json.is_null() { + LlmContinuationOutcomeV2::Failure { + error: parse_json_arg(error_json, "typed LLM failure") + .and_then(|value| { + serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) + }) + .expect("host emitted a valid typed LLM failure"), + } + } else { + LlmContinuationOutcomeV2::Failure { + error: non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + "host could not allocate a typed LLM callback value".into(), + ), + } + } +} + +unsafe extern "C" fn reject_unexpected_typed_llm_result( + _user_data: *mut c_void, + _response_json: *const NemoRelayNativeString, + _error_json: *const NemoRelayNativeString, +) { + panic!("invalid dispatch unexpectedly invoked its callback"); +} + +unsafe extern "C" fn record_typed_llm_stream_open( + user_data: *mut c_void, + stream: *const NemoRelayNativeLlmStreamV2, + error_json: *const NemoRelayNativeString, +) { + let sender = unsafe { + Box::from_raw( + user_data as *mut tokio::sync::oneshot::Sender>, + ) + }; + let result = typed_llm_stream_open_result(stream, error_json); + let _ = sender.send(result); +} + +fn typed_llm_stream_open_result( + stream: *const NemoRelayNativeLlmStreamV2, + error_json: *const NemoRelayNativeString, +) -> Result { + if !stream.is_null() { + Ok(stream as usize) + } else if !error_json.is_null() { + match parse_json_arg(error_json, "typed LLM stream open error").and_then(|value| { + serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) + }) { + Ok(error) => Err(error), + Err(status) => Err(non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!("invalid typed stream open error: {status:?}"), + )), + } + } else { + Err(non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + "host could not allocate a typed stream-open callback value".into(), + )) + } +} + +fn test_dispatch_target(url: &str) -> nemo_relay_plugin::LlmContinuationTargetV2 { + nemo_relay_plugin::LlmContinuationTargetV2 { + url: url.into(), + headers: BTreeMap::new(), + } +} + +unsafe extern "C" fn record_typed_llm_stream_next( + user_data: *mut c_void, + chunk_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, + done: bool, +) { + let sender = unsafe { + Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) + }; + let event = typed_llm_stream_event(chunk_json, error_json, done); + let _ = sender.send(event); +} + +fn typed_llm_stream_event( + chunk_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, + done: bool, +) -> LlmContinuationStreamEventV2 { + if !chunk_json.is_null() { + LlmContinuationStreamEventV2::Chunk { + chunk: parse_json_arg(chunk_json, "typed LLM stream chunk") + .expect("host emitted a valid typed LLM stream chunk"), + } + } else if !error_json.is_null() { + LlmContinuationStreamEventV2::Failure { + error: parse_json_arg(error_json, "typed LLM stream failure") + .and_then(|value| { + serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) + }) + .expect("host emitted a valid typed LLM stream failure"), + } + } else if done { + LlmContinuationStreamEventV2::Done + } else { + LlmContinuationStreamEventV2::Failure { + error: non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + "host could not allocate a typed provider-stream callback value".into(), + ), + } + } +} + +fn start_typed_llm_call( + next: *const NemoRelayNativeAsyncNext, + invocation: *const NemoRelayNativeString, +) -> tokio::sync::oneshot::Receiver { + let (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next, + invocation, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + receiver +} + +fn start_typed_llm_stream( + next: *const NemoRelayNativeAsyncNext, + invocation: *const NemoRelayNativeString, + output: *const NemoRelayNativeAsyncStream, +) -> tokio::sync::oneshot::Receiver> { + let (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + next, + invocation, + output, + record_typed_llm_stream_open, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + receiver +} + +fn pull_typed_llm_stream( + stream: *const NemoRelayNativeLlmStreamV2, +) -> tokio::sync::oneshot::Receiver { + let (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + stream, + record_typed_llm_stream_next, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + receiver +} + +#[derive(Default)] +struct TypedLlmResultCallbackState { + callbacks: AtomicUsize, + outcome: Mutex>, + notified: tokio::sync::Notify, +} + +unsafe extern "C" fn record_typed_llm_result_state( + user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, +) { + let state = unsafe { &*user_data.cast::() }; + state.callbacks.fetch_add(1, Ordering::SeqCst); + *state + .outcome + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some(typed_llm_outcome(response_json, error_json)); + state.notified.notify_one(); +} + +#[derive(Default)] +struct TypedLlmStreamOpenCallbackState { + callbacks: AtomicUsize, + result: Mutex>>, + notified: tokio::sync::Notify, +} + +unsafe extern "C" fn record_typed_llm_stream_open_state( + user_data: *mut c_void, + stream: *const NemoRelayNativeLlmStreamV2, + error_json: *const NemoRelayNativeString, +) { + let state = unsafe { &*user_data.cast::() }; + state.callbacks.fetch_add(1, Ordering::SeqCst); + *state + .result + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some(typed_llm_stream_open_result(stream, error_json)); + state.notified.notify_one(); +} + +#[derive(Default)] +struct TypedLlmStreamNextCallbackState { + callbacks: AtomicUsize, + event: Mutex>, + notified: tokio::sync::Notify, +} + +unsafe extern "C" fn record_typed_llm_stream_next_state( + user_data: *mut c_void, + chunk_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, + done: bool, +) { + let state = unsafe { &*user_data.cast::() }; + state.callbacks.fetch_add(1, Ordering::SeqCst); + *state + .event + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some(typed_llm_stream_event(chunk_json, error_json, done)); + state.notified.notify_one(); +} + +#[derive(Default)] +struct NativeForwardTerminalState { + callbacks: AtomicUsize, + notified: tokio::sync::Notify, +} + +unsafe extern "C" fn record_native_forward_terminal(user_data: *mut c_void) { + let state = unsafe { &*user_data.cast::() }; + state.callbacks.fetch_add(1, Ordering::SeqCst); + state.notified.notify_one(); +} + +unsafe extern "C" fn return_v2_callback_state_and_release_next( + user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + _completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let state = unsafe { &*user_data.cast::() }.load(Ordering::Acquire) as u32; + unsafe { native_async_next_release(next) }; + state +} + +unsafe extern "C" fn complete_v2_callback_and_release_next( + _user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let response = native_string_from_json(&json!({"safe": true})).unwrap(); + assert_eq!( + unsafe { native_async_completion_resolve_json(completion, response) }, + NemoRelayStatus::Ok + ); + unsafe { + native_string_free(response); + native_async_next_release(next); + } + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +struct V2StreamContractState { + callback_state: u32, + finish: bool, +} + +unsafe extern "C" fn return_v2_stream_state_and_release_handles( + user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + stream: *const NemoRelayNativeAsyncStream, +) -> u32 { + let state = unsafe { &*user_data.cast::() }; + if state.finish { + assert_eq!( + unsafe { native_async_stream_finish(stream) }, + NemoRelayStatus::Ok + ); + } + unsafe { + native_async_next_release(next); + native_async_stream_release(stream); + } + state.callback_state +} + +unsafe extern "C" fn pending_native_task( + _user_data: *mut c_void, + _task: *const NemoRelayNativeAsyncTaskV2, +) -> u32 { + NemoRelayNativeAsyncCallbackState::Pending as u32 +} + +#[derive(Clone, Copy)] +enum CooperativeTaskMode { + Pending, + SettleCompletionOnSecondPoll, + CompleteWithoutSettlement, + InvalidState, + PushTwoChunks, +} + +struct CooperativeTaskState { + mode: CooperativeTaskMode, + polls: AtomicUsize, + completion: usize, + stream: usize, + started: Option>, + frees: Arc, +} + +unsafe extern "C" fn poll_cooperative_task( + user_data: *mut c_void, + task: *const NemoRelayNativeAsyncTaskV2, +) -> u32 { + let state = unsafe { &*user_data.cast::() }; + let poll = state.polls.fetch_add(1, Ordering::SeqCst); + match state.mode { + CooperativeTaskMode::Pending => { + if poll == 0 + && let Some(started) = &state.started + { + let _ = started.send(0); + } + NemoRelayNativeAsyncCallbackState::Pending as u32 + } + CooperativeTaskMode::SettleCompletionOnSecondPoll if poll == 0 => { + unsafe { native_async_task_retain_v2(task) }; + if state + .started + .as_ref() + .is_none_or(|started| started.send(task as usize).is_err()) + { + unsafe { native_async_task_release_v2(task) }; + } + NemoRelayNativeAsyncCallbackState::Pending as u32 + } + CooperativeTaskMode::SettleCompletionOnSecondPoll => { + let result = native_string_from_json(&json!({"cooperative": true})).unwrap(); + assert_eq!( + unsafe { + native_async_completion_resolve_json( + state.completion as *const NemoRelayNativeAsyncCompletion, + result, + ) + }, + NemoRelayStatus::Ok + ); + unsafe { native_string_free(result) }; + NemoRelayNativeAsyncCallbackState::Complete as u32 + } + CooperativeTaskMode::CompleteWithoutSettlement => { + NemoRelayNativeAsyncCallbackState::Complete as u32 + } + CooperativeTaskMode::InvalidState => 99, + CooperativeTaskMode::PushTwoChunks => { + let first = native_string_from_json(&json!({"chunk": 1})).unwrap(); + let second = native_string_from_json(&json!({"chunk": 2})).unwrap(); + let stream = state.stream as *const NemoRelayNativeAsyncStream; + let push_json = build_native_host_api_v4().v3.async_stream_push_json; + let result = if poll == 0 { + assert_eq!(unsafe { push_json(stream, first) }, NemoRelayStatus::Ok); + assert_eq!( + unsafe { push_json(stream, second) }, + NemoRelayStatus::WouldBlock + ); + if let Some(started) = &state.started { + let _ = started.send(0); + } + NemoRelayNativeAsyncCallbackState::Pending as u32 + } else { + assert_eq!(unsafe { push_json(stream, second) }, NemoRelayStatus::Ok); + assert_eq!( + unsafe { native_async_stream_finish(stream) }, + NemoRelayStatus::Ok + ); + NemoRelayNativeAsyncCallbackState::Complete as u32 + }; + unsafe { + native_string_free(first); + native_string_free(second); + } + result + } + } +} + +unsafe extern "C" fn free_cooperative_task(user_data: *mut c_void) { + let state = unsafe { Box::from_raw(user_data.cast::()) }; + if state.completion != 0 { + unsafe { + native_async_completion_release( + state.completion as *const NemoRelayNativeAsyncCompletion, + ) + }; + } + if state.stream != 0 { + unsafe { native_async_stream_release(state.stream as *const NemoRelayNativeAsyncStream) }; + } + state.frees.fetch_add(1, Ordering::SeqCst); +} + #[derive(Default)] struct NativeStreamCallbackState { error: Mutex>, @@ -309,7 +763,8 @@ fn native_test_adapter( relay_compat: "^0.7".into(), allows_multiple_components: false, plugin: Mutex::new(plugin), - _library: libloading::os::unix::Library::this().into(), + library: Some(libloading::os::unix::Library::this().into()), + retain_library_on_drop: AtomicBool::new(false), }), } } @@ -418,7 +873,7 @@ fn assert_native_digest_edges() { fn assert_native_host_api_versions() { let current = native_host_api(); - let legacy = native_host_api_legacy(); + let legacy = native_host_api_v2(); assert!(!current.is_null()); assert!(!legacy.is_null()); assert_eq!( @@ -437,6 +892,9 @@ async fn native_async_wait_and_rejection_cover_dropped_and_aborted_continuations cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), next_abort: Mutex::new(None), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }), @@ -457,6 +915,9 @@ async fn native_async_wait_and_rejection_cover_dropped_and_aborted_continuations cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(true), next_abort: Mutex::new(Some(abort)), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -481,8 +942,8 @@ async fn native_async_wait_and_rejection_cover_dropped_and_aborted_continuations unsafe { native_async_completion_release(completion_ref) }; } -#[test] -fn native_stream_callback_guard_covers_terminal_drop_modes() { +#[tokio::test] +async fn native_stream_callback_guard_covers_terminal_drop_modes() { let (sender, _receiver) = tokio::sync::mpsc::channel(1); let stream = Arc::new(NativeAsyncStream { sender: Mutex::new(Some(sender)), @@ -490,6 +951,9 @@ fn native_stream_callback_guard_covers_terminal_drop_modes() { settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -548,6 +1012,9 @@ async fn native_async_stream_forwarding_reports_conversion_and_stream_errors() { settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }) @@ -681,8 +1148,8 @@ async fn native_async_result_entrypoint_covers_llm_and_stream_continuations() { } } -#[test] -fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { +#[tokio::test] +async fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { let chunk = native_string("null"); let no_sender = Arc::new(NativeAsyncStream { @@ -691,6 +1158,9 @@ fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -717,6 +1187,9 @@ fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -739,6 +1212,9 @@ fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -760,6 +1236,9 @@ fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { settled: AtomicBool::new(true), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -835,6 +1314,9 @@ async fn native_async_stream_next_entrypoint_validates_handle_kind_and_request() settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: tokio::runtime::Handle::current(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -1168,6 +1650,14 @@ fn assert_native_json_parsing_boundaries() { assert_last_error_contains("optional is not valid JSON"); unsafe { native_string_free(invalid_json) }; + let invalid_utf8 = + Box::into_raw(Box::new(NativeHostString(vec![0xff]))) as *mut NemoRelayNativeString; + assert_eq!( + parse_json_arg(invalid_utf8, "invalid UTF-8 JSON").unwrap_err(), + NemoRelayStatus::InvalidUtf8 + ); + unsafe { native_string_free(invalid_utf8) }; + assert_eq!( parse_json_arg(ptr::null(), "null JSON").unwrap_err(), NemoRelayStatus::InvalidJson @@ -1205,12 +1695,22 @@ fn assert_native_json_output_and_host_api() { take_json_from_native_string(json_out, "unused").unwrap(), json!({"ok": true}) ); + fail_native_string_allocation_after(0); + json_out = ptr::null_mut(); + assert_eq!( + write_native_json(&json!({"ok": true}), &mut json_out), + NemoRelayStatus::Internal + ); + assert!(json_out.is_null()); let host_api = unsafe { &*native_host_api() }; - assert_eq!(host_api.abi_version, NEMO_RELAY_NATIVE_ABI_VERSION); + assert_eq!( + host_api.abi_version, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS + ); assert_eq!( host_api.struct_size, - std::mem::size_of::() + std::mem::size_of::() ); } @@ -1244,6 +1744,9 @@ fn native_async_next_abi_runs_tool_llm_and_stream_continuations() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -1281,6 +1784,9 @@ fn native_async_next_abi_runs_tool_llm_and_stream_continuations() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, _receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -1338,8 +1844,11 @@ fn native_async_next_reports_a_revoked_continuation_without_calling_the_provider None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::oneshot::channel(); + let (sender, receiver) = tokio::sync::oneshot::channel::>(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -1468,27 +1977,288 @@ fn native_async_next_result_supports_repeated_concurrent_calls() { } #[test] -fn native_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { +fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); - let captured_stack = create_scope_stack(); - let captured_scope = captured_stack - .read() - .unwrap_or_else(|error| error.into_inner()) - .top() - .uuid - .to_string(); - let next = with_scope_stack(captured_stack, || { - Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::Tool(Arc::new(|value| { - Box::pin(async move { - Ok(json!({ - "value": value, - "scope": crate::api::runtime::task_scope_top().uuid.to_string(), - })) - }) + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|request| { + let target = current_llm_dispatch_target().expect("typed target is bound"); + Box::pin(async move { + assert!(request.headers.is_empty()); + assert_eq!( + target, + LlmDispatchTargetContext::try_new( + "https://provider.example/v1/chat/completions".into(), + BTreeMap::from([("authorization".into(), "Bearer target-secret".into(),)]), + ) + .unwrap() + ); + Err(FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(429), + body: "rate limited".into(), + headers: BTreeMap::from([("retry-after".into(), "1".into())]), + class: UpstreamFailureClass::RetryableStatus, + })) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model"}), + }, + target: nemo_relay_plugin::LlmContinuationTargetV2 { + headers: BTreeMap::from([("authorization".into(), "Bearer target-secret".into())]), + ..test_dispatch_target("https://provider.example/v1/chat/completions") + }, + }) + .unwrap(), + ) + .unwrap(); + let receiver = start_typed_llm_call(next_ref, dispatch); + let outcome = runtime.block_on(receiver).unwrap(); + assert_eq!( + outcome, + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::Http { + status: 429, + body: "rate limited".into(), + headers: BTreeMap::from([("retry-after".into(), "1".into())]), + }, + } + ); + + unsafe { + native_string_free(dispatch); + native_async_next_release(next_ref); + } +} + +#[test] +fn native_api_v2_releasing_next_cancels_a_pending_targeted_call() { + struct DropSignal(std::sync::mpsc::Sender<()>); + + impl Drop for DropSignal { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(move |_| { + let started_tx = started_tx.clone(); + let guard = DropSignal(dropped_tx.clone()); + Box::pin(async move { + let _guard = guard; + let _ = started_tx.send(()); + std::future::pending::>().await + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({}), + }, + target: test_dispatch_target("https://provider.example/v1/chat/completions"), + }) + .unwrap(), + ) + .unwrap(); + let receiver = start_typed_llm_call(next_ref, dispatch); + runtime + .block_on(async { tokio::task::spawn_blocking(move || started_rx.recv()).await }) + .unwrap() + .unwrap(); + unsafe { native_async_next_release(next_ref) }; + + let outcome = runtime.block_on(receiver).unwrap(); + assert!(matches!( + outcome, + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + } + )); + dropped_rx + .recv_timeout(Duration::from_secs(1)) + .expect("releasing next should abort and drop targeted provider work"); + unsafe { native_string_free(dispatch) }; +} + +#[test] +fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let provider_calls = Arc::new(AtomicUsize::new(0)); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm({ + let provider_calls = Arc::clone(&provider_calls); + Arc::new(move |_| { + provider_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(Json::Null) }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({}), + }, + target: test_dispatch_target("/v1/chat/completions"), + }) + .unwrap(), + ) + .unwrap(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + reject_unexpected_typed_llm_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("absolute HTTP(S) URL"); + assert_eq!(provider_calls.load(Ordering::SeqCst), 0); + + unsafe { + native_string_free(dispatch); + native_async_next_release(next_ref); + } +} + +#[test] +fn native_api_v2_rejects_prohibited_target_headers() { + let request = LlmRequest { + headers: Map::new(), + content: json!({}), + }; + let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); + target + .headers + .insert("x-nemo-relay-internal-dispatch-url".into(), "secret".into()); + assert_eq!( + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { + request: request.clone(), + target, + }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("host-owned or prohibited"); + + let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); + target + .headers + .insert("host".into(), "attacker.invalid".into()); + assert_eq!( + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request, target }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("host-owned or prohibited"); +} + +#[test] +fn native_api_v2_rejects_target_url_credentials() { + let target = test_dispatch_target("https://user:secret@provider.example/v1/chat/completions"); + assert_eq!( + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({}), + }, + target, + }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("absolute HTTP(S) URL"); +} + +#[test] +fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { + let error = typed_llm_failure(FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(429), + body: "é".repeat(crate::error::MAX_UPSTREAM_FAILURE_BODY_BYTES), + headers: BTreeMap::from([ + ("Authorization".into(), "Bearer secret".into()), + ("Set-Cookie".into(), "session=secret".into()), + ("Retry-After".into(), "1".into()), + ( + "X-Request-ID".into(), + format!( + "{}é", + "x".repeat(crate::error::MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) + ), + ), + ]), + class: UpstreamFailureClass::RetryableStatus, + })); + + let LlmContinuationFailureV2::Http { body, headers, .. } = error else { + panic!("expected an upstream failure"); + }; + assert_eq!(body.len(), crate::error::MAX_UPSTREAM_FAILURE_BODY_BYTES); + assert!(body.is_char_boundary(body.len())); + assert_eq!(headers.get("retry-after").map(String::as_str), Some("1")); + assert_eq!( + headers.get("x-request-id").map(String::len), + Some(crate::error::MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) + ); + assert!(!headers.contains_key("authorization")); + assert!(!headers.contains_key("set-cookie")); +} + +#[test] +fn native_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let captured_stack = create_scope_stack(); + let captured_scope = captured_stack + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid + .to_string(); + let next = with_scope_stack(captured_stack, || { + Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(Arc::new(|value| { + Box::pin(async move { + Ok(json!({ + "value": value, + "scope": crate::api::runtime::task_scope_top().uuid.to_string(), + })) + }) })), runtime.handle().clone(), None, @@ -1528,51 +2298,2450 @@ fn native_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { } } -fn native_continuation_context_observation( - expected_stack: &ScopeStackHandle, - expected_event_uuid: uuid::Uuid, -) -> Json { - let expected_scope_uuid = expected_stack - .read() - .unwrap_or_else(|error| error.into_inner()) - .top() - .uuid; - let visible_scope_uuid = current_scope_stack() +#[test] +fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let provider_polls = Arc::new(AtomicUsize::new(0)); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream({ + let provider_polls = Arc::clone(&provider_polls); + Arc::new(move |request| { + let provider_polls = Arc::clone(&provider_polls); + assert!(request.headers.is_empty()); + Box::pin(async move { + assert_eq!( + current_llm_dispatch_target().expect("typed target is bound"), + LlmDispatchTargetContext::try_new( + "https://provider.example/v1/messages".into(), + BTreeMap::new(), + ) + .unwrap() + ); + let mut events = VecDeque::from(vec![ + Ok(json!({"type": "content_block_delta", "delta": {"text": "hi"}})), + Err(FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(503), + body: "unavailable".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::ModelUnavailable, + })), + ]); + Ok(LlmJsonStream::new(futures_util::stream::poll_fn( + move |_| { + provider_polls.fetch_add(1, Ordering::SeqCst); + std::task::Poll::Ready(events.pop_front()) + }, + ))) + }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (output_stream, receiver) = test_native_output_stream(1); + let output_stream_ref = + Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model", "stream": true}), + }, + target: test_dispatch_target("https://provider.example/v1/messages"), + }) + .unwrap(), + ) + .unwrap(); + + let open_receiver = start_typed_llm_stream(next_ref, dispatch, output_stream_ref); + let provider_stream = runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), open_receiver) + .await + .expect("typed LLM stream should open") + .expect("open callback should be delivered") + .expect("provider stream should open") + }) as *const NemoRelayNativeLlmStreamV2; + assert_eq!( + provider_polls.load(Ordering::SeqCst), + 0, + "opening a provider stream must not read ahead" + ); + let mut events = Vec::new(); + loop { + let receiver = pull_typed_llm_stream(provider_stream); + let event = runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), receiver) + .await + .expect("typed LLM stream should produce an event") + .expect("next callback should be delivered") + }); + let terminal = matches!( + event, + LlmContinuationStreamEventV2::Done | LlmContinuationStreamEventV2::Failure { .. } + ); + events.push(event); + assert_eq!( + provider_polls.load(Ordering::SeqCst), + events.len(), + "each ABI next must poll exactly one immediately-ready provider item" + ); + if terminal { + break; + } + } + assert_eq!( + events, + vec![ + LlmContinuationStreamEventV2::Chunk { + chunk: json!({ + "type": "content_block_delta", + "delta": {"text": "hi"}, + }), + }, + LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::Http { + status: 503, + body: "unavailable".into(), + headers: BTreeMap::new(), + }, + }, + ] + ); + + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&output_stream), + }); + unsafe { + native_async_llm_stream_release_v2(provider_stream); + native_string_free(dispatch); + native_async_next_release(next_ref); + native_async_stream_release(output_stream_ref); + } +} + +#[test] +fn native_api_v2_provider_pulls_restore_scope_and_target_after_pending_wakes() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let ambient_scope = crate::api::runtime::task_scope_top().uuid; + let captured_stack = create_scope_stack(); + let captured_scope = captured_stack .read() .unwrap_or_else(|error| error.into_inner()) .top() .uuid; - json!({ - "scope_stack": visible_scope_uuid == expected_scope_uuid, - "active_event_uuid": active_event_uuid() == Some(expected_event_uuid), - "publication_context": crate::api::runtime::subscriber_dispatcher::publication_context::() - .is_some_and(|context| context.as_str() == "native-continuation"), - "publication_buffer": capture_nested_publication_buffer().is_some(), - "optimization_recorder": current_llm_optimization_recorder().is_some(), - }) + assert_ne!(ambient_scope, captured_scope); + let context = with_scope_stack(captured_stack, MiddlewareContinuationContext::capture); + let target_url = "https://provider.example/v1/context"; + let target = LlmDispatchTargetContext::try_new(target_url.into(), BTreeMap::new()).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + let observations_for_stream = Arc::clone(&observations); + let mut pending = true; + let mut index = 0; + let provider_stream = LlmJsonStream::new(futures_util::stream::poll_fn(move |cx| { + if pending { + pending = false; + cx.waker().wake_by_ref(); + return std::task::Poll::Pending; + } + pending = true; + let scope = crate::api::runtime::task_scope_top().uuid; + let target = current_llm_dispatch_target() + .expect("provider pull should restore its dispatch target"); + observations_for_stream + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push((scope, target)); + let chunk = json!({"index": index}); + index += 1; + std::task::Poll::Ready(Some(Ok(chunk))) + })); + let (output, _output_receiver) = test_native_output_stream(1); + let provider = Arc::new(NativeLlmProviderStreamV2 { + stream: tokio::sync::Mutex::new(Some(provider_stream)), + runtime: runtime.handle().clone(), + context, + target, + output, + lifecycle: Mutex::new(NativeLlmProviderStreamLifecycleV2::Idle), + }); + let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; + + for index in 0..2 { + let receiver = pull_typed_llm_stream(provider_ref); + assert_eq!( + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), receiver) + .await + .expect("pending provider pull should be woken") + .expect("provider pull callback should be delivered") + }), + LlmContinuationStreamEventV2::Chunk { + chunk: json!({"index": index}) + } + ); + assert!(current_llm_dispatch_target().is_none()); + assert_eq!(crate::api::runtime::task_scope_top().uuid, ambient_scope); + } + + assert_eq!( + *observations + .lock() + .unwrap_or_else(|error| error.into_inner()), + vec![ + ( + captured_scope, + LlmDispatchTargetContext::try_new(target_url.into(), BTreeMap::new()).unwrap(), + ), + ( + captured_scope, + LlmDispatchTargetContext::try_new(target_url.into(), BTreeMap::new()).unwrap(), + ), + ] + ); + unsafe { native_async_llm_stream_release_v2(provider_ref) }; } #[test] -fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuations() { +fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); - let expected_stack = create_scope_stack(); - let expected_event_uuid = uuid::Uuid::now_v7(); - let expected = json!({ - "scope_stack": true, - "active_event_uuid": true, - "publication_context": true, - "publication_buffer": true, - "optimization_recorder": true, - }); - - runtime.block_on(TASK_SCOPE_STACK.scope( - expected_stack.clone(), - with_task_publication_context( - Some(Arc::new(String::from("native-continuation"))), - scope_llm_optimization_recorder( + let provider_calls = Arc::new(AtomicUsize::new(0)); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream({ + let provider_calls = Arc::clone(&provider_calls); + Arc::new(move |_request| { + provider_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(json!({"chunk": 1})), + Ok(json!({"chunk": 2})), + ]))) + }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (stream, receiver) = test_native_output_stream(1); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let mut output = NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }; + let request = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + let terminal = Arc::new(NativeForwardTerminalState::default()); + + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + request, + stream_ref, + record_native_forward_terminal, + Arc::as_ptr(&terminal).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), async { + while provider_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + for _ in 0..10 { + tokio::task::yield_now().await; + } + }) + .await + .expect("pass-through provider should start"); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 0); + + assert_eq!(output.next().await.unwrap().unwrap(), json!({"chunk": 1})); + tokio::time::timeout(Duration::from_secs(1), terminal.notified.notified()) + .await + .expect("terminal callback should follow output settlement"); + assert_eq!(output.next().await.unwrap().unwrap(), json!({"chunk": 2})); + assert!(output.next().await.is_none()); + }); + assert!(stream.settled.load(Ordering::Acquire)); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + unsafe { + native_string_free(request); + native_async_next_release(next_ref); + native_async_stream_release(stream_ref); + } +} + +#[test] +fn native_api_v2_direct_stream_forwarding_preserves_downstream_failure() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_request| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(json!({"chunk": 1})), + Err(FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(503), + body: "provider unavailable".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::ModelUnavailable, + })), + ]))) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (stream, receiver) = test_native_output_stream(2); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let mut output = NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }; + let request = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + let terminal = Arc::new(NativeForwardTerminalState::default()); + + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + request, + stream_ref, + record_native_forward_terminal, + Arc::as_ptr(&terminal).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + assert_eq!(output.next().await.unwrap().unwrap(), json!({"chunk": 1})); + let error = output + .next() + .await + .expect("forwarded failure should be emitted") + .unwrap_err(); + assert!(matches!( + error, + FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(503), + .. + }) + )); + assert!(output.next().await.is_none()); + tokio::time::timeout(Duration::from_secs(1), terminal.notified.notified()) + .await + .expect("failure should settle the terminal callback"); + }); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + unsafe { + native_string_free(request); + native_async_next_release(next_ref); + native_async_stream_release(stream_ref); + } +} + +#[test] +fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(Mutex::new(Some(started_tx))); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(move |_request| { + let started_tx = Arc::clone(&started_tx); + Box::pin(async move { + if let Some(started_tx) = started_tx + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = started_tx.send(()); + } + Ok(LlmJsonStream::new(futures_util::stream::pending())) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (stream, receiver) = test_native_output_stream(1); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let output = NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }; + let request = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + let terminal = Arc::new(NativeForwardTerminalState::default()); + + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + request, + stream_ref, + record_native_forward_terminal, + Arc::as_ptr(&terminal).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(started_rx).unwrap(); + drop(output); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), terminal.notified.notified()) + .await + .expect("consumer cancellation should settle the terminal callback"); + }); + assert!(stream.cancelled.load(Ordering::Acquire)); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + unsafe { + native_string_free(request); + native_async_next_release(next_ref); + native_async_stream_release(stream_ref); + } +} + +#[test] +fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight() { + struct StreamDropSignal(std::sync::mpsc::Sender<()>); + + impl Drop for StreamDropSignal { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let polls = Arc::new(AtomicUsize::new(0)); + let polls_for_stream = Arc::clone(&polls); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); + let drop_signal = StreamDropSignal(dropped_tx); + let provider_stream = LlmJsonStream::new(futures_util::stream::poll_fn(move |_| { + let _ = &drop_signal; + if polls_for_stream.fetch_add(1, Ordering::SeqCst) == 0 { + let _ = started_tx.send(()); + } + std::task::Poll::Pending + })); + let (provider, _output_receiver) = test_native_provider_stream(&runtime, provider_stream); + let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; + + let callback = Arc::new(TypedLlmStreamNextCallbackState::default()); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next_state, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime + .block_on(async { tokio::task::spawn_blocking(move || started_rx.recv()).await }) + .unwrap() + .unwrap(); + assert_eq!(polls.load(Ordering::SeqCst), 1); + + let (overlap_sender, _overlap_receiver) = + tokio::sync::oneshot::channel::(); + let overlap_state = Box::into_raw(Box::new(overlap_sender)); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next, + overlap_state.cast(), + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("pending next"); + unsafe { drop(Box::from_raw(overlap_state)) }; + + // Releasing the plugin's reference while `next` is pending is safe because + // the callback task retains its own provider-stream reference. Release + // cancels the pull before dropping that reference. + unsafe { native_async_llm_stream_release_v2(provider_ref) }; + + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("cancelled provider next should settle"); + }); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + assert!(matches!( + callback + .event + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + }) + )); + dropped_rx + .recv_timeout(Duration::from_secs(1)) + .expect("release should abort the pull and drop the provider stream"); + runtime.block_on(tokio::task::yield_now()); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + assert_eq!(polls.load(Ordering::SeqCst), 1); +} + +#[test] +fn native_api_v2_output_drop_cancels_pending_provider_next_once() { + struct StreamDropSignal(std::sync::mpsc::Sender<()>); + + impl Drop for StreamDropSignal { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); + let drop_signal = StreamDropSignal(dropped_tx); + let mut started_tx = Some(started_tx); + let provider_stream = LlmJsonStream::new(futures_util::stream::poll_fn(move |_| { + let _ = &drop_signal; + if let Some(started_tx) = started_tx.take() { + let _ = started_tx.send(()); + } + std::task::Poll::Pending + })); + let (provider, output_receiver) = test_native_provider_stream(&runtime, provider_stream); + let output = Arc::clone(&provider.output); + let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; + let callback = Arc::new(TypedLlmStreamNextCallbackState::default()); + + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next_state, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime + .block_on(async { tokio::task::spawn_blocking(move || started_rx.recv()).await }) + .unwrap() + .unwrap(); + + drop(NativeAsyncStreamReceiver { + receiver: output_receiver, + stream: output, + }); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("output cancellation should settle the pending next callback"); + }); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + assert!(matches!( + callback + .event + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + }) + )); + + unsafe { native_async_llm_stream_release_v2(provider_ref) }; + dropped_rx + .recv_timeout(Duration::from_secs(1)) + .expect("releasing the cancelled provider should drop its stream"); + runtime.block_on(tokio::task::yield_now()); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); +} + +#[test] +fn native_api_v2_handles_256_concurrent_buffered_dispatches() { + const DISPATCH_COUNT: usize = 256; + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + let provider_calls = Arc::new(AtomicUsize::new(0)); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm({ + let provider_calls = Arc::clone(&provider_calls); + Arc::new(move |_| { + provider_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { + tokio::task::yield_now().await; + Ok(json!({"ok": true})) + }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model"}), + }, + target: test_dispatch_target("https://provider.example/v1/chat/completions"), + }) + .unwrap(), + ) + .unwrap(); + + let mut receivers = Vec::with_capacity(DISPATCH_COUNT); + for _ in 0..DISPATCH_COUNT { + receivers.push(start_typed_llm_call(next_ref, dispatch)); + } + + let outcomes = runtime.block_on(async { + tokio::time::timeout( + Duration::from_secs(5), + futures_util::future::join_all(receivers), + ) + .await + .expect("buffered dispatch stress test should not deadlock") + }); + assert_eq!(outcomes.len(), DISPATCH_COUNT); + assert!(outcomes.into_iter().all(|outcome| { + outcome + == Ok(LlmContinuationOutcomeV2::Success { + response: json!({"ok": true}), + }) + })); + assert_eq!(provider_calls.load(Ordering::SeqCst), DISPATCH_COUNT); + + unsafe { + native_string_free(dispatch); + native_async_next_release(next_ref); + } +} + +#[test] +fn native_api_v2_isolates_concurrent_dispatch_targets() { + const DISPATCH_COUNT: usize = 64; + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|request| { + Box::pin(async move { + tokio::task::yield_now().await; + let target = current_llm_dispatch_target().expect("typed target is bound"); + let index = request + .content + .get("index") + .and_then(|index| index.as_u64()) + .expect("test request should contain an integer index"); + let url = format!("https://provider-{index}.example/v1/chat/completions"); + let authorization = format!("Bearer target-{index}"); + assert_eq!( + target, + LlmDispatchTargetContext::try_new( + url.clone(), + BTreeMap::from([("authorization".into(), authorization.clone())]), + ) + .unwrap() + ); + Ok(json!({ + "request": request.content, + "url": url, + "authorization": authorization, + })) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let mut receivers = Vec::with_capacity(DISPATCH_COUNT); + + for index in 0..DISPATCH_COUNT { + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"index": index}), + }, + target: nemo_relay_plugin::LlmContinuationTargetV2 { + headers: BTreeMap::from([( + "authorization".into(), + format!("Bearer target-{index}"), + )]), + ..test_dispatch_target(&format!( + "https://provider-{index}.example/v1/chat/completions" + )) + }, + }) + .unwrap(), + ) + .unwrap(); + let receiver = start_typed_llm_call(next_ref, dispatch); + unsafe { native_string_free(dispatch) }; + receivers.push((index, receiver)); + } + + runtime.block_on(async { + for (index, receiver) in receivers { + assert_eq!( + receiver + .await + .expect("dispatch callback should be delivered"), + LlmContinuationOutcomeV2::Success { + response: json!({ + "request": {"index": index}, + "url": format!("https://provider-{index}.example/v1/chat/completions"), + "authorization": format!("Bearer target-{index}"), + }), + } + ); + } + }); + + unsafe { native_async_next_release(next_ref) }; +} + +#[test] +fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_deadlock() { + const STREAM_COUNT: usize = 64; + const EVENT_COUNT: usize = 100; + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter( + (0..EVENT_COUNT).map(|index| Ok(json!({"index": index}))), + ))) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (output_stream, output_receiver) = + test_native_output_stream(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY); + let output_stream_ref = + Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model", "stream": true}), + }, + target: test_dispatch_target("https://provider.example/v1/chat/completions"), + }) + .unwrap(), + ) + .unwrap(); + + let mut open_receivers = Vec::with_capacity(STREAM_COUNT); + for _ in 0..STREAM_COUNT { + open_receivers.push(start_typed_llm_stream( + next_ref, + dispatch, + output_stream_ref, + )); + } + + let event_counts = runtime.block_on(async { + let streams = tokio::time::timeout( + Duration::from_secs(5), + futures_util::future::join_all(open_receivers), + ) + .await + .expect("all provider streams should open") + .into_iter() + .map(|result| { + result + .expect("open callback should be delivered") + .expect("provider stream should open") + as *const NemoRelayNativeLlmStreamV2 + }) + .collect::>(); + let drains = streams.into_iter().map(|provider_stream| async move { + let mut chunks = 0; + loop { + let receiver = pull_typed_llm_stream(provider_stream); + match receiver.await.expect("next callback should be delivered") { + LlmContinuationStreamEventV2::Chunk { .. } => chunks += 1, + LlmContinuationStreamEventV2::Done => break, + LlmContinuationStreamEventV2::Failure { error } => { + panic!("provider stream failed during stress test: {error:?}") + } + } + } + unsafe { native_async_llm_stream_release_v2(provider_stream) }; + chunks + }); + tokio::time::timeout( + Duration::from_secs(10), + futures_util::future::join_all(drains), + ) + .await + .expect("provider stream stress test should not deadlock") + }); + assert_eq!(event_counts, vec![EVENT_COUNT; STREAM_COUNT]); + + drop(NativeAsyncStreamReceiver { + receiver: output_receiver, + stream: Arc::clone(&output_stream), + }); + unsafe { + native_string_free(dispatch); + native_async_next_release(next_ref); + native_async_stream_release(output_stream_ref); + } +} + +fn test_v2_callback_user_data(ptr: *mut c_void) -> Arc { + Arc::new(NativeCallbackUserData { + ptr, + free_fn: None, + _instance: None, + }) +} + +fn empty_llm_next() -> LlmExecutionNextFn { + Arc::new(|request| Box::pin(async move { Ok(request.content) })) +} + +fn empty_llm_stream_next() -> LlmStreamExecutionNextFn { + Arc::new(|_| Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) })) +} + +fn test_llm_request() -> LlmRequest { + LlmRequest { + headers: Map::new(), + content: json!({"model": "test"}), + } +} + +fn test_native_output_stream( + capacity: usize, +) -> ( + Arc, + tokio::sync::mpsc::Receiver>, +) { + let (sender, receiver) = tokio::sync::mpsc::channel(capacity); + ( + Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + settled: AtomicBool::new(false), + downstream_aborts: Mutex::new(HashMap::new()), + settlement: Mutex::new(()), + before_settlement_lock: None, + _callback_user_data: None, + }), + receiver, + ) +} + +#[test] +fn cooperative_stream_error_settles_once_and_aborts_downstream() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (stream, mut receiver) = test_native_output_stream(1); + let downstream = runtime.spawn(std::future::pending::<()>()); + stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert(downstream.id(), downstream.abort_handle()); + + runtime.block_on(settle_native_async_stream_error( + Arc::clone(&stream), + "cooperative policy failed".into(), + )); + + let error = runtime + .block_on(receiver.recv()) + .expect("error settlement should emit one terminal item") + .expect_err("settlement item should be an error"); + assert!(matches!( + error, + FlowError::Internal(message) if message == "cooperative policy failed" + )); + assert!(stream.settled.load(Ordering::Acquire)); + assert!( + stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + ); + assert!( + stream + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_empty() + ); + assert!(runtime.block_on(downstream).unwrap_err().is_cancelled()); + + runtime.block_on(settle_native_async_stream_error( + Arc::clone(&stream), + "duplicate".into(), + )); + assert!(runtime.block_on(receiver.recv()).is_none()); +} + +fn test_native_provider_stream( + runtime: &tokio::runtime::Runtime, + stream: LlmJsonStream, +) -> ( + Arc, + tokio::sync::mpsc::Receiver>, +) { + let (output, receiver) = test_native_output_stream(1); + ( + Arc::new(NativeLlmProviderStreamV2 { + stream: tokio::sync::Mutex::new(Some(stream)), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + target: LlmDispatchTargetContext::try_new( + "https://provider.example/v1/chat/completions".into(), + BTreeMap::new(), + ) + .unwrap(), + output, + lifecycle: Mutex::new(NativeLlmProviderStreamLifecycleV2::Idle), + }), + receiver, + ) +} + +#[test] +fn native_async_buffered_wrapper_enforces_complete_callback_contract() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + + let user_data = test_v2_callback_user_data(ptr::null_mut()); + let wrapped = wrap_native_async_llm_execution_with_user_data( + complete_v2_callback_and_release_next, + user_data, + ); + assert_eq!( + runtime + .block_on(wrapped("safe", test_llm_request(), empty_llm_next())) + .unwrap(), + json!({"safe": true}) + ); + + for (state, expected) in [ + ( + NemoRelayNativeAsyncCallbackState::Complete as usize, + "returned Complete without settling", + ), + (99, "returned an invalid state"), + ] { + let callback_state = AtomicUsize::new(state); + let user_data = + test_v2_callback_user_data((&callback_state as *const AtomicUsize).cast_mut().cast()); + let wrapped = wrap_native_async_llm_execution_with_user_data( + return_v2_callback_state_and_release_next, + user_data, + ); + let error = runtime + .block_on(wrapped("invalid", test_llm_request(), empty_llm_next())) + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + } + + let callback_state = AtomicUsize::new(NemoRelayNativeAsyncCallbackState::Complete as usize); + let user_data = + test_v2_callback_user_data((&callback_state as *const AtomicUsize).cast_mut().cast()); + let wrapped = wrap_native_async_llm_execution_with_user_data( + return_v2_callback_state_and_release_next, + user_data, + ); + fail_native_string_allocation_after(0); + let error = runtime + .block_on(wrapped("allocation", test_llm_request(), empty_llm_next())) + .unwrap_err(); + assert!(error.to_string().contains("failed to allocate"), "{error}"); + + let callback_state = AtomicUsize::new(NemoRelayNativeAsyncCallbackState::Complete as usize); + let user_data = + test_v2_callback_user_data((&callback_state as *const AtomicUsize).cast_mut().cast()); + let wrapped = wrap_native_async_llm_execution_with_user_data( + return_v2_callback_state_and_release_next, + user_data, + ); + let error = + futures::executor::block_on(wrapped("no-runtime", test_llm_request(), empty_llm_next())) + .unwrap_err(); + assert!( + error.to_string().contains("requires a Tokio runtime"), + "{error}" + ); +} + +#[test] +fn native_async_stream_wrapper_enforces_callback_and_ownership_contracts() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + + for (state, expected_error) in [ + ( + V2StreamContractState { + callback_state: NemoRelayNativeAsyncCallbackState::Complete as u32, + finish: true, + }, + None, + ), + ( + V2StreamContractState { + callback_state: NemoRelayNativeAsyncCallbackState::Complete as u32, + finish: false, + }, + Some("returned Complete without finishing"), + ), + ( + V2StreamContractState { + callback_state: 99, + finish: false, + }, + Some("panicked or returned an invalid state"), + ), + ] { + let user_data = + test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( + return_v2_stream_state_and_release_handles, + user_data, + ); + let result = runtime.block_on(wrapped( + "stream-contract", + test_llm_request(), + empty_llm_stream_next(), + )); + match expected_error { + Some(expected) => { + let error = result.err().expect("contract failure should reject setup"); + assert!(error.to_string().contains(expected), "{error}"); + } + None => { + let mut stream = result.expect("finished callback should return a stream"); + assert!(runtime.block_on(stream.next()).is_none()); + } + } + } + + let state = V2StreamContractState { + callback_state: NemoRelayNativeAsyncCallbackState::Complete as u32, + finish: true, + }; + let user_data = + test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( + return_v2_stream_state_and_release_handles, + user_data, + ); + fail_native_string_allocation_after(0); + let error = runtime + .block_on(wrapped( + "allocation", + test_llm_request(), + empty_llm_stream_next(), + )) + .err() + .expect("injected allocation failure should reject stream setup"); + assert!(error.to_string().contains("failed to allocate"), "{error}"); + + std::thread::spawn(|| { + let state = V2StreamContractState { + callback_state: NemoRelayNativeAsyncCallbackState::Complete as u32, + finish: true, + }; + let user_data = + test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( + return_v2_stream_state_and_release_handles, + user_data, + ); + assert_eq!(native_string_live_allocations(), 0); + let error = futures::executor::block_on(wrapped( + "no-runtime", + test_llm_request(), + empty_llm_stream_next(), + )) + .err() + .expect("stream setup without a runtime should fail"); + assert!( + error.to_string().contains("requires a Tokio runtime"), + "{error}" + ); + assert_eq!(native_string_live_allocations(), 0); + }) + .join() + .expect("no-runtime stream validation thread should not panic"); +} + +#[test] +fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { + let (sender, receiver) = tokio::sync::oneshot::channel::>(); + let mut result_guard = NativeAsyncResultCallbackGuard { + cb: complete_native_next_result, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }; + result_guard.complete(Err(FlowError::InvalidArgument("invalid".into()))); + result_guard.complete(Ok(Json::Null)); + assert!( + receiver + .blocking_recv() + .unwrap() + .unwrap_err() + .contains("invalid") + ); + + let (sender, receiver) = tokio::sync::oneshot::channel::>(); + drop(NativeAsyncResultCallbackGuard { + cb: complete_native_next_result, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }); + assert!( + receiver + .blocking_recv() + .unwrap() + .unwrap_err() + .contains("cancelled") + ); + + let (sender, receiver) = tokio::sync::oneshot::channel::(); + let mut typed_guard = NativeLlmResultCallbackGuardV2 { + cb: complete_typed_llm_result, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }; + typed_guard.complete(Ok(json!({"ok": true}))); + typed_guard.complete(Ok(Json::Null)); + assert_eq!( + receiver.blocking_recv().unwrap(), + LlmContinuationOutcomeV2::Success { + response: json!({"ok": true}) + } + ); + + let (sender, receiver) = tokio::sync::oneshot::channel::(); + drop(NativeLlmResultCallbackGuardV2 { + cb: complete_typed_llm_result, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }); + assert!(matches!( + receiver.blocking_recv().unwrap(), + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + } + )); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (provider, _output_receiver) = + test_native_provider_stream(&runtime, LlmJsonStream::new(tokio_stream::empty())); + let (sender, receiver) = + tokio::sync::oneshot::channel::>(); + let mut open_guard = NativeLlmStreamOpenCallbackGuardV2 { + cb: record_typed_llm_stream_open, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }; + open_guard.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Transport, + "offline".into(), + )); + open_guard.success(Arc::clone(&provider)); + assert!(matches!( + receiver.blocking_recv().unwrap(), + Err(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Transport, + .. + }) + )); + + let (sender, receiver) = + tokio::sync::oneshot::channel::>(); + drop(NativeLlmStreamOpenCallbackGuardV2 { + cb: record_typed_llm_stream_open, + user_data: Box::into_raw(Box::new(sender)).cast::() as usize, + active: true, + _library_guard: None, + }); + assert!(matches!( + receiver.blocking_recv().unwrap(), + Err(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + }) + )); + drop(provider); + + let (stream, _receiver) = test_native_output_stream(1); + let terminal = NativeForwardTerminalState::default(); + let mut forward_guard = NativeLlmStreamForwardCallbackGuardV2 { + cb: record_native_forward_terminal, + user_data: (&terminal as *const NativeForwardTerminalState) as usize, + stream: Arc::clone(&stream), + active: true, + }; + forward_guard.settle(); + forward_guard.settle(); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + let (stream, _receiver) = test_native_output_stream(1); + let terminal = NativeForwardTerminalState::default(); + drop(NativeLlmStreamForwardCallbackGuardV2 { + cb: record_native_forward_terminal, + user_data: (&terminal as *const NativeForwardTerminalState) as usize, + stream: Arc::clone(&stream), + active: true, + }); + assert!(stream.cancelled.load(Ordering::Acquire)); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); +} + +#[test] +fn native_api_v2_split_callbacks_settle_when_native_string_allocation_fails() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let live_before = native_string_live_allocations(); + let dispatch = test_v2_dispatch_json(); + let live_with_dispatch = native_string_live_allocations(); + + assert_unary_allocation_failures_settle(&runtime, dispatch, live_with_dispatch); + + let open_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_| { + Box::pin(async { Err(FlowError::InvalidArgument("bad stream request".into())) }) + })), + runtime.handle().clone(), + None, + )); + let open_next_ref = Arc::into_raw(open_next) as *const NemoRelayNativeAsyncNext; + let (output, output_receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let open_callback = Arc::new(TypedLlmStreamOpenCallbackState::default()); + fail_native_string_allocation_after(0); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + open_next_ref, + dispatch, + output_ref, + record_typed_llm_stream_open_state, + Arc::as_ptr(&open_callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), open_callback.notified.notified()) + .await + .expect("stream-open allocation failure must still invoke its callback"); + }); + assert_eq!(open_callback.callbacks.load(Ordering::SeqCst), 1); + assert!(matches!( + open_callback + .result + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(Err(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + .. + })) + )); + assert_eq!(native_string_live_allocations(), live_with_dispatch); + drop(output_receiver); + unsafe { + native_async_next_release(open_next_ref); + native_async_stream_release(output_ref); + } + + for stream in [ + LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])), + LlmJsonStream::new(tokio_stream::iter(vec![Err(FlowError::InvalidArgument( + "provider failure".into(), + ))])), + ] { + let (provider, output_receiver) = test_native_provider_stream(&runtime, stream); + let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; + let callback = Arc::new(TypedLlmStreamNextCallbackState::default()); + fail_native_string_allocation_after(0); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next_state, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("provider-next allocation failure must still invoke its callback"); + }); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + assert!(matches!( + callback + .event + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + .. + } + }) + )); + assert_eq!(native_string_live_allocations(), live_with_dispatch); + drop(output_receiver); + unsafe { native_async_llm_stream_release_v2(provider_ref) }; + } + + unsafe { native_string_free(dispatch) }; + assert_eq!(native_string_live_allocations(), live_before); +} + +fn assert_unary_allocation_failures_settle( + runtime: &tokio::runtime::Runtime, + dispatch: *const NemoRelayNativeString, + live_with_dispatch: usize, +) { + for next in [ + Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|_| Box::pin(async { Ok(json!({"ok": true})) }))), + runtime.handle().clone(), + None, + )), + Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|_| { + Box::pin(async { Err(FlowError::InvalidArgument("bad request".into())) }) + })), + runtime.handle().clone(), + None, + )), + ] { + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let callback = Arc::new(TypedLlmResultCallbackState::default()); + fail_native_string_allocation_after(0); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + record_typed_llm_result_state, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("unary allocation failure must still invoke its callback"); + }); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + assert!(matches!( + callback + .outcome + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + .. + } + }) + )); + assert_eq!(native_string_live_allocations(), live_with_dispatch); + unsafe { native_async_next_release(next_ref) }; + } +} + +#[test] +fn native_api_v2_failure_mapping_covers_http_and_non_http_kinds() { + let cases = [ + ( + FlowError::Upstream(crate::error::UpstreamFailure { + status: None, + body: "timeout".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Timeout, + }), + LlmNonHttpFailureKindV2::Timeout, + ), + ( + FlowError::Upstream(crate::error::UpstreamFailure { + status: None, + body: "transport".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Connection, + }), + LlmNonHttpFailureKindV2::Transport, + ), + ( + FlowError::GuardrailRejected("blocked".into()), + LlmNonHttpFailureKindV2::Guardrail, + ), + ( + FlowError::InvalidArgument("bad request".into()), + LlmNonHttpFailureKindV2::InvalidRequest, + ), + ( + FlowError::Internal("broken".into()), + LlmNonHttpFailureKindV2::Internal, + ), + ]; + for (error, expected_kind) in cases { + let LlmContinuationFailureV2::NonHttp { kind, .. } = typed_llm_failure(error) else { + panic!("expected non-HTTP failure") + }; + assert_eq!(kind, expected_kind); + } + + let message = format!("{}é", "x".repeat(NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES)); + let LlmContinuationFailureV2::NonHttp { message, .. } = + non_http_llm_failure(LlmNonHttpFailureKindV2::Internal, message) + else { + panic!("expected non-HTTP failure") + }; + assert_eq!(message.len(), NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES); + assert!(message.is_char_boundary(message.len())); +} + +fn test_v2_dispatch_json() -> *mut NemoRelayNativeString { + native_string_from_json( + &serde_json::to_value(LlmContinuationInvocationV2 { + request: test_llm_request(), + target: test_dispatch_target("https://provider.example/v1/chat/completions"), + }) + .unwrap(), + ) + .unwrap() +} + +#[test] +fn native_api_v2_unary_entrypoint_reports_validation_panics_and_typed_failures() { + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + ptr::null(), + ptr::null(), + reject_unexpected_typed_llm_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let dispatch = test_v2_dispatch_json(); + let wrong_kind = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(Arc::new(|value| Box::pin(async move { Ok(value) }))), + runtime.handle().clone(), + None, + )); + let wrong_kind_ref = Arc::into_raw(wrong_kind) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + wrong_kind_ref, + dispatch, + reject_unexpected_typed_llm_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("requires an LLM execution continuation"); + + let valid_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(empty_llm_next()), + runtime.handle().clone(), + None, + )); + let valid_next_ref = Arc::into_raw(valid_next) as *const NemoRelayNativeAsyncNext; + let malformed = native_string("not-json"); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + valid_next_ref, + malformed, + reject_unexpected_typed_llm_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + let wrong_shape = native_string_from_json(&json!({"request": {}})).unwrap(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + valid_next_ref, + wrong_shape, + reject_unexpected_typed_llm_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + + let failures = [ + ( + FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(503), + body: "unavailable".into(), + headers: BTreeMap::from([("retry-after".into(), "1".into())]), + class: UpstreamFailureClass::RetryableStatus, + }), + Some(503), + None, + ), + ( + FlowError::Upstream(crate::error::UpstreamFailure { + status: None, + body: "timed out".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Timeout, + }), + None, + Some(LlmNonHttpFailureKindV2::Timeout), + ), + ( + FlowError::GuardrailRejected("blocked".into()), + None, + Some(LlmNonHttpFailureKindV2::Guardrail), + ), + ]; + for (failure, http_status, non_http_kind) in failures { + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(move |_| { + let failure = failure.clone(); + Box::pin(async move { Err(failure) }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let receiver = start_typed_llm_call(next_ref, dispatch); + let outcome = runtime.block_on(receiver).unwrap(); + match outcome { + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::Http { status, .. }, + } => assert_eq!(Some(status), http_status), + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { kind, .. }, + } => assert_eq!(Some(kind), non_http_kind), + other => panic!("unexpected continuation outcome: {other:?}"), + } + unsafe { native_async_next_release(next_ref) }; + } + + let panicking = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|_| { + Box::pin(async { + panic!("targeted provider panic"); + }) + })), + runtime.handle().clone(), + None, + )); + let panicking_ref = Arc::into_raw(panicking) as *const NemoRelayNativeAsyncNext; + let receiver = start_typed_llm_call(panicking_ref, dispatch); + assert!(matches!( + runtime.block_on(receiver).unwrap(), + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + .. + } + } + )); + + unsafe { + native_string_free(dispatch); + native_string_free(malformed); + native_string_free(wrong_shape); + native_async_next_release(wrong_kind_ref); + native_async_next_release(valid_next_ref); + native_async_next_release(panicking_ref); + } +} + +#[test] +fn native_api_v2_stream_open_reports_validation_setup_and_provider_failures() { + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + ptr::null(), + ptr::null(), + ptr::null(), + record_typed_llm_stream_open, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let dispatch = test_v2_dispatch_json(); + let stream_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(empty_llm_stream_next()), + runtime.handle().clone(), + None, + )); + let stream_next_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + stream_next_ref, + dispatch, + ptr::null(), + record_typed_llm_stream_open, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let (output, receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let wrong_kind = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(empty_llm_next()), + runtime.handle().clone(), + None, + )); + let wrong_kind_ref = Arc::into_raw(wrong_kind) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + wrong_kind_ref, + dispatch, + output_ref, + record_typed_llm_stream_open, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + let malformed = native_string("not-json"); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + stream_next_ref, + malformed, + output_ref, + record_typed_llm_stream_open, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + + output.settled.store(true, Ordering::Release); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + stream_next_ref, + dispatch, + output_ref, + record_typed_llm_stream_open, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + drop(receiver); + + for failure in [ + FlowError::Upstream(crate::error::UpstreamFailure { + status: Some(429), + body: "rate limited".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::RetryableStatus, + }), + FlowError::InvalidArgument("bad stream request".into()), + ] { + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(move |_| { + let failure = failure.clone(); + Box::pin(async move { Err(failure) }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (output, receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let open_receiver = start_typed_llm_stream(next_ref, dispatch, output_ref); + assert!(runtime.block_on(open_receiver).unwrap().is_err()); + drop(receiver); + unsafe { + native_async_next_release(next_ref); + native_async_stream_release(output_ref); + } + } + + let panicking = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_| { + Box::pin(async { + panic!("stream setup panic"); + }) + })), + runtime.handle().clone(), + None, + )); + let panicking_ref = Arc::into_raw(panicking) as *const NemoRelayNativeAsyncNext; + let (panic_output, panic_receiver) = test_native_output_stream(1); + let panic_output_ref = + Arc::into_raw(Arc::clone(&panic_output)) as *const NemoRelayNativeAsyncStream; + let open_receiver = start_typed_llm_stream(panicking_ref, dispatch, panic_output_ref); + assert!(matches!( + runtime.block_on(open_receiver).unwrap(), + Err(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + .. + }) + )); + drop(panic_receiver); + + unsafe { + native_string_free(dispatch); + native_string_free(malformed); + native_async_next_release(stream_next_ref); + native_async_next_release(wrong_kind_ref); + native_async_stream_release(output_ref); + native_async_next_release(panicking_ref); + native_async_stream_release(panic_output_ref); + } +} + +#[test] +fn native_api_v2_stream_forwarding_validates_handles_and_contains_panics() { + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + ptr::null(), + ptr::null(), + ptr::null(), + record_native_forward_terminal, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(empty_llm_stream_next()), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let request = + native_string_from_json(&serde_json::to_value(test_llm_request()).unwrap()).unwrap(); + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + request, + ptr::null(), + record_native_forward_terminal, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let (output, receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let wrong_kind = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(empty_llm_next()), + runtime.handle().clone(), + None, + )); + let wrong_kind_ref = Arc::into_raw(wrong_kind) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + wrong_kind_ref, + request, + output_ref, + record_native_forward_terminal, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + let malformed = native_string("not-json"); + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + malformed, + output_ref, + record_native_forward_terminal, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + output.cancelled.store(true, Ordering::Release); + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + next_ref, + request, + output_ref, + record_native_forward_terminal, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + drop(receiver); + + let (settled, _receiver) = test_native_output_stream(1); + finish_forwarded_native_stream(&settled); + finish_forwarded_native_stream(&settled); + assert!( + runtime + .block_on(push_forwarded_native_stream_chunk(&settled, Json::Null)) + .is_err() + ); + runtime.block_on(reject_forwarded_native_stream( + &settled, + FlowError::Internal("late".into()), + )); + + let panicking = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(futures_util::stream::once(async { + panic!("forwarded stream panic"); + #[allow(unreachable_code)] + Ok(Json::Null) + }))) + }) + })), + runtime.handle().clone(), + None, + )); + let panicking_ref = Arc::into_raw(panicking) as *const NemoRelayNativeAsyncNext; + let (panic_output, mut panic_receiver) = test_native_output_stream(1); + let panic_output_ref = + Arc::into_raw(Arc::clone(&panic_output)) as *const NemoRelayNativeAsyncStream; + let terminal = Arc::new(NativeForwardTerminalState::default()); + assert_eq!( + unsafe { + native_async_llm_next_forward_stream_v2( + panicking_ref, + request, + panic_output_ref, + record_native_forward_terminal, + Arc::as_ptr(&terminal).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + let error = runtime + .block_on(async { + tokio::time::timeout(Duration::from_secs(1), panic_receiver.recv()) + .await + .expect("panicking stream should settle") + .expect("panicking stream should emit a failure") + }) + .unwrap_err(); + assert!(error.to_string().contains("forwarded stream panic")); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), terminal.notified.notified()) + .await + .expect("terminal callback should be delivered"); + }); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + unsafe { + native_string_free(request); + native_string_free(malformed); + native_async_next_release(next_ref); + native_async_next_release(wrong_kind_ref); + native_async_stream_release(output_ref); + native_async_next_release(panicking_ref); + native_async_stream_release(panic_output_ref); + } +} + +#[test] +fn native_api_v2_provider_stream_reports_terminal_and_release_states() { + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + ptr::null(), + record_typed_llm_stream_next, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + unsafe { native_async_llm_stream_release_v2(ptr::null()) }; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (stream, _output_receiver) = + test_native_provider_stream(&runtime, LlmJsonStream::new(tokio_stream::empty())); + let stream_ref = Arc::into_raw(stream) as *const NemoRelayNativeLlmStreamV2; + let receiver = pull_typed_llm_stream(stream_ref); + assert!(matches!( + runtime.block_on(receiver).unwrap(), + LlmContinuationStreamEventV2::Done + )); + let callback_state = Box::into_raw(Box::new( + tokio::sync::oneshot::channel::().0, + )); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + stream_ref, + record_typed_llm_stream_next, + callback_state.cast(), + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("terminal"); + unsafe { + drop(Box::from_raw(callback_state)); + native_async_llm_stream_release_v2(stream_ref); + } + + struct StreamDropSignal(std::sync::mpsc::Sender<()>); + + impl Drop for StreamDropSignal { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + + let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); + let drop_signal = StreamDropSignal(dropped_tx); + let never_polled = LlmJsonStream::new(futures_util::stream::poll_fn(move |_| { + let _ = &drop_signal; + std::task::Poll::Pending + })); + let (stream, _output_receiver) = test_native_provider_stream(&runtime, never_polled); + let stream_ref = Arc::into_raw(stream) as *const NemoRelayNativeLlmStreamV2; + unsafe { native_async_llm_stream_release_v2(stream_ref) }; + dropped_rx + .recv_timeout(Duration::from_secs(1)) + .expect("release before first next should drop the provider stream"); +} + +#[test] +fn native_legacy_stream_next_reports_validation_item_and_setup_failures() { + assert_eq!( + unsafe { + native_async_next_invoke_stream( + ptr::null(), + ptr::null(), + ptr::null(), + accept_native_stream_item, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let request = + native_string_from_json(&serde_json::to_value(test_llm_request()).unwrap()).unwrap(); + let stream_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(empty_llm_stream_next()), + runtime.handle().clone(), + None, + )); + let stream_next_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_next_invoke_stream( + stream_next_ref, + request, + ptr::null(), + accept_native_stream_item, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + + let (output, receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let wrong_kind = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(empty_llm_next()), + runtime.handle().clone(), + None, + )); + let wrong_kind_ref = Arc::into_raw(wrong_kind) as *const NemoRelayNativeAsyncNext; + assert_eq!( + unsafe { + native_async_next_invoke_stream( + wrong_kind_ref, + request, + output_ref, + accept_native_stream_item, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + let wrong_shape = native_string_from_json(&json!({"not": "an llm request"})).unwrap(); + assert_eq!( + unsafe { + native_async_next_invoke_stream( + stream_next_ref, + wrong_shape, + output_ref, + accept_native_stream_item, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + drop(receiver); + + for next_fn in [ + Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Err( + FlowError::Internal("stream item failed".into()), + )]))) + }) as Pin> + Send>> + }) as LlmStreamExecutionNextFn, + Arc::new(|_| { + Box::pin(async { Err(FlowError::Internal("stream setup failed".into())) }) + as Pin> + Send>> + }), + ] { + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(next_fn), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (output, receiver) = test_native_output_stream(1); + let output_ref = Arc::into_raw(Arc::clone(&output)) as *const NemoRelayNativeAsyncStream; + let callback = Arc::new(NativeStreamCallbackState::default()); + assert_eq!( + unsafe { + native_async_next_invoke_stream( + next_ref, + request, + output_ref, + record_native_stream_result, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("legacy stream failure callback should run"); + }); + assert!( + callback + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_deref() + .is_some_and(|error| error.contains("stream")) + ); + drop(receiver); + unsafe { + native_async_next_release(next_ref); + native_async_stream_release(output_ref); + } + } + + let allocation_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": true + }))]))) + }) + })), + runtime.handle().clone(), + None, + )); + let allocation_next_ref = Arc::into_raw(allocation_next) as *const NemoRelayNativeAsyncNext; + let (allocation_output, allocation_receiver) = test_native_output_stream(1); + let allocation_output_ref = + Arc::into_raw(Arc::clone(&allocation_output)) as *const NemoRelayNativeAsyncStream; + let callback = Arc::new(NativeStreamCallbackState::default()); + fail_native_string_allocation_after(0); + assert_eq!( + unsafe { + native_async_next_invoke_stream( + allocation_next_ref, + request, + allocation_output_ref, + record_native_stream_result, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback.notified.notified()) + .await + .expect("allocation failure should terminate the legacy stream"); + }); + assert!(!callback.done.load(Ordering::Acquire)); + assert!( + callback + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_deref() + .is_some_and(|error| error.contains("failed to serialize or allocate")) + ); + assert_eq!(callback.callbacks.load(Ordering::SeqCst), 1); + drop(allocation_receiver); + + unsafe { + native_string_free(request); + native_string_free(wrong_shape); + native_async_next_release(stream_next_ref); + native_async_next_release(wrong_kind_ref); + native_async_stream_release(output_ref); + native_async_next_release(allocation_next_ref); + native_async_stream_release(allocation_output_ref); + } +} + +#[test] +fn native_legacy_unary_result_reports_invalid_kinds_and_panics() { + assert_eq!( + unsafe { + native_async_next_invoke_result( + ptr::null(), + ptr::null(), + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let llm = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(empty_llm_next()), + runtime.handle().clone(), + None, + )); + let llm_ref = Arc::into_raw(llm) as *const NemoRelayNativeAsyncNext; + let malformed = native_string("not-json"); + assert_eq!( + unsafe { + native_async_next_invoke_result( + llm_ref, + malformed, + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + let wrong_shape = native_string_from_json(&json!({"not": "an llm request"})).unwrap(); + assert_eq!( + unsafe { + native_async_next_invoke_result( + llm_ref, + wrong_shape, + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson + ); + let stream = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(empty_llm_stream_next()), + runtime.handle().clone(), + None, + )); + let stream_ref = Arc::into_raw(stream) as *const NemoRelayNativeAsyncNext; + let valid = + native_string_from_json(&serde_json::to_value(test_llm_request()).unwrap()).unwrap(); + assert_eq!( + unsafe { + native_async_next_invoke_result( + stream_ref, + valid, + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg + ); + + let panicking = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(Arc::new(|_| { + Box::pin(async { + panic!("legacy result panic"); + }) + })), + runtime.handle().clone(), + None, + )); + let panicking_ref = Arc::into_raw(panicking) as *const NemoRelayNativeAsyncNext; + let tool_value = native_string_from_json(&Json::Null).unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::>(); + assert_eq!( + unsafe { + native_async_next_invoke_result( + panicking_ref, + tool_value, + complete_native_next_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + assert!( + runtime + .block_on(receiver) + .unwrap() + .unwrap_err() + .contains("legacy result panic") + ); + + unsafe { + native_string_free(malformed); + native_string_free(wrong_shape); + native_string_free(valid); + native_string_free(tool_value); + native_async_next_release(llm_ref); + native_async_next_release(stream_ref); + native_async_next_release(panicking_ref); + } +} + +#[test] +fn native_async_completion_reject_covers_null_invalid_and_abort_paths() { + assert_eq!( + unsafe { native_async_completion_reject(ptr::null(), ptr::null()) }, + NemoRelayStatus::NullPointer + ); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel(); + let pending = runtime.spawn(std::future::pending::<()>()); + let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(Some(pending.abort_handle())), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let invalid_utf8 = + Box::into_raw(Box::new(NativeHostString(vec![0xff]))) as *mut NemoRelayNativeString; + assert_eq!( + unsafe { native_async_completion_reject(completion_ref, invalid_utf8) }, + NemoRelayStatus::InvalidArg + ); + assert_eq!( + unsafe { native_async_completion_reject(completion_ref, ptr::null()) }, + NemoRelayStatus::Ok + ); + assert_eq!( + unsafe { native_async_completion_reject(completion_ref, ptr::null()) }, + NemoRelayStatus::InvalidArg + ); + assert!( + runtime + .block_on(receiver) + .unwrap() + .unwrap_err() + .to_string() + .contains("native async callback rejected") + ); + assert!(completion.next_abort.lock().unwrap().is_none()); + assert!(runtime.block_on(pending).unwrap_err().is_cancelled()); + + unsafe { + native_string_free(invalid_utf8); + native_async_completion_release(completion_ref); + } +} + +fn native_continuation_context_observation( + expected_stack: &ScopeStackHandle, + expected_event_uuid: uuid::Uuid, +) -> Json { + let expected_scope_uuid = expected_stack + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid; + let visible_scope_uuid = current_scope_stack() + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid; + json!({ + "scope_stack": visible_scope_uuid == expected_scope_uuid, + "active_event_uuid": active_event_uuid() == Some(expected_event_uuid), + "publication_context": crate::api::runtime::subscriber_dispatcher::publication_context::() + .is_some_and(|context| context.as_str() == "native-continuation"), + "publication_buffer": capture_nested_publication_buffer().is_some(), + "optimization_recorder": current_llm_optimization_recorder().is_some(), + }) +} + +#[test] +fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuations() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let expected_stack = create_scope_stack(); + let expected_event_uuid = uuid::Uuid::now_v7(); + let expected = json!({ + "scope_stack": true, + "active_event_uuid": true, + "publication_context": true, + "publication_buffer": true, + "optimization_recorder": true, + }); + + runtime.block_on(TASK_SCOPE_STACK.scope( + expected_stack.clone(), + with_task_publication_context( + Some(Arc::new(String::from("native-continuation"))), + scope_llm_optimization_recorder( LlmOptimizationRecorder::default(), with_active_event_uuid( expected_event_uuid, @@ -1596,6 +4765,9 @@ fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuation let unary_ref = Arc::into_raw(unary) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -1650,16 +4822,7 @@ fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuation )); let stream_next_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + let (stream, receiver) = test_native_output_stream(1); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; let invocation = native_string_from_json( @@ -1836,6 +4999,9 @@ fn native_async_next_panics_settle_unary_and_stream_errors() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -1876,16 +5042,7 @@ fn native_async_next_panics_settle_unary_and_stream_errors() { None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, _receiver) = tokio::sync::mpsc::channel(1); - let output_stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + let (output_stream, _receiver) = test_native_output_stream(1); let output_stream_ref = Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; let callback_state = Arc::new(NativeStreamCallbackState::default()); @@ -1954,6 +5111,9 @@ fn native_async_next_is_permanently_one_shot() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -2005,6 +5165,9 @@ fn cancelled_native_async_next_does_not_start_unary_or_stream_continuations() { let unary_ref = Arc::into_raw(unary) as *const NemoRelayNativeAsyncNext; let (sender, _receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(true), next_invoked: AtomicBool::new(false), @@ -2035,6 +5198,9 @@ fn cancelled_native_async_next_does_not_start_unary_or_stream_continuations() { let stream_next_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::mpsc::channel(1); let stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(true), settled: AtomicBool::new(false), @@ -2098,6 +5264,9 @@ fn malformed_llm_next_does_not_consume_the_completion() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, _receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -2143,16 +5312,7 @@ fn native_async_stream_next_supports_repeated_concurrent_calls() { None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + let (stream, receiver) = test_native_output_stream(1); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; let invocation = native_string_from_json( &serde_json::to_value(LlmRequest { @@ -2259,16 +5419,7 @@ fn native_async_stream_settlement_rejects_late_next_and_aborts_in_flight_next() None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + let (stream, receiver) = test_native_output_stream(1); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; let invocation = native_string_from_json( &serde_json::to_value(LlmRequest { @@ -2348,47 +5499,174 @@ fn native_async_stream_settlement_rejects_late_next_and_aborts_in_flight_next() runtime.block_on(tokio::task::yield_now()); assert_eq!(provider_calls.load(Ordering::SeqCst), 1); - drop(NativeAsyncStreamReceiver { - receiver, - stream: Arc::clone(&stream), - }); + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }); + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + native_async_stream_release(stream_ref); + } + } +} + +#[test] +fn native_async_stream_next_stops_callbacks_after_false() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_request| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(json!({"chunk": 1})), + Ok(json!({"chunk": 2})), + ]))) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (stream, receiver) = test_native_output_stream(1); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let invocation = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + let callbacks = AtomicUsize::new(0); + + assert_eq!( unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); - native_async_stream_release(stream_ref); + native_async_next_invoke_stream( + next_ref, + invocation, + stream_ref, + stop_after_first_native_stream_item, + (&callbacks as *const AtomicUsize).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + for _ in 0..10 { + tokio::task::yield_now().await; } + }); + assert_eq!(callbacks.load(Ordering::SeqCst), 1); + + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }); + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + native_async_stream_release(stream_ref); } } #[test] -fn native_async_stream_next_stops_callbacks_after_false() { +fn native_async_stream_in_flight_cancellation_releases_callback_state() { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(Mutex::new(Some(started_tx))); let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(Arc::new(|_request| { - Box::pin(async { - Ok(LlmJsonStream::new(tokio_stream::iter(vec![ - Ok(json!({"chunk": 1})), - Ok(json!({"chunk": 2})), - ]))) + NativeAsyncNextInner::LlmStream(Arc::new(move |_request| { + let started_tx = Arc::clone(&started_tx); + Box::pin(async move { + if let Some(started_tx) = started_tx.lock().unwrap().take() { + let _ = started_tx.send(()); + } + std::future::pending().await }) })), runtime.handle().clone(), None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, + let (stream, receiver) = test_native_output_stream(1); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let invocation = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + let callback_state = Arc::new(NativeStreamCallbackState::default()); + let drop_count = Arc::new(AtomicUsize::new(0)); + let callback_user_data = Box::into_raw(Box::new(OwnedNativeStreamCallbackState { + result: Arc::clone(&callback_state), + drop_count: Arc::clone(&drop_count), + stream: stream_ref, + })); + + assert_eq!( + unsafe { + native_async_next_invoke_stream( + next_ref, + invocation, + stream_ref, + record_and_release_native_stream_result, + callback_user_data.cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(started_rx).unwrap(); + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), }); + runtime + .block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback_state.notified.notified()).await + }) + .expect("in-flight cancellation should deliver a terminal callback"); + runtime.block_on(tokio::task::yield_now()); + assert_eq!(callback_state.callbacks.load(Ordering::SeqCst), 1); + assert_eq!(drop_count.load(Ordering::SeqCst), 1); + assert!(!callback_state.done.load(Ordering::Acquire)); + assert!( + callback_state + .error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_deref() + .is_some_and(|error| error.contains("cancelled")) + ); + assert_eq!(Arc::strong_count(&stream), 1); + + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + } +} + +#[test] +fn native_async_stream_cancellation_before_first_poll_releases_callback_state() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(move |_request| Box::pin(std::future::pending()))), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (stream, receiver) = test_native_output_stream(1); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; let invocation = native_string_from_json( &serde_json::to_value(LlmRequest { @@ -2398,141 +5676,371 @@ fn native_async_stream_next_stops_callbacks_after_false() { .unwrap(), ) .unwrap(); - let callbacks = AtomicUsize::new(0); + let callback_state = Arc::new(NativeStreamCallbackState::default()); + let drop_count = Arc::new(AtomicUsize::new(0)); + let callback_user_data = Box::into_raw(Box::new(OwnedNativeStreamCallbackState { + result: Arc::clone(&callback_state), + drop_count: Arc::clone(&drop_count), + stream: stream_ref, + })); + + assert_eq!( + unsafe { + native_async_next_invoke_stream( + next_ref, + invocation, + stream_ref, + record_and_release_native_stream_result, + callback_user_data.cast(), + ) + }, + NemoRelayStatus::Ok + ); + // The current-thread runtime has not been driven, so cancellation happens + // before the spawned continuation can be polled for the first time. + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }); + runtime + .block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback_state.notified.notified()).await + }) + .expect("pre-poll cancellation should deliver a terminal callback"); + runtime.block_on(tokio::task::yield_now()); + assert_eq!(callback_state.callbacks.load(Ordering::SeqCst), 1); + assert_eq!(drop_count.load(Ordering::SeqCst), 1); + assert!(!callback_state.done.load(Ordering::Acquire)); + assert!( + callback_state + .error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_deref() + .is_some_and(|error| error.contains("cancelled")) + ); + assert_eq!(Arc::strong_count(&stream), 1); + + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + } +} + +fn wait_for_task_free(runtime: &Runtime, frees: &AtomicUsize) { + runtime + .block_on(async { + tokio::time::timeout(Duration::from_secs(1), async { + while frees.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + }) + .expect("cooperative task state should be freed"); +} + +#[test] +fn native_async_task_pending_wake_settles_and_frees_once() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel(); + let completion = Arc::new(NativeAsyncCompletion { + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let frees = Arc::new(AtomicUsize::new(0)); + let state = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::SettleCompletionOnSecondPoll, + polls: AtomicUsize::new(0), + completion: completion_ref as usize, + stream: 0, + started: Some(started_tx), + frees: Arc::clone(&frees), + })); + assert_eq!( + unsafe { + native_async_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + + let task = started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("first poll should retain a task waker") + as *const NemoRelayNativeAsyncTaskV2; + unsafe { native_async_task_wake_v2(task) }; + unsafe { native_async_task_release_v2(task) }; + let result = runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(1), receiver).await }) + .expect("completion task should settle") + .unwrap() + .unwrap(); + assert_eq!(result, json!({"cooperative": true})); + wait_for_task_free(&runtime, &frees); + assert_eq!(frees.load(Ordering::SeqCst), 1); +} + +#[test] +fn native_async_task_cancellation_reclaims_never_woken_completion_and_stream_tasks() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + + let (sender, receiver) = tokio::sync::oneshot::channel(); + let completion = Arc::new(NativeAsyncCompletion { + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let (completion_started_tx, completion_started_rx) = std::sync::mpsc::channel(); + let completion_frees = Arc::new(AtomicUsize::new(0)); + let state = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::Pending, + polls: AtomicUsize::new(0), + completion: completion_ref as usize, + stream: 0, + started: Some(completion_started_tx), + frees: Arc::clone(&completion_frees), + })); + assert_eq!( + unsafe { + native_async_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + completion_started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("completion task should poll once"); + drop(NativeAsyncWait { + completion: Arc::clone(&completion), + receiver, + completed: false, + }); + wait_for_task_free(&runtime, &completion_frees); + let (sender, receiver) = tokio::sync::mpsc::channel(1); + let stream = Arc::new(NativeAsyncStream { + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + settled: AtomicBool::new(false), + downstream_aborts: Mutex::new(HashMap::new()), + settlement: Mutex::new(()), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let (stream_started_tx, stream_started_rx) = std::sync::mpsc::channel(); + let stream_frees = Arc::new(AtomicUsize::new(0)); + let state = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::Pending, + polls: AtomicUsize::new(0), + completion: 0, + stream: stream_ref as usize, + started: Some(stream_started_tx), + frees: Arc::clone(&stream_frees), + })); assert_eq!( unsafe { - native_async_next_invoke_stream( - next_ref, - invocation, + native_async_stream_spawn_task_v2( stream_ref, - stop_after_first_native_stream_item, - (&callbacks as *const AtomicUsize).cast_mut().cast(), + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), ) }, NemoRelayStatus::Ok ); - runtime.block_on(async { - for _ in 0..10 { - tokio::task::yield_now().await; - } - }); - assert_eq!(callbacks.load(Ordering::SeqCst), 1); - + stream_started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("stream task should poll once"); drop(NativeAsyncStreamReceiver { receiver, stream: Arc::clone(&stream), }); - unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); - native_async_stream_release(stream_ref); - } + wait_for_task_free(&runtime, &stream_frees); + assert_eq!(completion_frees.load(Ordering::SeqCst), 1); + assert_eq!(stream_frees.load(Ordering::SeqCst), 1); } #[test] -fn native_async_stream_in_flight_cancellation_releases_callback_state() { - let runtime = tokio::runtime::Builder::new_current_thread() +fn native_async_task_duplicate_spawn_does_not_consume_user_data() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) .enable_all() .build() .unwrap(); - let (started_tx, started_rx) = tokio::sync::oneshot::channel(); - let started_tx = Arc::new(Mutex::new(Some(started_tx))); - let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(Arc::new(move |_request| { - let started_tx = Arc::clone(&started_tx); - Box::pin(async move { - if let Some(started_tx) = started_tx.lock().unwrap().take() { - let _ = started_tx.send(()); - } - std::future::pending().await - }) - })), - runtime.handle().clone(), - None, - )); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { + let (sender, receiver) = tokio::sync::oneshot::channel(); + let completion = Arc::new(NativeAsyncCompletion { sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); - let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; - let invocation = native_string_from_json( - &serde_json::to_value(LlmRequest { - headers: Map::new(), - content: json!({"stream": true}), - }) - .unwrap(), - ) - .unwrap(); - let callback_state = Arc::new(NativeStreamCallbackState::default()); - let drop_count = Arc::new(AtomicUsize::new(0)); - let callback_user_data = Box::into_raw(Box::new(OwnedNativeStreamCallbackState { - result: Arc::clone(&callback_state), - drop_count: Arc::clone(&drop_count), - stream: stream_ref, + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let first_frees = Arc::new(AtomicUsize::new(0)); + let first = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::Pending, + polls: AtomicUsize::new(0), + completion: completion_ref as usize, + stream: 0, + started: Some(started_tx), + frees: Arc::clone(&first_frees), })); - assert_eq!( unsafe { - native_async_next_invoke_stream( - next_ref, - invocation, - stream_ref, - record_and_release_native_stream_result, - callback_user_data.cast(), + native_async_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + first.cast(), + Some(free_cooperative_task), ) }, NemoRelayStatus::Ok ); - runtime.block_on(started_rx).unwrap(); - drop(NativeAsyncStreamReceiver { + started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("first task should be active"); + + let duplicate_frees = Arc::new(AtomicUsize::new(0)); + let duplicate = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::Pending, + polls: AtomicUsize::new(0), + completion: 0, + stream: 0, + started: None, + frees: Arc::clone(&duplicate_frees), + })); + assert_eq!( + unsafe { + native_async_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + duplicate.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_eq!(duplicate_frees.load(Ordering::SeqCst), 0); + unsafe { free_cooperative_task(duplicate.cast()) }; + assert_eq!(duplicate_frees.load(Ordering::SeqCst), 1); + drop(NativeAsyncWait { + completion: Arc::clone(&completion), receiver, - stream: Arc::clone(&stream), + completed: false, }); - runtime - .block_on(async { - tokio::time::timeout(Duration::from_secs(1), callback_state.notified.notified()).await - }) - .expect("in-flight cancellation should deliver a terminal callback"); - runtime.block_on(tokio::task::yield_now()); - assert_eq!(callback_state.callbacks.load(Ordering::SeqCst), 1); - assert_eq!(drop_count.load(Ordering::SeqCst), 1); - assert!(!callback_state.done.load(Ordering::Acquire)); - assert!( - callback_state - .error - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .as_deref() - .is_some_and(|error| error.contains("cancelled")) - ); - assert_eq!(Arc::strong_count(&stream), 1); + wait_for_task_free(&runtime, &first_frees); +} - unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); +#[test] +fn native_async_task_contract_errors_settle_completion_as_internal() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + for (mode, expected) in [ + ( + CooperativeTaskMode::CompleteWithoutSettlement, + "Complete without settling", + ), + (CooperativeTaskMode::InvalidState, "invalid callback state"), + ] { + let (sender, receiver) = tokio::sync::oneshot::channel(); + let completion = Arc::new(NativeAsyncCompletion { + sender: Mutex::new(Some(sender)), + cancelled: AtomicBool::new(false), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let frees = Arc::new(AtomicUsize::new(0)); + let state = Box::into_raw(Box::new(CooperativeTaskState { + mode, + polls: AtomicUsize::new(0), + completion: completion_ref as usize, + stream: 0, + started: None, + frees: Arc::clone(&frees), + })); + assert_eq!( + unsafe { + native_async_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + let error = runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(1), receiver).await }) + .expect("contract error should settle completion") + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + wait_for_task_free(&runtime, &frees); } } #[test] -fn native_async_stream_cancellation_before_first_poll_releases_callback_state() { - let runtime = tokio::runtime::Builder::new_current_thread() +fn native_async_task_stream_backpressure_wakes_after_consumer_drain() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) .enable_all() .build() .unwrap(); - let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(Arc::new(move |_request| Box::pin(std::future::pending()))), - runtime.handle().clone(), - None, - )); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::mpsc::channel(1); let stream = Arc::new(NativeAsyncStream { sender: Mutex::new(Some(sender)), @@ -2540,67 +6048,57 @@ fn native_async_stream_cancellation_before_first_poll_releases_callback_state() settled: AtomicBool::new(false), downstream_aborts: Mutex::new(HashMap::new()), settlement: Mutex::new(()), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; - let invocation = native_string_from_json( - &serde_json::to_value(LlmRequest { - headers: Map::new(), - content: json!({"stream": true}), - }) - .unwrap(), - ) - .unwrap(); - let callback_state = Arc::new(NativeStreamCallbackState::default()); - let drop_count = Arc::new(AtomicUsize::new(0)); - let callback_user_data = Box::into_raw(Box::new(OwnedNativeStreamCallbackState { - result: Arc::clone(&callback_state), - drop_count: Arc::clone(&drop_count), - stream: stream_ref, + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let frees = Arc::new(AtomicUsize::new(0)); + let state = Box::into_raw(Box::new(CooperativeTaskState { + mode: CooperativeTaskMode::PushTwoChunks, + polls: AtomicUsize::new(0), + completion: 0, + stream: stream_ref as usize, + started: Some(started_tx), + frees: Arc::clone(&frees), })); - assert_eq!( unsafe { - native_async_next_invoke_stream( - next_ref, - invocation, + native_async_stream_spawn_task_v2( stream_ref, - record_and_release_native_stream_result, - callback_user_data.cast(), + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), ) }, NemoRelayStatus::Ok ); - // The current-thread runtime has not been driven, so cancellation happens - // before the spawned continuation can be polled for the first time. - drop(NativeAsyncStreamReceiver { + started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("producer should observe bounded backpressure"); + let mut output = NativeAsyncStreamReceiver { receiver, stream: Arc::clone(&stream), + }; + let chunks = runtime.block_on(async { + let first = tokio::time::timeout(Duration::from_secs(1), output.next()) + .await + .expect("first chunk should be available") + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(1), output.next()) + .await + .expect("draining should wake the pending producer") + .unwrap() + .unwrap(); + assert!(output.next().await.is_none()); + [first, second] }); - runtime - .block_on(async { - tokio::time::timeout(Duration::from_secs(1), callback_state.notified.notified()).await - }) - .expect("pre-poll cancellation should deliver a terminal callback"); - runtime.block_on(tokio::task::yield_now()); - assert_eq!(callback_state.callbacks.load(Ordering::SeqCst), 1); - assert_eq!(drop_count.load(Ordering::SeqCst), 1); - assert!(!callback_state.done.load(Ordering::Acquire)); - assert!( - callback_state - .error - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .as_deref() - .is_some_and(|error| error.contains("cancelled")) - ); - assert_eq!(Arc::strong_count(&stream), 1); - - unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); - } + assert_eq!(chunks, [json!({"chunk": 1}), json!({"chunk": 2})]); + wait_for_task_free(&runtime, &frees); } #[test] @@ -2611,6 +6109,9 @@ fn native_async_completion_abi_rejects_invalid_duplicate_and_cancelled_settlemen .unwrap(); let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -2646,6 +6147,9 @@ fn native_async_completion_abi_rejects_invalid_duplicate_and_cancelled_settlemen let (sender, _receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(true), next_invoked: AtomicBool::new(false), @@ -2672,6 +6176,9 @@ fn completed_native_async_wait_is_not_marked_cancelled() { .unwrap(); let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -2726,6 +6233,9 @@ fn native_async_completion_cancellation_wins_resolve_and_reject_settlement_races let (sender, _receiver) = tokio::sync::oneshot::channel(); let settlement_checkpoint = Arc::new(std::sync::Barrier::new(2)); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -2807,6 +6317,9 @@ fn cancelling_completion_aborts_pending_native_next() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::oneshot::channel(); let completion = Arc::new(NativeAsyncCompletion { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), next_invoked: AtomicBool::new(false), @@ -3081,6 +6594,9 @@ fn native_async_stream_settlement_cannot_succeed_after_cancellation() { let (sender, receiver) = tokio::sync::mpsc::channel(1); let settlement_checkpoint = Arc::new(std::sync::Barrier::new(2)); let stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -3142,33 +6658,79 @@ fn native_async_stream_settlement_cannot_succeed_after_cancellation() { } #[test] -fn native_async_stream_push_is_bounded_retryable_and_incremental() { +fn native_async_stream_backpressure_preserves_v3_internal_and_v4_would_block() { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); - let (sender, receiver) = tokio::sync::mpsc::channel(1); - let stream = Arc::new(NativeAsyncStream { - sender: Mutex::new(Some(sender)), - cancelled: AtomicBool::new(false), - settled: AtomicBool::new(false), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + let (stream, receiver) = test_native_output_stream(1); let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; let first_chunk = native_string(r#"{"chunk":1}"#); let second_chunk = native_string(r#"{"chunk":2}"#); + let host_v3 = build_native_host_api_v3(); + let host_v4 = build_native_host_api_v4(); + assert_eq!( + unsafe { (host_v3.async_stream_push_json)(stream_ref, first_chunk) }, + NemoRelayStatus::Ok + ); + assert_eq!( + unsafe { (host_v3.async_stream_push_json)(stream_ref, second_chunk) }, + NemoRelayStatus::Internal + ); + assert_last_error_contains("backpressured"); + assert_eq!( + unsafe { (host_v4.v3.async_stream_push_json)(stream_ref, second_chunk) }, + NemoRelayStatus::WouldBlock + ); + let mut receiver = NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }; + assert_eq!( + runtime.block_on(receiver.next()).unwrap().unwrap(), + json!({"chunk": 1}) + ); + assert_eq!( + unsafe { (host_v4.v3.async_stream_push_json)(stream_ref, second_chunk) }, + NemoRelayStatus::Ok + ); + assert_eq!( + runtime.block_on(receiver.next()).unwrap().unwrap(), + json!({"chunk": 2}) + ); + assert_eq!( + unsafe { native_async_stream_finish(stream_ref) }, + NemoRelayStatus::Ok + ); + assert!(runtime.block_on(receiver.next()).is_none()); + drop(receiver); + assert!(unsafe { native_async_stream_is_cancelled(stream_ref) }); assert_eq!( unsafe { native_async_stream_push_json(stream_ref, first_chunk) }, + NemoRelayStatus::InvalidArg + ); + unsafe { + native_string_free(first_chunk); + native_string_free(second_chunk); + native_async_stream_release(stream_ref); + } + + let (stream, receiver) = test_native_output_stream(1); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + let chunk = native_string(r#"{"chunk":1}"#); + let message = native_string("provider failed"); + assert_eq!( + unsafe { (host_v3.async_stream_push_json)(stream_ref, chunk) }, NemoRelayStatus::Ok ); assert_eq!( - unsafe { native_async_stream_push_json(stream_ref, second_chunk) }, + unsafe { (host_v3.async_stream_reject)(stream_ref, message) }, NemoRelayStatus::Internal ); - assert_last_error_contains("backpressured"); + assert_eq!( + unsafe { (host_v4.v3.async_stream_reject)(stream_ref, message) }, + NemoRelayStatus::WouldBlock + ); let mut receiver = NativeAsyncStreamReceiver { receiver, stream: Arc::clone(&stream), @@ -3178,27 +6740,22 @@ fn native_async_stream_push_is_bounded_retryable_and_incremental() { json!({"chunk": 1}) ); assert_eq!( - unsafe { native_async_stream_push_json(stream_ref, second_chunk) }, + unsafe { (host_v4.v3.async_stream_reject)(stream_ref, message) }, NemoRelayStatus::Ok ); - assert_eq!( - runtime.block_on(receiver.next()).unwrap().unwrap(), - json!({"chunk": 2}) - ); - assert_eq!( - unsafe { native_async_stream_finish(stream_ref) }, - NemoRelayStatus::Ok + assert!( + runtime + .block_on(receiver.next()) + .expect("rejection should be emitted") + .unwrap_err() + .to_string() + .contains("provider failed") ); assert!(runtime.block_on(receiver.next()).is_none()); drop(receiver); - assert!(unsafe { native_async_stream_is_cancelled(stream_ref) }); - assert_eq!( - unsafe { native_async_stream_push_json(stream_ref, first_chunk) }, - NemoRelayStatus::InvalidArg - ); unsafe { - native_string_free(first_chunk); - native_string_free(second_chunk); + native_string_free(chunk); + native_string_free(message); native_async_stream_release(stream_ref); } } @@ -3912,6 +7469,37 @@ fn native_registration_entrypoints_reject_null_contexts() { ), NemoRelayStatus::NullPointer ); + assert_eq!( + native_async_completion_spawn_task_v2( + ptr::null(), + pending_native_task, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_async_stream_spawn_task_v2( + ptr::null(), + pending_native_task, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + // Task-spawn validation clears its own last-error slot. End on a + // registration entry point so this test's shared diagnostic assertion + // still verifies the null plugin-context contract. + assert_eq!( + native_plugin_context_register_subscriber( + ptr::null_mut(), + ptr::null(), + noop_subscriber, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); } assert_last_error_contains("plugin context is null"); } @@ -3924,7 +7512,8 @@ fn native_registration_entrypoints_reject_invalid_host_contexts_and_names() { relay_compat: "^0.7".into(), allows_multiple_components: false, plugin: Mutex::new(NemoRelayNativePluginV1::default()), - _library: libloading::os::unix::Library::this().into(), + library: Some(libloading::os::unix::Library::this().into()), + retain_library_on_drop: AtomicBool::new(false), }); let mut invalid_host = NativeHostPluginContext { ctx: ptr::null_mut(), @@ -4381,7 +7970,8 @@ fn assert_async_request_registration_rejects_legacy_relay_contract() { relay_compat: "^0.5".into(), allows_multiple_components: false, plugin: Mutex::new(NemoRelayNativePluginV1::default()), - _library: libloading::os::unix::Library::this().into(), + library: Some(libloading::os::unix::Library::this().into()), + retain_library_on_drop: AtomicBool::new(false), }); let mut registration = PluginRegistrationContext::new(); let mut host = NativeHostPluginContext { @@ -4430,7 +8020,8 @@ async fn native_async_wrappers_validate_callback_result_shapes() { relay_compat: "^0.7".into(), allows_multiple_components: false, plugin: Mutex::new(NemoRelayNativePluginV1::default()), - _library: libloading::os::unix::Library::this().into(), + library: Some(libloading::os::unix::Library::this().into()), + retain_library_on_drop: AtomicBool::new(false), }); let result = native_string("true"); let user_data = result.cast(); @@ -4773,45 +8364,50 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { }, NemoRelayStatus::NullPointer ); + assert_last_error_contains("request codec decode output pointer is null"); + + let mut output = ptr::null_mut(); + assert_eq!( + unsafe { native_llm_request_codec_decode(ptr::null(), request_json, &mut output) }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode capability is null"); assert_eq!( unsafe { native_llm_request_codec_encode( - ptr::null(), + ptr::from_ref(&request_codec).cast(), annotated_json, request_json, - &mut ptr::null_mut(), + ptr::null_mut(), ) }, NemoRelayStatus::NullPointer ); + assert_last_error_contains("request codec encode output pointer is null"); + assert_eq!( unsafe { - native_llm_request_codec_encode( - ptr::from_ref(&request_codec).cast(), - annotated_json, - ptr::null(), - &mut ptr::null_mut(), - ) + native_llm_request_codec_encode(ptr::null(), annotated_json, request_json, &mut output) }, NemoRelayStatus::NullPointer ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode capability is null"); + assert_eq!( unsafe { native_llm_request_codec_encode( ptr::from_ref(&request_codec).cast(), annotated_json, - request_json, - ptr::null_mut(), + ptr::null(), + &mut output, ) }, NemoRelayStatus::NullPointer ); - assert_eq!( - unsafe { - native_llm_response_codec_decode(ptr::null(), request_json, &mut ptr::null_mut()) - }, - NemoRelayStatus::NullPointer - ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode original request is null"); assert_eq!( unsafe { native_llm_response_codec_decode( @@ -4822,14 +8418,34 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { }, NemoRelayStatus::NullPointer ); + assert_last_error_contains("response codec decode output pointer is null"); + + assert_eq!( + unsafe { native_llm_response_codec_decode(ptr::null(), request_json, &mut output) }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode capability is null"); + + assert_null_codec_inputs_clear_existing_outputs(&request_codec, &response_codec, request_json); + unsafe { + native_string_free(annotated_json); + native_string_free(request_json); + } +} +fn assert_null_codec_inputs_clear_existing_outputs( + request_codec: &NativeHostLlmRequestCodec, + response_codec: &NativeHostLlmResponseCodec, + request_json: *mut NemoRelayNativeString, +) { let request_decode_sentinel = native_string("request-decode-sentinel"); let mut output = request_decode_sentinel; set_native_last_error("stale request decode error"); assert_eq!( unsafe { native_llm_request_codec_decode( - ptr::from_ref(&request_codec).cast(), + ptr::from_ref(request_codec).cast(), ptr::null(), &mut output, ) @@ -4846,7 +8462,7 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { assert_eq!( unsafe { native_llm_request_codec_encode( - ptr::from_ref(&request_codec).cast(), + ptr::from_ref(request_codec).cast(), ptr::null(), request_json, &mut output, @@ -4864,7 +8480,7 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { assert_eq!( unsafe { native_llm_response_codec_decode( - ptr::from_ref(&response_codec).cast(), + ptr::from_ref(response_codec).cast(), ptr::null(), &mut output, ) @@ -4873,11 +8489,7 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { ); assert!(output.is_null()); assert_last_error_contains("response codec decode response is null"); - unsafe { - native_string_free(response_decode_sentinel); - native_string_free(annotated_json); - native_string_free(request_json); - } + unsafe { native_string_free(response_decode_sentinel) }; } unsafe extern "C" fn tool_json_echo( @@ -5023,6 +8635,16 @@ unsafe extern "C" fn llm_stream_execution_error( NemoRelayStatus::InvalidArg } +unsafe extern "C" fn event_sanitize_error_with_output( + _user_data: *mut c_void, + _event_json: *const NemoRelayNativeString, + _fields_json: *const NemoRelayNativeString, + out_fields_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_fields_json = native_string(r#"{"discarded":true}"#) }; + NemoRelayStatus::InvalidArg +} + unsafe extern "C" fn llm_request_echo( _user_data: *mut c_void, request_json: *const NemoRelayNativeString, @@ -5044,6 +8666,25 @@ unsafe extern "C" fn llm_request_alias( NemoRelayStatus::Ok } +unsafe extern "C" fn llm_request_error_with_output( + _user_data: *mut c_void, + _request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, + out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_request_json = native_string(r#"{"discarded":true}"#) }; + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn llm_request_none( + _user_data: *mut c_void, + _request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, + _out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} + unsafe extern "C" fn llm_request_codec_round_trip( _user_data: *mut c_void, request_json: *const NemoRelayNativeString, @@ -5093,6 +8734,40 @@ unsafe extern "C" fn llm_response_alias( NemoRelayStatus::Ok } +#[cfg(unix)] +unsafe extern "C" fn llm_response_error_with_output( + _user_data: *mut c_void, + _response_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, + out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_response_json = native_string(r#"{"discarded":true}"#) }; + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn llm_response_none( + _user_data: *mut c_void, + _response_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, + _out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} + +#[cfg(unix)] +unsafe extern "C" fn llm_execution_error_with_output( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _request_json: *const NemoRelayNativeString, + _next_fn: NemoRelayNativeLlmNextFn, + _next_ctx: *mut c_void, + out_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_json = native_string(r#"{"discarded":true}"#) }; + NemoRelayStatus::InvalidArg +} + #[test] fn native_callback_helpers_cover_success_error_and_invalid_output() { assert_eq!( @@ -5106,6 +8781,37 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { .contains("tool callback rejected input") ); + let event = Event::Mark(crate::api::event::MarkEvent { + base: crate::api::event::BaseEvent::builder() + .name("native-test") + .build(), + category: None, + category_profile: None, + }); + let fields = EventSanitizeFields::default(); + let live_before = native_string_live_allocations(); + fail_native_string_allocation_after(1); + assert!( + call_event_sanitize_callback( + event_sanitize_error_with_output, + ptr::null_mut(), + &event, + &fields, + ) + .is_err() + ); + assert_eq!(native_string_live_allocations(), live_before); + assert!( + call_event_sanitize_callback( + event_sanitize_error_with_output, + ptr::null_mut(), + &event, + &fields, + ) + .is_err() + ); + assert_eq!(native_string_live_allocations(), live_before); + let request = LlmRequest { headers: Map::new(), content: json!({"model": "test"}), @@ -5135,6 +8841,25 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { .unwrap(), Some(request.clone()) ); + assert!( + call_llm_sanitize_request_callback( + llm_request_error_with_output, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .is_err() + ); + assert_eq!( + call_llm_sanitize_request_callback( + llm_request_none, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + None + ); let response = json!({"message": "alias"}); assert_eq!( @@ -5200,7 +8925,8 @@ async fn native_callback_wrappers_release_error_outputs_and_preserve_reasons() { relay_compat: "^0.7".into(), allows_multiple_components: false, plugin: Mutex::new(NemoRelayNativePluginV1::default()), - _library: libloading::os::unix::Library::this().into(), + library: Some(libloading::os::unix::Library::this().into()), + retain_library_on_drop: AtomicBool::new(false), }); let request = LlmRequest { headers: Map::new(), @@ -5309,6 +9035,47 @@ async fn native_callback_wrappers_release_error_outputs_and_preserve_reasons() { .to_string() .contains("LLM stream execution failed") ); + assert!( + call_llm_sanitize_response_callback( + llm_response_error_with_output, + ptr::null_mut(), + &json!({"message": "discarded"}), + LlmSanitizeResponseContext::default(), + ) + .is_err() + ); + assert_eq!( + call_llm_sanitize_response_callback( + llm_response_none, + ptr::null_mut(), + &json!({"message": "none"}), + LlmSanitizeResponseContext::default(), + ) + .unwrap(), + None + ); + + let user_data = test_v2_callback_user_data(ptr::null_mut()); + assert!( + call_llm_execution_callback( + llm_execution_error_with_output, + &user_data, + "llm", + &test_llm_request(), + empty_llm_next(), + ) + .is_err() + ); + assert!( + call_llm_stream_execution_callback( + llm_stream_execution_error, + user_data, + "llm", + &test_llm_request(), + empty_llm_stream_next(), + ) + .is_err() + ); } #[test] @@ -5456,6 +9223,10 @@ fn native_llm_sanitizer_input_allocation_failures_release_codec_ids() { let identity = LlmCodecIdentity::Runtime("com.example.chat.v1".into()); let live_before = native_string_live_allocations(); + fail_native_string_allocation_after(0); + assert!(native_llm_codec_identity(&identity).is_err()); + assert_eq!(native_string_live_allocations(), live_before); + fail_native_string_allocation_after(1); let request_error = call_llm_sanitize_request_callback( llm_request_alias, @@ -5532,9 +9303,21 @@ fn native_non_streaming_continuations_cover_success_and_error_paths() { ); unsafe { drop(Box::from_raw(next as *mut ToolExecutionNextFn)) }; - let panicking_next: ToolExecutionNextFn = - Arc::new(|_| Box::pin(async { panic!("tool next panic") })); - let next = Box::into_raw(Box::new(panicking_next)) as *mut c_void; + let invalid = native_string("not-json"); + let next = Box::into_raw(Box::new(tool_next(Ok(Json::Null)))) as *mut c_void; + out = ptr::null_mut(); + assert_eq!( + unsafe { native_tool_next(invalid, next, &mut out) }, + NemoRelayStatus::InvalidJson + ); + unsafe { + drop(Box::from_raw(next as *mut ToolExecutionNextFn)); + native_string_free(invalid); + } + + let next: ToolExecutionNextFn = Arc::new(|_| Box::pin(async { panic!("tool next panic") })); + let next = Box::into_raw(Box::new(next)) as *mut c_void; + out = ptr::null_mut(); assert_eq!( unsafe { native_tool_next(args, next, &mut out) }, NemoRelayStatus::Internal @@ -5548,6 +9331,10 @@ fn native_non_streaming_continuations_cover_success_and_error_paths() { content: json!({"model": "test"}), }; let request_json = native_string_from_json(&serde_json::to_value(&request).unwrap()).unwrap(); + assert_eq!( + unsafe { native_llm_next(request_json, ptr::null_mut(), &mut out) }, + NemoRelayStatus::NullPointer + ); let next = Box::into_raw(Box::new(llm_next(Ok(json!({"answer": 42}))))) as *mut c_void; out = ptr::null_mut(); assert_eq!( @@ -5570,9 +9357,21 @@ fn native_non_streaming_continuations_cover_success_and_error_paths() { ); unsafe { drop(Box::from_raw(next as *mut LlmExecutionNextFn)) }; - let panicking_next: LlmExecutionNextFn = - Arc::new(|_| Box::pin(async { panic!("LLM next panic") })); - let next = Box::into_raw(Box::new(panicking_next)) as *mut c_void; + let invalid = native_string("not-json"); + let next = Box::into_raw(Box::new(llm_next(Ok(Json::Null)))) as *mut c_void; + out = ptr::null_mut(); + assert_eq!( + unsafe { native_llm_next(invalid, next, &mut out) }, + NemoRelayStatus::InvalidJson + ); + unsafe { + drop(Box::from_raw(next as *mut LlmExecutionNextFn)); + native_string_free(invalid); + } + + let next: LlmExecutionNextFn = Arc::new(|_| Box::pin(async { panic!("LLM next panic") })); + let next = Box::into_raw(Box::new(next)) as *mut c_void; + out = ptr::null_mut(); assert_eq!( unsafe { native_llm_next(request_json, next, &mut out) }, NemoRelayStatus::Internal @@ -5698,6 +9497,11 @@ async fn native_stream_adapter_covers_chunks_end_errors_and_cancellation() { assert_eq!(drop_count.load(Ordering::SeqCst), 1); } + let (raw, _, drop_count) = test_native_stream([NativeStreamItem::EndWithJson]); + let mut stream = native_stream_to_relay_stream(raw, None, None).unwrap(); + assert!(stream.next().await.is_none()); + assert_eq!(drop_count.load(Ordering::SeqCst), 1); + let (mut raw, _, drop_count) = test_native_stream([]); raw.struct_size = 0; assert!(NativeRelayLlmStream::from_raw(raw, None, None).is_err()); @@ -5708,19 +9512,17 @@ async fn native_stream_adapter_covers_chunks_end_errors_and_cancellation() { assert!(NativeRelayLlmStream::from_raw(raw, None, None).is_err()); assert_eq!(drop_count.load(Ordering::SeqCst), 1); - let (raw, _, drop_count) = test_native_stream([NativeStreamItem::EndWithJson]); - let mut stream = native_stream_to_relay_stream(raw, None, None).unwrap(); - assert!(stream.next().await.is_none()); - assert_eq!(drop_count.load(Ordering::SeqCst), 1); - - let mut invalid = NativeRelayLlmStream { - raw: NemoRelayNativeLlmStreamV1::default(), + let (mut raw, _, drop_count) = test_native_stream([]); + raw.next = None; + let mut stream = NativeRelayLlmStream { + raw, finished: false, _next_ctx: None, _callback_user_data: None, }; - assert!(invalid.next().await.unwrap().is_err()); - assert!(invalid.next().await.is_none()); + assert!(stream.next().await.unwrap().is_err()); + assert!(stream.next().await.is_none()); + assert_eq!(drop_count.load(Ordering::SeqCst), 1); } #[tokio::test] @@ -5837,9 +9639,22 @@ fn native_stream_continuation_covers_success_and_error() { ); unsafe { drop(Box::from_raw(next_ctx as *mut LlmStreamExecutionNextFn)) }; + let invalid = native_string("not-json"); + let next_ctx = Box::into_raw(Box::new(empty_llm_stream_next())) as *mut c_void; + raw = NemoRelayNativeLlmStreamV1::default(); + assert_eq!( + unsafe { native_llm_stream_next(invalid, next_ctx, &mut raw) }, + NemoRelayStatus::InvalidJson + ); + unsafe { + drop(Box::from_raw(next_ctx as *mut LlmStreamExecutionNextFn)); + native_string_free(invalid); + } + let next: LlmStreamExecutionNextFn = - Arc::new(|_| Box::pin(async { panic!("stream next panic") })); + Arc::new(|_| Box::pin(async { panic!("LLM stream next panic") })); let next_ctx = Box::into_raw(Box::new(next)) as *mut c_void; + raw = NemoRelayNativeLlmStreamV1::default(); assert_eq!( unsafe { native_llm_stream_next(request_json, next_ctx, &mut raw) }, NemoRelayStatus::Internal diff --git a/crates/plugin/Cargo.toml b/crates/plugin/Cargo.toml index 78c2cabba..a7e1af47a 100644 --- a/crates/plugin/Cargo.toml +++ b/crates/plugin/Cargo.toml @@ -13,6 +13,7 @@ description = "Rust plugin authoring SDK and stable native plugin ABI for NeMo R workspace = true [dependencies] +futures = "0.3" nemo-relay-types.workspace = true serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 8df78823f..c666836df 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -31,7 +31,7 @@ the dynamic-library boundary on the stable C-compatible ABI. - **Register real runtime behavior**: Use `PluginContext` for subscribers, guardrails, and intercepts. - **Keep a stable boundary**: Export one versioned native entry point through - the `nemo_relay_plugin!` macro. + the `nemo_relay_plugin!` or `nemo_relay_plugin_v2!` macro. - **Use host runtime helpers**: Emit events and manage scope state through the high-level `PluginRuntime` wrapper. @@ -42,20 +42,24 @@ the dynamic-library boundary on the stable C-compatible ABI. - **`PluginContext`**: Component-scoped registration APIs for middleware and subscribers. - **`PluginRuntime`**: Typed helpers for Relay-owned scopes and marks. -- **Stable native ABI v3**: C-compatible host and plugin tables behind the - safe Rust authoring interface. The v3 tables preserve a v2-compatible field - prefix, but native plugins must still be rebuilt for v3 as described in the +- **Versioned native C ABI**: C-compatible host and plugin tables behind the + safe Rust authoring interface. Manifest native API v1 uses the V3 table and + native API v2 uses its append-only V4 extension. Plugins must still be + rebuilt for V3 as described in the [0.7 migration guide](https://docs.nvidia.com/nemo/relay/reference/migration-guides#upgrade-to-nemo-relay-07). - **Raw async middleware**: Completion-based raw registrations for plugins that need asynchronous guardrails, intercepts, or event sanitizers. Typed - Rust callbacks remain synchronous convenience APIs. + Rust callbacks on native API v1 remain synchronous convenience APIs. +- **Safe native API v2 LLM continuations**: Register future-returning Rust + callbacks, dispatch explicit provider targets, and consume provider events + as Rust streams without writing C callback or handle-management code. ## Installation Add the SDK to a Rust dynamic-plugin project: ```bash -cargo add nemo-relay-plugin serde_json +cargo add nemo-relay-plugin futures serde_json ``` Configure the library as a dynamic library: @@ -94,6 +98,119 @@ Build the `cdylib`, describe its entry symbol and compatibility in a `relay-plugin.toml` manifest, then register it through the Relay CLI. See the complete example for platform-specific artifact and manifest setup. +## Native API v2 + +Existing plugins continue to use manifest `compat.native_api = "1"` and +`nemo_relay_plugin!`. A plugin that needs Relay-owned provider dispatch exports +only native API v2: + +```rust +nemo_relay_plugin::nemo_relay_plugin_v2!( + nemo_relay_register_plugin, + || ExamplePlugin +); +``` + +Set `compat.native_api = "2"` in `relay-plugin.toml`. Rust plugins normally use +`PluginContext::register_async_llm_execution_v2` and +`PluginContext::register_async_llm_stream_execution_v2`. The SDK owns the C +callback trampolines, host strings, JSON conversion, panic isolation, output +settlement, cancellation, and handle release. + +```rust +use futures::StreamExt; +use nemo_relay_plugin::{ + LlmContinuationInvocationV2, LlmContinuationTargetV2, + LlmStreamExecutionOutcomeV2, +}; + +let buffered_target = target.clone(); +ctx.register_async_llm_execution_v2("route", 0, move |_name, request, next| { + let target = buffered_target.clone(); + async move { + next.call(LlmContinuationInvocationV2 { request, target }) + .await + .map_err(|failure| format!("{failure:?}")) + } +})?; + +ctx.register_async_llm_stream_execution_v2("route-stream", 0, move |_name, request, next| { + let target = target.clone(); + async move { + let provider = next + .open_stream(LlmContinuationInvocationV2 { request, target }) + .await + .map_err(|failure| format!("{failure:?}"))?; + let stream = provider.map(|item| item.map_err(|failure| format!("{failure:?}"))); + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream))) + } +})?; +``` + +To invoke the ordinary untargeted downstream continuation for a buffered +request, call `LlmContinuationV2::call_passthrough`. For streaming, return +`LlmStreamExecutionOutcomeV2::Passthrough(request)`. The call remains inside +its managed Relay LLM lifecycle, but Relay pumps the downstream stream directly +through its bounded queue; provider events do not cross into the plugin merely +to be forwarded. + +Continuation operations follow the outer callback mode. A streaming callback +can open targeted streams or pass through to the downstream stream, but it +cannot invoke a buffered continuation. Aggregate a streamed judge or side call +inside the policy when its result is needed before returning the caller stream. + +The plugin provides JSON, an absolute target URL, and explicit target headers. +Relay sends the request with HTTP `POST` and binds that transport target to the +current LLM continuation without storing it in `LlmRequest.headers`. +Successful buffered calls return provider JSON up to 16 MiB. Provider +rejections return an HTTP status, bounded body, and safe response headers; +failures without an HTTP response use a small transport-oriented kind. The +plugin owns its retry and fallback policy. + +Relay core performs the terminal targeted HTTP request after the remaining LLM +execution intercepts run. This contract is host-independent: it works through +the CLI gateway and through SDK-embedded Relay hosts that call the managed LLM +execution APIs directly. + +Streaming dispatch returns `LlmProviderStreamV2`, which implements Rust +`Stream` and cancels unfinished provider work on drop. Targeted streaming +endpoints must return SSE with a JSON value in each `data` frame; the plugin +receives those JSON events rather than raw SSE framing. Provider streams permit +at most one outstanding pull; plugin output and direct pass-through use bounded +queues. Safe callbacks register through the generic V3 +asynchronous-middleware APIs and return `Pending`. Relay then polls their Rust +futures and returned streams cooperatively on its Tokio runtime. Each resumed +poll restores the captured Relay continuation and scope context, and a pending +callback does not occupy a blocking worker. Output backpressure parks the task +until the bounded host queue can accept more data. No Rust future, trait object, +`serde_json::Value`, or allocator-owned Rust string crosses the C ABI boundary. + +Callback futures and streams must be executor-neutral. A native plugin shared +library can link a different copy of an async runtime than the Relay host, so +host-side polling does not enter plugin-local runtime state. An integration may +instead own and bridge its own runtime explicitly, but it must not assume that +Tokio APIs such as `tokio::spawn` can discover Relay's runtime across the +dynamic-library boundary. The SDK itself has no Tokio dependency. + +The raw `PluginContext::host_api_v4` table and generic V3 `Pending` +registration methods remain available for advanced ABI consumers and non-Rust +bindings. V4 adds targeted LLM continuation and host-task operations; it does +not define a separate blocking registration model. Code using the raw tables is +responsible for every callback lifetime, host string, completion, task and +stream settlement, cancellation, and release operation. + +The manifest API number is distinct from the internal host-table ABI number: +native API v1 negotiates the V3 host table and native API v2 negotiates V4. +Native plugins are trusted in-process extensions. A v2 plugin owns its target +credentials; Relay transports them but excludes their values from diagnostics +and observability. + +Clean plugin teardown unloads the native library normally. If teardown finds +an opaque callback, task, continuation, or stream handle that still owns plugin +code, Relay conservatively keeps only that library mapping loaded for the rest +of the process. All descriptor and handle state is still released. This avoids +unmapping a plugin while an escaped handle is returning through its own code. + ## Documentation - [NeMo Relay documentation](https://docs.nvidia.com/nemo/relay) diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index b6bf3da7d..18e0f1d32 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -9,7 +9,9 @@ //! Native plugins built with it communicate with a host through versioned //! C-compatible tables and host-owned string handles. +use std::collections::BTreeMap; use std::ffi::{c_char, c_void}; +use std::fmt; use std::marker::{PhantomData, PhantomPinned}; use std::panic::{AssertUnwindSafe, catch_unwind}; use std::ptr; @@ -35,13 +37,22 @@ pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; -/// Native plugin ABI version supported by this crate. +mod native_v2; + +pub use native_v2::{ + LlmContinuationV2, LlmJsonAsyncStreamV2, LlmProviderStreamV2, LlmStreamContinuationV2, + LlmStreamExecutionOutcomeV2, +}; + +/// Native ABI used by the original native API v1 SDK and export macro. /// -/// Version 3 reserves the native async middleware extension. Hosts retain a -/// version-2 table for already-built plugins during entry-point negotiation. +/// Relay preserves this value so plugins rebuilt with the current SDK and the +/// existing [`nemo_relay_plugin!`] macro remain native API v1 plugins. pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 3; /// ABI version that introduced completion-based asynchronous middleware. pub const NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE: u32 = 3; +/// ABI version that introduced targeted LLM continuations and structured outcomes. +pub const NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS: u32 = 4; /// Legacy native plugin ABI accepted by Relay hosts for compatibility. pub const NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY: u32 = 2; @@ -117,6 +128,8 @@ pub enum NemoRelayStatus { InvalidArg = 9, /// A stream reached end-of-stream and has no chunk to return. StreamEnd = 10, + /// The operation would block until host-side capacity becomes available. + WouldBlock = 11, } /// Opaque host-owned UTF-8 string or JSON byte buffer. @@ -871,6 +884,43 @@ pub struct NemoRelayNativeAsyncStream { _marker: PhantomData<(*mut u8, PhantomPinned)>, } +/// Opaque host-scheduled cooperative task for native API v2 callbacks. +/// +/// A task is polled on Relay's Tokio runtime with the middleware invocation's +/// full continuation context restored. The task pointer passed to its poll +/// callback is borrowed. A plugin that stores it in a Rust `Waker` must retain +/// one reference for every stored clone and release those references exactly +/// once. +/// +/// Host polling does not enter runtime-local state linked into a plugin shared +/// library. The polled task must therefore be executor-neutral unless the +/// plugin explicitly owns and bridges its own runtime. +#[repr(C)] +pub struct NemoRelayNativeAsyncTaskV2 { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + +/// Polls one host-scheduled native API v2 task. +/// +/// The callback returns [`NemoRelayNativeAsyncCallbackState::Pending`] after +/// arranging for a later `async_task_wake_v2`, or `Complete` after settling its +/// associated completion or output stream. Relay serializes polls and restores +/// the invocation context before every call. +pub type NemoRelayNativeAsyncTaskPollCbV2 = + unsafe extern "C" fn(user_data: *mut c_void, task: *const NemoRelayNativeAsyncTaskV2) -> u32; + +/// Opaque host-owned provider stream returned by a native API v2 LLM continuation. +/// +/// The plugin requests one item at a time with the v2 host table, then releases +/// the handle exactly once. Releasing an unfinished stream cancels provider +/// production. +#[repr(C)] +pub struct NemoRelayNativeLlmStreamV2 { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + /// Receives one downstream stream item. `chunk_json` is non-null for a chunk, /// `error` is non-null for failure or consumer cancellation, and `done` marks /// clean completion. Unless the callback itself returns `false`, the host @@ -897,6 +947,130 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( error: *const NemoRelayNativeString, ); +/// Explicit provider target for one native API v2 LLM continuation. +/// +/// Relay sends LLM continuation requests with HTTP `POST`. +#[derive(Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +pub struct LlmContinuationTargetV2 { + /// Absolute HTTP(S) provider URL including the selected endpoint. + pub url: String, + /// Explicit outbound provider headers, including target credentials. + /// + /// Relay validates and transports these headers but never records their + /// values in plugin diagnostics or observability events. + pub headers: BTreeMap, +} + +impl fmt::Debug for LlmContinuationTargetV2 { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("LlmContinuationTargetV2") + .field("url", &"") + .field("header_names", &self.headers.keys().collect::>()) + .finish() + } +} + +/// Typed LLM continuation invocation supplied through native API v2. +#[derive(Clone, PartialEq, Serialize, serde::Deserialize)] +pub struct LlmContinuationInvocationV2 { + /// Replacement request passed to the Relay execution continuation. + pub request: LlmRequest, + /// Explicit provider target selected by the plugin. + pub target: LlmContinuationTargetV2, +} + +impl fmt::Debug for LlmContinuationInvocationV2 { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("LlmContinuationInvocationV2") + .field("request", &"") + .field("target", &self.target) + .finish() + } +} + +/// Stable non-HTTP failure classification exposed through native API v2. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LlmNonHttpFailureKindV2 { + /// Provider connection could not be established or was interrupted. + Transport, + /// Provider request timed out. + Timeout, + /// The caller cancelled the operation. + Cancelled, + /// Relay rejected an invalid continuation invocation. + InvalidRequest, + /// A guardrail rejected the provider call. + Guardrail, + /// Relay could not complete the operation. + Internal, +} + +/// Structured LLM continuation failure exposed through native API v2. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(tag = "failure_type", rename_all = "snake_case")] +pub enum LlmContinuationFailureV2 { + /// A provider returned a non-success HTTP response. + Http { + /// Provider HTTP status. + status: u16, + /// Bounded provider response body. + body: String, + /// Safe response headers with credential-bearing fields removed. + headers: BTreeMap, + }, + /// No provider HTTP response was available. + NonHttp { + /// Stable failure kind. + kind: LlmNonHttpFailureKindV2, + /// Bounded human-readable context. + message: String, + }, +} + +/// Receives one typed unary LLM continuation outcome. +/// +/// Exactly one of `response_json` and `error_json` is non-null. A response is +/// provider JSON; an error is a serialized [`LlmContinuationFailureV2`]. Both +/// strings are borrowed for the callback. +pub type NemoRelayNativeAsyncLlmResultCbV2 = unsafe extern "C" fn( + user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, +); + +/// Receives the result of opening one typed streaming LLM continuation. +/// +/// Exactly one of `stream` and `error_json` is non-null. A non-null stream is +/// an owned plugin reference that must be released exactly once. +pub type NemoRelayNativeAsyncLlmStreamOpenCbV2 = unsafe extern "C" fn( + user_data: *mut c_void, + stream: *const NemoRelayNativeLlmStreamV2, + error_json: *const NemoRelayNativeString, +); + +/// Receives one item from a native API v2 provider stream. +/// +/// A chunk has non-null `chunk_json`, null `error_json`, and `done = false`. +/// Clean completion has both strings null and `done = true`. Failure has null +/// `chunk_json`, a serialized [`LlmContinuationFailureV2`] in `error_json`, and +/// `done = true`. Both strings are borrowed for the callback. Only one `next` +/// operation may be active per stream. +pub type NemoRelayNativeAsyncLlmStreamNextCbV2 = unsafe extern "C" fn( + user_data: *mut c_void, + chunk_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, + done: bool, +); + +/// Receives terminal settlement of a directly forwarded downstream stream. +/// +/// Relay settles the output stream before invoking this callback, so the +/// callback is only a wake-up that lets the plugin reclaim `user_data`. +pub type NemoRelayNativeAsyncLlmStreamForwardCbV2 = unsafe extern "C" fn(user_data: *mut c_void); + /// Incremental native LLM stream intercept callback. /// /// The callback owns `next` and `stream` and must release each exactly once. @@ -991,6 +1165,8 @@ pub struct NemoRelayNativeHostApiV3 { /// discriminant. The host rejects unknown `u32` values and /// [`NemoRelayNativeAsyncMiddlewareKind::LlmStreamExecutionIntercept`], /// which must use `plugin_context_register_async_stream_middleware`. + /// The host consumes `user_data` once this function is called and invokes + /// `free_fn` exactly once, including when validation or registration fails. pub plugin_context_register_async_middleware: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, kind: u32, @@ -1003,9 +1179,10 @@ pub struct NemoRelayNativeHostApiV3 { ) -> NemoRelayStatus, /// Pushes one JSON chunk to an incremental native stream without blocking. /// - /// A full bounded host queue returns [`NemoRelayStatus::Internal`] and - /// records a backpressure message in the host's last-error slot. The - /// producer may retry the logical chunk after the consumer advances. + /// A full bounded host queue returns [`NemoRelayStatus::Internal`] for a + /// native API v1 host or [`NemoRelayStatus::WouldBlock`] through the ABI-v4 + /// native API v2 table. The producer may retry the logical chunk after the + /// consumer advances. pub async_stream_push_json: unsafe extern "C" fn( stream: *const NemoRelayNativeAsyncStream, chunk_json: *const NemoRelayNativeString, @@ -1015,8 +1192,10 @@ pub struct NemoRelayNativeHostApiV3 { unsafe extern "C" fn(stream: *const NemoRelayNativeAsyncStream) -> NemoRelayStatus, /// Rejects an incremental native stream without blocking. /// - /// A full bounded queue returns [`NemoRelayStatus::Internal`]; the caller - /// may retry the rejection after the consumer advances. + /// A full bounded host queue returns [`NemoRelayStatus::Internal`] for a + /// native API v1 host or [`NemoRelayStatus::WouldBlock`] through the ABI-v4 + /// native API v2 table. The caller may retry the rejection after the + /// consumer advances. pub async_stream_reject: unsafe extern "C" fn( stream: *const NemoRelayNativeAsyncStream, message: *const NemoRelayNativeString, @@ -1041,6 +1220,9 @@ pub struct NemoRelayNativeHostApiV3 { user_data: *mut c_void, ) -> NemoRelayStatus, /// Registers an incremental asynchronous LLM stream intercept. + /// + /// The host consumes `user_data` once this function is called and invokes + /// `free_fn` exactly once, including when validation or registration fails. pub plugin_context_register_async_stream_middleware: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -1062,6 +1244,106 @@ pub struct NemoRelayNativeHostApiV3 { ) -> NemoRelayStatus, } +/// ABI-v4 host extension implementing native API v2 targeted LLM continuations +/// and cooperative plugin-task scheduling. +/// +/// Its first field is the complete ABI-v3 table. Native API v1 plugins +/// continue to receive ABI-v3 or ABI-v2 tables during entry-point negotiation. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct NemoRelayNativeHostApiV4 { + /// Compatibility prefix containing the ABI-v3 asynchronous middleware API. + pub v3: NemoRelayNativeHostApiV3, + /// Invokes a unary LLM continuation with an explicit target and structured + /// outcome. + /// + /// When the function returns [`NemoRelayStatus::Ok`], the host consumes + /// `user_data` and invokes `cb` exactly once. On any other status, the + /// caller retains `user_data` and the callback is not invoked. + pub async_llm_next_invoke_result_v2: unsafe extern "C" fn( + next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Opens a streaming LLM continuation with an explicit target. + /// + /// On success the callback receives an owned provider-stream handle whose + /// items are read with `async_llm_stream_next_v2`. + /// When the function returns [`NemoRelayStatus::Ok`], the host consumes + /// `user_data` and invokes `cb` exactly once. On any other status, the + /// caller retains `user_data` and the callback is not invoked. + pub async_llm_next_open_stream_v2: unsafe extern "C" fn( + next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Requests one provider event from a native API v2 stream. + /// + /// When the function returns [`NemoRelayStatus::Ok`], the host consumes + /// `user_data` and invokes `cb` exactly once. On any other status, the + /// caller retains `user_data` and the callback is not invoked. + pub async_llm_stream_next_v2: unsafe extern "C" fn( + stream: *const NemoRelayNativeLlmStreamV2, + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Releases the plugin-owned provider-stream reference, cancelling provider + /// production first when the stream has not reached a terminal item. If a + /// `next` operation is active, its callback receives the cancellation + /// outcome and remains the sole owner responsible for reclaiming its + /// `user_data`. + pub async_llm_stream_release_v2: + unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2), + /// Starts a cooperative task associated with an async completion. + /// + /// Relay polls the task on its Tokio runtime and wakes it when the owning + /// completion is cancelled. A successful call consumes `user_data` and + /// invokes `free_fn` exactly once after completion or cancellation. A + /// failed call does not consume either value. + pub async_completion_spawn_task_v2: unsafe extern "C" fn( + completion: *const NemoRelayNativeAsyncCompletion, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, + /// Starts a cooperative task associated with an incremental output stream. + /// + /// Relay polls the task on its Tokio runtime and wakes it when the consumer + /// cancels or when bounded output backpressure clears. A successful call + /// consumes `user_data` and invokes `free_fn` exactly once after settlement + /// or cancellation. A failed call does not consume either value. + pub async_stream_spawn_task_v2: unsafe extern "C" fn( + stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, + /// Retains one reference to a cooperative task for a stored waker clone. + pub async_task_retain_v2: unsafe extern "C" fn(task: *const NemoRelayNativeAsyncTaskV2), + /// Schedules a cooperative task for another serialized poll. + pub async_task_wake_v2: unsafe extern "C" fn(task: *const NemoRelayNativeAsyncTaskV2), + /// Releases one previously retained cooperative-task reference. + pub async_task_release_v2: unsafe extern "C" fn(task: *const NemoRelayNativeAsyncTaskV2), + /// Forwards the ordinary downstream LLM stream directly into a native + /// interceptor's output stream. + /// + /// Relay owns event pumping and bounded backpressure. No provider event + /// crosses the plugin boundary. The terminal callback runs exactly once + /// after output settlement when this function returns [`NemoRelayStatus::Ok`]. + /// An accepted call consumes `user_data`; a failed call does not consume it + /// and does not invoke the callback. + pub async_llm_next_forward_stream_v2: unsafe extern "C" fn( + next: *const NemoRelayNativeAsyncNext, + request_json: *const NemoRelayNativeString, + output_stream: *const NemoRelayNativeAsyncStream, + terminal_callback: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, +} + unsafe impl Send for NemoRelayNativeHostApiV3 {} unsafe impl Sync for NemoRelayNativeHostApiV3 {} @@ -1789,6 +2071,14 @@ impl<'a> PluginContext<'a> { self.host } + /// Returns the native API v2 host extension when the plugin was loaded + /// through ABI v4. + pub fn host_api_v4(&self) -> Option<&'a NemoRelayNativeHostApiV4> { + (self.host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS + && self.host.struct_size >= std::mem::size_of::()) + .then(|| unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV4) }) + } + /// Returns a cloneable high-level runtime handle. pub fn runtime(&self) -> PluginRuntime { PluginRuntime::new(self.host) @@ -2184,8 +2474,8 @@ impl<'a> PluginContext<'a> { /// Registers a typed LLM stream execution intercept. /// - /// Native ABI v2 represents stream execution as one JSON result. The host - /// wraps that result as a one-chunk stream. + /// The host pulls the returned [`LlmJsonStream`] incrementally through an + /// opaque native stream handle. pub fn register_llm_stream_execution_intercept( &mut self, name: &str, @@ -2590,12 +2880,12 @@ impl<'a> PluginContext<'a> { /// # Safety /// The callback and user data must remain valid until deregistration or /// `free_fn`; callback-owned `next` and `stream` handles must each be - /// released exactly once. Stream pushes and rejection are nonblocking: - /// `Internal` with a host last-error containing `backpressured` means the - /// bounded queue is full and the operation may be retried. The output - /// stream owns the callback lifetime. `next` may be invoked repeatedly or - /// concurrently until that stream settles; Relay then rejects or cancels - /// unfinished and later calls. + /// released exactly once. Stream pushes and rejection are nonblocking. A + /// full host queue returns [`NemoRelayStatus::Internal`] for native API v1 + /// or [`NemoRelayStatus::WouldBlock`] for native API v2, and the operation + /// may be retried. The output stream owns the callback lifetime. `next` may + /// be invoked repeatedly or concurrently until that stream settles; Relay + /// then rejects or cancels unfinished and later calls. pub unsafe fn register_async_stream_middleware_raw( &mut self, name: &str, @@ -3205,6 +3495,11 @@ struct HostString<'a> { ptr: *mut NemoRelayNativeString, } +// The string is an exclusively owned host allocation. Native API string +// allocation and release are thread-safe, and the borrowed immutable host +// table remains valid for this value's lifetime. +unsafe impl Send for HostString<'_> {} + impl<'a> HostString<'a> { fn try_new( host: &'a NemoRelayNativeHostApiV1, @@ -3288,11 +3583,16 @@ impl<'a> OptionalHostJson<'a> { enum OwnedHostApi { V1(NemoRelayNativeHostApiV1), V3(NemoRelayNativeHostApiV3), + V4(NemoRelayNativeHostApiV4), } impl OwnedHostApi { unsafe fn copy_from(host: &NemoRelayNativeHostApiV1) -> Self { - if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE + if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS + && host.struct_size >= std::mem::size_of::() + { + Self::V4(unsafe { *(host as *const _ as *const NemoRelayNativeHostApiV4) }) + } else if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE && host.struct_size >= std::mem::size_of::() { Self::V3(unsafe { *(host as *const _ as *const NemoRelayNativeHostApiV3) }) @@ -3305,6 +3605,7 @@ impl OwnedHostApi { match self { Self::V1(host) => host, Self::V3(host) => &host.v1, + Self::V4(host) => &host.v3.v1, } } } @@ -3582,7 +3883,13 @@ pub unsafe fn export_plugin( } unsafe { *out = NemoRelayNativePluginV1::default() }; let host_ref = unsafe { &*host }; - export_plugin_checked(host_ref, out, || plugin) + export_plugin_checked( + host_ref, + out, + NEMO_RELAY_NATIVE_ABI_VERSION, + std::mem::size_of::(), + || plugin, + ) } /// Initializes a native plugin descriptor from a constructor callback. @@ -3606,22 +3913,60 @@ where } unsafe { *out = NemoRelayNativePluginV1::default() }; let host_ref = unsafe { &*host }; - export_plugin_checked(host_ref, out, constructor) + export_plugin_checked( + host_ref, + out, + NEMO_RELAY_NATIVE_ABI_VERSION, + std::mem::size_of::(), + constructor, + ) +} + +/// Initializes a native API v2 plugin descriptor from a constructor callback. +/// +/// # Safety +/// `host` must point to a complete [`NemoRelayNativeHostApiV4`] table for the +/// duration of the call, and `out` must point to writable memory for one +/// [`NemoRelayNativePluginV1`] descriptor. +#[doc(hidden)] +pub unsafe fn __export_plugin_v2_from_constructor( + host: *const NemoRelayNativeHostApiV1, + out: *mut NemoRelayNativePluginV1, + constructor: F, +) -> NemoRelayStatus +where + P: NativePlugin, + F: FnOnce() -> P, +{ + if host.is_null() || out.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { *out = NemoRelayNativePluginV1::default() }; + let host_ref = unsafe { &*host }; + export_plugin_checked( + host_ref, + out, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, + std::mem::size_of::(), + constructor, + ) } fn export_plugin_checked( host_ref: &NemoRelayNativeHostApiV1, out: *mut NemoRelayNativePluginV1, + required_abi_version: u32, + required_struct_size: usize, constructor: F, ) -> NemoRelayStatus where P: NativePlugin, F: FnOnce() -> P, { - if host_ref.abi_version != NEMO_RELAY_NATIVE_ABI_VERSION { + if host_ref.abi_version != required_abi_version { return NemoRelayStatus::InvalidArg; } - if host_ref.struct_size < std::mem::size_of::() { + if host_ref.struct_size < required_struct_size { return NemoRelayStatus::InvalidArg; } @@ -3677,3 +4022,34 @@ macro_rules! nemo_relay_plugin { } }; } + +/// Exports a native API v2-only plugin entry symbol. +/// +/// The generated entry rejects native API v1 host tables. Use this macro when +/// the plugin requires targeted LLM continuations from [`NemoRelayNativeHostApiV4`]. +#[macro_export] +macro_rules! nemo_relay_plugin_v2 { + ($symbol:ident, $constructor:expr) => { + #[doc = "Native API v2 plugin entry symbol generated by `nemo_relay_plugin_v2!`."] + #[unsafe(no_mangle)] + pub unsafe extern "C" fn $symbol( + host: *const $crate::NemoRelayNativeHostApiV1, + out: *mut $crate::NemoRelayNativePluginV1, + ) -> $crate::NemoRelayStatus { + match ::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(|| unsafe { + $crate::__export_plugin_v2_from_constructor(host, out, $constructor) + })) { + Ok(status) => status, + Err(_) => { + unsafe { + $crate::__set_last_error_from_entry( + host, + "native API v2 plugin entry callback panicked", + ) + }; + $crate::NemoRelayStatus::Internal + } + } + } + }; +} diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs new file mode 100644 index 000000000..03d66eec5 --- /dev/null +++ b/crates/plugin/src/native_v2.rs @@ -0,0 +1,1204 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Safe Rust facade for native API v2 LLM continuations. + +use std::ffi::c_void; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use futures::channel::oneshot; +use futures::future::FutureExt; +use futures::task::{ArcWake, waker_ref}; +use futures::{Stream, StreamExt}; +use serde::Deserialize; + +use super::{ + HostString, Json, LlmContinuationFailureV2, LlmContinuationInvocationV2, + LlmNonHttpFailureKindV2, LlmRequest, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncStream, NemoRelayNativeAsyncTaskV2, NemoRelayNativeHostApiV1, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmStreamV2, NemoRelayNativeString, NemoRelayStatus, + PluginContext, Result, read_json_value, read_required_host_string, set_last_error, +}; + +const MAX_SDK_ERROR_BYTES: usize = 4 * 1024; + +/// Asynchronous JSON stream returned by a safe native API v2 stream callback. +pub type LlmJsonAsyncStreamV2 = Pin> + Send>>; + +/// Result of a safe native API v2 streaming execution callback. +pub enum LlmStreamExecutionOutcomeV2 { + /// Relay emits the plugin-produced stream. + Stream(LlmJsonAsyncStreamV2), + /// Relay forwards the ordinary downstream continuation itself. + /// + /// Provider events stay inside Relay and are copied into the caller stream + /// with the host's bounded backpressure path. + Passthrough(LlmRequest), +} + +/// Cloneable targeted buffered LLM continuation. +/// +/// Clones may be called repeatedly or concurrently. The underlying C handle +/// is released exactly once after the final clone and all in-flight calls are +/// dropped. +#[derive(Clone)] +pub struct LlmContinuationV2 { + inner: Arc, +} + +/// Cloneable targeted streaming LLM continuation. +/// +/// Clones may open provider streams repeatedly or concurrently. Each returned +/// provider stream has independent pull and cancellation state. +#[derive(Clone)] +pub struct LlmStreamContinuationV2 { + inner: Arc, +} + +/// Provider stream returned by [`LlmStreamContinuationV2::open_stream`]. +/// +/// Dropping an unfinished stream cancels provider production. At most one host +/// pull is outstanding for a stream at any time. +pub struct LlmProviderStreamV2 { + host: NemoRelayNativeHostApiV4, + raw: *const NemoRelayNativeLlmStreamV2, + pending: Option, LlmContinuationFailureV2>>>, + finished: bool, +} + +struct ContinuationInner { + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, +} + +enum HostTaskCancellation { + Completion(*const NemoRelayNativeAsyncCompletion), + Stream(*const NemoRelayNativeAsyncStream), +} + +#[derive(Clone, Copy)] +struct CompletionHandle(*const NemoRelayNativeAsyncCompletion); + +#[derive(Clone, Copy)] +struct OutputHandle(*const NemoRelayNativeAsyncStream); + +// The host owns these opaque handles and documents their operations as +// thread-safe for a pending callback's lifetime. +unsafe impl Send for CompletionHandle {} +unsafe impl Sync for CompletionHandle {} +unsafe impl Send for OutputHandle {} +unsafe impl Sync for OutputHandle {} + +impl CompletionHandle { + fn as_ptr(&self) -> *const NemoRelayNativeAsyncCompletion { + self.0 + } +} + +impl OutputHandle { + fn as_ptr(&self) -> *const NemoRelayNativeAsyncStream { + self.0 + } +} + +struct HostFutureTask { + host: NemoRelayNativeHostApiV4, + cancellation: HostTaskCancellation, + completion_to_release: Option<*const NemoRelayNativeAsyncCompletion>, + future: Option + Send>>>, + waker: Option>, +} + +struct CompletionReleaseGuard { + host: NemoRelayNativeHostApiV4, + completion: *const NemoRelayNativeAsyncCompletion, +} + +struct HostTaskWaker { + host: NemoRelayNativeHostApiV4, + raw: *const NemoRelayNativeAsyncTaskV2, +} + +// Relay serializes task polling and documents retained task handles as safe to +// wake from any thread. +unsafe impl Send for HostFutureTask {} +unsafe impl Send for HostTaskWaker {} +unsafe impl Sync for HostTaskWaker {} + +impl HostFutureTask { + fn is_cancelled(&self) -> bool { + unsafe { + match self.cancellation { + HostTaskCancellation::Completion(completion) => { + (self.host.v3.async_completion_is_cancelled)(completion) + } + HostTaskCancellation::Stream(stream) => { + (self.host.v3.async_stream_is_cancelled)(stream) + } + } + } + } +} + +impl Drop for HostFutureTask { + fn drop(&mut self) { + // Keep the callback-owned completion (and therefore the host's plugin + // library guard) alive until every plugin-owned destructor has run. + // User futures and streams may panic from Drop, so catch those panics + // long enough for the completion release guard to run exactly once, + // then resume unwinding for the outer FFI panic fence to report it. + let completion_release = + self.completion_to_release + .take() + .map(|completion| CompletionReleaseGuard { + host: self.host, + completion, + }); + let future_panic = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(self.future.take()))) + .err(); + let waker_panic = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(self.waker.take()))) + .err(); + drop(completion_release); + if let Some(payload) = future_panic.or(waker_panic) { + std::panic::resume_unwind(payload); + } + } +} + +impl Drop for CompletionReleaseGuard { + fn drop(&mut self) { + unsafe { (self.host.v3.async_completion_release)(self.completion) }; + } +} + +impl HostTaskWaker { + unsafe fn new( + host: NemoRelayNativeHostApiV4, + raw: *const NemoRelayNativeAsyncTaskV2, + ) -> Arc { + unsafe { (host.async_task_retain_v2)(raw) }; + Arc::new(Self { host, raw }) + } +} + +impl ArcWake for HostTaskWaker { + fn wake_by_ref(arc_self: &Arc) { + unsafe { (arc_self.host.async_task_wake_v2)(arc_self.raw) }; + } +} + +impl Drop for HostTaskWaker { + fn drop(&mut self) { + unsafe { (self.host.async_task_release_v2)(self.raw) }; + } +} + +unsafe extern "C" fn poll_host_future_task( + user_data: *mut c_void, + task: *const NemoRelayNativeAsyncTaskV2, +) -> u32 { + if user_data.is_null() || task.is_null() { + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + let state = unsafe { &mut *user_data.cast::() }; + if state.is_cancelled() { + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + if state.waker.is_none() { + state.waker = Some(unsafe { HostTaskWaker::new(state.host, task) }); + } + let waker = waker_ref( + state + .waker + .as_ref() + .expect("host task waker was initialized"), + ); + let mut context = Context::from_waker(&waker); + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + state + .future + .as_mut() + .expect("host task future was initialized") + .as_mut() + .poll(&mut context) + })) { + Ok(Poll::Ready(())) => NemoRelayNativeAsyncCallbackState::Complete as u32, + Ok(Poll::Pending) => NemoRelayNativeAsyncCallbackState::Pending as u32, + Err(_) => { + set_last_error(&state.host.v3.v1, "native API v2 host task panicked"); + NemoRelayNativeAsyncCallbackState::Complete as u32 + } + } +} + +unsafe extern "C" fn drop_host_future_task(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let host = state.host.v3.v1; + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(state))).is_err() { + set_last_error(&host, "native API v2 host task state drop panicked"); + } +} + +// Host tables are immutable, and Relay documents retained continuation +// handles as thread-safe for repeated and concurrent invocation. +unsafe impl Send for ContinuationInner {} +unsafe impl Sync for ContinuationInner {} + +impl Drop for ContinuationInner { + fn drop(&mut self) { + unsafe { (self.host.v3.async_next_release)(self.next) }; + } +} + +struct StreamContinuationInner { + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, + output: *const NemoRelayNativeAsyncStream, +} + +// The output handle is owned exclusively by this shared RAII state. Host +// operations synchronize settlement and cancellation. +unsafe impl Send for StreamContinuationInner {} +unsafe impl Sync for StreamContinuationInner {} + +impl Drop for StreamContinuationInner { + fn drop(&mut self) { + unsafe { + (self.host.v3.async_stream_release)(self.output); + (self.host.v3.async_next_release)(self.next); + } + } +} + +impl LlmContinuationV2 { + unsafe fn from_raw( + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, + ) -> std::result::Result { + if next.is_null() { + return Err(NemoRelayStatus::NullPointer); + } + Ok(Self { + inner: Arc::new(ContinuationInner { host, next }), + }) + } + + /// Dispatches one explicitly targeted LLM continuation. + pub async fn call( + &self, + invocation: LlmContinuationInvocationV2, + ) -> std::result::Result { + let (sender, receiver) = oneshot::channel(); + let state = Box::new(TargetedResultCallback { + host: self.inner.host.v3.v1, + sender, + }); + let state = Box::into_raw(state).cast::(); + // An accepted host call consumes this state and returns it through the + // callback exactly once. A failed call leaves it with this caller. + let status = match HostString::from_json(&self.inner.host.v3.v1, &invocation) { + Some(invocation) => unsafe { + (self.inner.host.async_llm_next_invoke_result_v2)( + self.inner.next, + invocation.as_ptr(), + targeted_result_callback, + state, + ) + }, + None => { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(internal_failure( + "failed to serialize targeted LLM continuation invocation", + )); + } + }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(status_failure("targeted LLM continuation", status)); + } + receiver.await.unwrap_or_else(|_| { + Err(internal_failure( + "targeted LLM continuation callback closed without an outcome", + )) + }) + } + + /// Invokes the ordinary buffered downstream continuation. + pub async fn call_passthrough(&self, request: LlmRequest) -> Result { + let (sender, receiver) = oneshot::channel(); + let state = Box::new(PassthroughResultCallback { + host: self.inner.host.v3.v1, + sender, + }); + let state = Box::into_raw(state).cast::(); + let status = match HostString::from_json(&self.inner.host.v3.v1, &request) { + Some(request) => unsafe { + (self.inner.host.v3.async_next_invoke_result)( + self.inner.next, + request.as_ptr(), + passthrough_result_callback, + state, + ) + }, + None => { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err("failed to serialize LLM pass-through request".into()); + } + }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(format!("LLM pass-through continuation failed: {status:?}")); + } + receiver.await.unwrap_or_else(|_| { + Err("LLM pass-through continuation callback closed without a result".into()) + }) + } +} + +impl LlmStreamContinuationV2 { + unsafe fn from_raw( + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, + output: *const NemoRelayNativeAsyncStream, + ) -> std::result::Result { + if next.is_null() || output.is_null() { + return Err(NemoRelayStatus::NullPointer); + } + Ok(Self { + inner: Arc::new(StreamContinuationInner { host, next, output }), + }) + } + + /// Opens one explicitly targeted provider stream. + pub async fn open_stream( + &self, + invocation: LlmContinuationInvocationV2, + ) -> std::result::Result { + let host = self.inner.host; + let (sender, receiver) = oneshot::channel(); + let state = Box::new(StreamOpenCallback { host, sender }); + let state = Box::into_raw(state).cast::(); + // An accepted host call consumes this state and returns it through the + // callback exactly once. A failed call leaves it with this caller. + let status = match HostString::from_json(&host.v3.v1, &invocation) { + Some(invocation) => unsafe { + (host.async_llm_next_open_stream_v2)( + self.inner.next, + invocation.as_ptr(), + self.inner.output, + stream_open_callback, + state, + ) + }, + None => { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(internal_failure( + "failed to serialize targeted LLM stream invocation", + )); + } + }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(status_failure("targeted LLM stream setup", status)); + } + receiver.await.unwrap_or_else(|_| { + Err(internal_failure( + "targeted LLM stream setup callback closed without an outcome", + )) + }) + } + + async fn forward_passthrough(&self, request: LlmRequest) -> Result<()> { + let host = self.inner.host; + let (sender, receiver) = oneshot::channel(); + let state = Box::new(ForwardStreamCallback { sender }); + let state = Box::into_raw(state).cast::(); + // An accepted host call consumes this state and returns it through the + // terminal callback exactly once. A failed call leaves it here. + let status = match HostString::from_json(&host.v3.v1, &request) { + Some(request) => unsafe { + (host.async_llm_next_forward_stream_v2)( + self.inner.next, + request.as_ptr(), + self.inner.output, + forward_stream_callback, + state, + ) + }, + None => { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err("failed to serialize streaming pass-through request".into()); + } + }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(state.cast::())) }; + return Err(format!( + "streaming pass-through continuation failed: {status:?}" + )); + } + receiver + .await + .map_err(|_| "streaming pass-through callback closed before settlement".into()) + } +} + +// Relay retains the provider stream handle until release and serializes host +// pulls. The Rust wrapper owns the sole plugin reference. +unsafe impl Send for LlmProviderStreamV2 {} + +impl Stream for LlmProviderStreamV2 { + type Item = std::result::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.finished { + return Poll::Ready(None); + } + if self.pending.is_none() { + let (sender, receiver) = oneshot::channel(); + let state = Box::new(ProviderNextCallback { + host: self.host.v3.v1, + sender, + }); + let state = Box::into_raw(state).cast::(); + // A successful poll transfers the callback state to the host. A + // failed poll leaves it here; stream release cancels an accepted + // poll and lets its callback reclaim the state exactly once. + let status = unsafe { + (self.host.async_llm_stream_next_v2)(self.raw, provider_next_callback, state) + }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(state.cast::())) }; + self.finished = true; + return Poll::Ready(Some(Err(status_failure( + "targeted provider stream poll", + status, + )))); + } + self.pending = Some(receiver); + } + + let receiver = self + .pending + .as_mut() + .expect("provider stream pending receiver was initialized"); + match Pin::new(receiver).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(result) => { + self.pending = None; + let item = result.unwrap_or_else(|_| { + Err(internal_failure( + "provider stream callback closed without an event", + )) + }); + match item { + Ok(Some(chunk)) => Poll::Ready(Some(Ok(chunk))), + Ok(None) => { + self.finished = true; + Poll::Ready(None) + } + Err(error) => { + self.finished = true; + Poll::Ready(Some(Err(error))) + } + } + } + } + } +} + +impl Drop for LlmProviderStreamV2 { + fn drop(&mut self) { + unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; + } +} + +#[derive(Deserialize)] +struct LlmCallbackInvocation { + name: String, + request: LlmRequest, +} + +struct SafeV2Callback { + host: NemoRelayNativeHostApiV4, + callback: Arc, +} + +unsafe extern "C" fn drop_safe_v2_callback(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::>()) }; + let host = state.host.v3.v1; + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(state))).is_err() { + set_last_error(&host, "native API v2 safe callback state drop panicked"); + } +} + +impl PluginContext<'_> { + /// Registers a safe asynchronous native API v2 buffered LLM execution callback. + /// + /// The callback receives owned Rust values and a cloneable continuation. + /// Relay cooperatively polls its future on the host runtime with the + /// invocation's active scope stack restored for every poll. + /// The future must remain executor-neutral: a plugin shared library may + /// link a different async-runtime instance, whose runtime-local state is + /// not entered merely because Relay polls the future on its host runtime. + pub fn register_async_llm_execution_v2( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: Fn(String, LlmRequest, LlmContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let host = self + .host_api_v4() + .copied() + .ok_or_else(|| "native API v2 requires a complete ABI-v4 host table".to_string())?; + let name = HostString::try_new(&host.v3.v1, name).map_err(|status| { + format!("native API v2 buffered LLM registration name failed: {status:?}") + })?; + let state = Box::new(SafeV2Callback { + host, + callback: Arc::new(callback), + }); + let user_data = Box::into_raw(state).cast::(); + let status = unsafe { + // Once invoked, the host consumes `user_data` on both success and + // failure and calls `free_fn` exactly once. Allocate the fallible + // name first so local ownership is never ambiguous. + (host.v3.plugin_context_register_async_middleware)( + self.raw, + NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept as u32, + name.as_ptr(), + priority, + false, + safe_buffered_trampoline::, + user_data, + Some(drop_safe_v2_callback::), + ) + }; + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(super::status_error( + &host.v3.v1, + status, + "native API v2 buffered LLM registration", + )) + } + } + + /// Registers a safe asynchronous native API v2 streaming LLM callback. + /// + /// The callback may return a Rust stream or request host-owned direct + /// pass-through. Relay cooperatively polls its future and returned stream + /// and wakes them when bounded output backpressure clears. + /// Both must remain executor-neutral because host runtime-local state does + /// not cross the dynamic-library boundary. + pub fn register_async_llm_stream_execution_v2( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: Fn(String, LlmRequest, LlmStreamContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let host = self + .host_api_v4() + .copied() + .ok_or_else(|| "native API v2 requires a complete ABI-v4 host table".to_string())?; + let name = HostString::try_new(&host.v3.v1, name).map_err(|status| { + format!("native API v2 streaming LLM registration name failed: {status:?}") + })?; + let state = Box::new(SafeV2Callback { + host, + callback: Arc::new(callback), + }); + let user_data = Box::into_raw(state).cast::(); + let status = unsafe { + // Once invoked, the host consumes `user_data` on both success and + // failure and calls `free_fn` exactly once. Allocate the fallible + // name first so local ownership is never ambiguous. + (host.v3.plugin_context_register_async_stream_middleware)( + self.raw, + name.as_ptr(), + priority, + safe_streaming_trampoline::, + user_data, + Some(drop_safe_v2_callback::), + ) + }; + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(super::status_error( + &host.v3.v1, + status, + "native API v2 streaming LLM registration", + )) + } + } +} + +unsafe extern "C" fn safe_buffered_trampoline( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 +where + F: Fn(String, LlmRequest, LlmContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + if user_data.is_null() { + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + let state = unsafe { &*user_data.cast::>() }; + if completion.is_null() { + if !next.is_null() { + unsafe { (state.host.v3.async_next_release)(next) }; + } + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + let continuation = match unsafe { LlmContinuationV2::from_raw(state.host, next) } { + Ok(continuation) => continuation, + Err(status) => { + reject_completion( + &state.host, + completion, + &format!("invalid native API v2 LLM continuation: {status:?}"), + ); + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + }; + let invocation: LlmCallbackInvocation = match read_json_value( + &state.host.v3.v1, + invocation_json, + "native API v2 LLM invocation", + ) { + Ok(invocation) => invocation, + Err(status) => { + reject_completion( + &state.host, + completion, + &format!("invalid native API v2 LLM invocation: {status:?}"), + ); + drop(continuation); + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + }; + let host = state.host; + let completion_handle = CompletionHandle(completion); + let callback = Arc::clone(&state.callback); + let future = async move { + let result = std::panic::AssertUnwindSafe(async move { + callback(invocation.name, invocation.request, continuation).await + }) + .catch_unwind() + .await; + match result { + Ok(Ok(value)) => resolve_completion(&host, completion_handle.as_ptr(), &value), + Ok(Err(error)) => reject_completion(&host, completion_handle.as_ptr(), &error), + Err(_) => reject_completion( + &host, + completion_handle.as_ptr(), + "native API v2 buffered LLM callback panicked", + ), + } + }; + let task = Box::new(HostFutureTask { + host, + cancellation: HostTaskCancellation::Completion(completion), + completion_to_release: Some(completion), + future: Some(Box::pin(future)), + waker: None, + }); + let task = Box::into_raw(task).cast::(); + let status = unsafe { + (host.async_completion_spawn_task_v2)( + completion, + poll_host_future_task, + task, + Some(drop_host_future_task), + ) + }; + if status == NemoRelayStatus::Ok { + NemoRelayNativeAsyncCallbackState::Pending as u32 + } else { + let mut task = unsafe { Box::from_raw(task.cast::()) }; + task.completion_to_release = None; + drop(task); + reject_completion( + &host, + completion, + &format!("native API v2 buffered task spawn failed: {status:?}"), + ); + NemoRelayNativeAsyncCallbackState::Complete as u32 + } +} + +unsafe extern "C" fn safe_streaming_trampoline( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + output: *const NemoRelayNativeAsyncStream, +) -> u32 +where + F: Fn(String, LlmRequest, LlmStreamContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + if user_data.is_null() { + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + let state = unsafe { &*user_data.cast::>() }; + if output.is_null() { + if !next.is_null() { + unsafe { (state.host.v3.async_next_release)(next) }; + } + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + let continuation = match unsafe { LlmStreamContinuationV2::from_raw(state.host, next, output) } + { + Ok(continuation) => continuation, + Err(status) => { + reject_output_once( + &state.host, + output, + &format!("invalid native API v2 stream continuation: {status:?}"), + ); + unsafe { (state.host.v3.async_stream_release)(output) }; + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + }; + let invocation: LlmCallbackInvocation = match read_json_value( + &state.host.v3.v1, + invocation_json, + "native API v2 streaming LLM invocation", + ) { + Ok(invocation) => invocation, + Err(status) => { + reject_output_once( + &state.host, + output, + &format!("invalid native API v2 stream invocation: {status:?}"), + ); + drop(continuation); + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + }; + let host = state.host; + let output_handle = OutputHandle(output); + let callback = Arc::clone(&state.callback); + let callback_continuation = continuation.clone(); + let future = async move { + let result = std::panic::AssertUnwindSafe(async move { + let outcome = + callback(invocation.name, invocation.request, callback_continuation).await?; + match outcome { + LlmStreamExecutionOutcomeV2::Stream(stream) => { + pump_output_stream(&continuation, stream).await + } + LlmStreamExecutionOutcomeV2::Passthrough(request) => { + continuation.forward_passthrough(request).await + } + } + }) + .catch_unwind() + .await; + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => reject_output(&host, output_handle, &error).await, + Err(_) => { + reject_output( + &host, + output_handle, + "native API v2 streaming LLM callback panicked", + ) + .await; + } + } + }; + let task = Box::new(HostFutureTask { + host, + cancellation: HostTaskCancellation::Stream(output), + completion_to_release: None, + future: Some(Box::pin(future)), + waker: None, + }); + let task = Box::into_raw(task).cast::(); + let status = unsafe { + (host.async_stream_spawn_task_v2)( + output, + poll_host_future_task, + task, + Some(drop_host_future_task), + ) + }; + if status == NemoRelayStatus::Ok { + NemoRelayNativeAsyncCallbackState::Pending as u32 + } else { + reject_output_once( + &host, + output, + &format!("native API v2 streaming task spawn failed: {status:?}"), + ); + unsafe { drop(Box::from_raw(task.cast::())) }; + NemoRelayNativeAsyncCallbackState::Complete as u32 + } +} + +async fn pump_output_stream( + continuation: &LlmStreamContinuationV2, + mut stream: LlmJsonAsyncStreamV2, +) -> Result<()> { + while let Some(item) = stream.next().await { + let chunk = item?; + push_output_json(continuation, &chunk).await?; + } + finish_output(continuation) +} + +async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { + let host = &continuation.inner.host; + let chunk = HostString::from_json(&host.v3.v1, chunk) + .ok_or_else(|| "failed to serialize native API v2 output chunk".to_string())?; + std::future::poll_fn(move |_context| { + if unsafe { (host.v3.async_stream_is_cancelled)(continuation.inner.output) } { + return Poll::Ready(Err("native API v2 output stream was cancelled".into())); + } + let status = + unsafe { (host.v3.async_stream_push_json)(continuation.inner.output, chunk.as_ptr()) }; + match status { + NemoRelayStatus::Ok => Poll::Ready(Ok(())), + NemoRelayStatus::WouldBlock => Poll::Pending, + status => Poll::Ready(Err(format!("native API v2 output push failed: {status:?}"))), + } + }) + .await +} + +fn finish_output(continuation: &LlmStreamContinuationV2) -> Result<()> { + let host = &continuation.inner.host; + if unsafe { (host.v3.async_stream_is_cancelled)(continuation.inner.output) } { + return Err("native API v2 output stream was cancelled".into()); + } + let status = unsafe { (host.v3.async_stream_finish)(continuation.inner.output) }; + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(format!("native API v2 output finish failed: {status:?}")) + } +} + +async fn reject_output(host: &NemoRelayNativeHostApiV4, output: OutputHandle, error: &str) { + if unsafe { (host.v3.async_stream_is_cancelled)(output.as_ptr()) } { + return; + } + let error = bounded_error(error); + let Some(message) = HostString::new(&host.v3.v1, &error) else { + set_last_error( + &host.v3.v1, + "failed to allocate native API v2 stream rejection", + ); + return; + }; + std::future::poll_fn(move |_context| { + if unsafe { (host.v3.async_stream_is_cancelled)(output.as_ptr()) } { + return Poll::Ready(()); + } + let status = unsafe { (host.v3.async_stream_reject)(output.as_ptr(), message.as_ptr()) }; + match status { + NemoRelayStatus::Ok => Poll::Ready(()), + NemoRelayStatus::WouldBlock => Poll::Pending, + status => { + set_last_error( + &host.v3.v1, + &format!("native API v2 output rejection failed: {status:?}"), + ); + Poll::Ready(()) + } + } + }) + .await +} + +fn reject_output_once( + host: &NemoRelayNativeHostApiV4, + output: *const NemoRelayNativeAsyncStream, + error: &str, +) { + if unsafe { (host.v3.async_stream_is_cancelled)(output) } { + return; + } + let error = bounded_error(error); + let Some(message) = HostString::new(&host.v3.v1, &error) else { + set_last_error( + &host.v3.v1, + "failed to allocate native API v2 stream rejection", + ); + return; + }; + let status = unsafe { (host.v3.async_stream_reject)(output, message.as_ptr()) }; + if status != NemoRelayStatus::Ok { + set_last_error( + &host.v3.v1, + &format!("native API v2 output rejection failed: {status:?}"), + ); + } +} + +fn resolve_completion( + host: &NemoRelayNativeHostApiV4, + completion: *const NemoRelayNativeAsyncCompletion, + value: &Json, +) { + if unsafe { (host.v3.async_completion_is_cancelled)(completion) } { + return; + } + let Some(value) = HostString::from_json(&host.v3.v1, value) else { + reject_completion( + host, + completion, + "failed to serialize native API v2 callback result", + ); + return; + }; + let status = unsafe { (host.v3.async_completion_resolve_json)(completion, value.as_ptr()) }; + if status != NemoRelayStatus::Ok { + set_last_error( + &host.v3.v1, + &format!("native API v2 callback completion failed: {status:?}"), + ); + } +} + +fn reject_completion( + host: &NemoRelayNativeHostApiV4, + completion: *const NemoRelayNativeAsyncCompletion, + error: &str, +) { + if unsafe { (host.v3.async_completion_is_cancelled)(completion) } { + return; + } + let error = bounded_error(error); + let Some(error) = HostString::new(&host.v3.v1, &error) else { + set_last_error( + &host.v3.v1, + "failed to allocate native API v2 callback rejection", + ); + return; + }; + let status = unsafe { (host.v3.async_completion_reject)(completion, error.as_ptr()) }; + if status != NemoRelayStatus::Ok { + set_last_error( + &host.v3.v1, + &format!("native API v2 callback rejection failed: {status:?}"), + ); + } +} + +struct TargetedResultCallback { + host: NemoRelayNativeHostApiV1, + sender: oneshot::Sender>, +} + +unsafe extern "C" fn targeted_result_callback( + user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, +) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match (response_json.is_null(), error_json.is_null()) { + (false, true) => read_json_value( + &state.host, + response_json, + "targeted LLM continuation response", + ) + .map_err(|status| status_failure("targeted LLM continuation response", status)), + (true, false) => { + read_json_value(&state.host, error_json, "targeted LLM continuation failure") + .map_or_else( + |status| Err(status_failure("targeted LLM continuation failure", status)), + Err, + ) + } + _ => Err(internal_failure( + "targeted LLM continuation returned an invalid outcome", + )), + } + })) + .unwrap_or_else(|_| Err(internal_failure("targeted LLM result callback panicked"))); + let _ = state.sender.send(result); +} + +struct PassthroughResultCallback { + host: NemoRelayNativeHostApiV1, + sender: oneshot::Sender>, +} + +unsafe extern "C" fn passthrough_result_callback( + user_data: *mut c_void, + value_json: *const NemoRelayNativeString, + error: *const NemoRelayNativeString, +) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + if !error.is_null() { + return read_required_host_string(&state.host, error, "LLM pass-through error") + .map_or_else( + |status| Err(format!("invalid LLM pass-through error: {status:?}")), + Err, + ); + } + read_json_value(&state.host, value_json, "LLM pass-through result") + .map_err(|status| format!("invalid LLM pass-through result: {status:?}")) + })) + .unwrap_or_else(|_| Err("LLM pass-through result callback panicked".into())); + let _ = state.sender.send(result); +} + +struct StreamOpenCallback { + host: NemoRelayNativeHostApiV4, + sender: oneshot::Sender>, +} + +unsafe extern "C" fn stream_open_callback( + user_data: *mut c_void, + stream: *const NemoRelayNativeLlmStreamV2, + error_json: *const NemoRelayNativeString, +) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let mut provider = (!stream.is_null()).then(|| LlmProviderStreamV2 { + host: state.host, + raw: stream, + pending: None, + finished: false, + }); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match (stream.is_null(), error_json.is_null()) { + (false, true) => Ok(()), + (true, false) => read_json_value( + &state.host.v3.v1, + error_json, + "targeted LLM stream setup failure", + ) + .map_or_else( + |status| Err(status_failure("targeted LLM stream setup failure", status)), + Err, + ), + _ => Err(internal_failure( + "targeted LLM stream setup returned an invalid outcome", + )), + } + })) + .unwrap_or_else(|_| { + Err(internal_failure( + "targeted LLM stream setup callback panicked", + )) + }); + let result = result.and_then(|()| { + provider + .take() + .ok_or_else(|| internal_failure("targeted LLM stream setup returned a null stream")) + }); + let _ = state.sender.send(result); +} + +struct ProviderNextCallback { + host: NemoRelayNativeHostApiV1, + sender: oneshot::Sender, LlmContinuationFailureV2>>, +} + +unsafe extern "C" fn provider_next_callback( + user_data: *mut c_void, + chunk_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, + done: bool, +) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match (chunk_json.is_null(), error_json.is_null(), done) { + (false, true, false) => { + read_json_value(&state.host, chunk_json, "targeted provider stream chunk") + .map(Some) + .map_err(|status| status_failure("targeted provider stream chunk", status)) + } + (true, true, true) => Ok(None), + (true, false, true) => { + read_json_value(&state.host, error_json, "targeted provider stream failure") + .map_err(|status| status_failure("targeted provider stream failure", status)) + .and_then(Err) + } + _ => Err(internal_failure( + "targeted provider stream returned an invalid event", + )), + } + })) + .unwrap_or_else(|_| { + Err(internal_failure( + "targeted provider stream callback panicked", + )) + }); + let _ = state.sender.send(result); +} + +struct ForwardStreamCallback { + sender: oneshot::Sender<()>, +} + +unsafe extern "C" fn forward_stream_callback(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let _ = state.sender.send(()); +} + +fn status_failure(label: &str, status: NemoRelayStatus) -> LlmContinuationFailureV2 { + internal_failure(format!("{label} failed at the native boundary: {status:?}")) +} + +fn internal_failure(message: impl Into) -> LlmContinuationFailureV2 { + LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + message: bounded_error(&message.into()), + } +} + +fn bounded_error(message: &str) -> String { + if message.len() <= MAX_SDK_ERROR_BYTES { + return message.to_owned(); + } + let mut boundary = MAX_SDK_ERROR_BYTES; + while !message.is_char_boundary(boundary) { + boundary -= 1; + } + message[..boundary].to_owned() +} diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 0161400fb..9714bde4b 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -3,34 +3,45 @@ //! Public-API tests for typed native plugin callback registration. -use std::collections::VecDeque; +use std::collections::{BTreeMap, VecDeque}; use std::ffi::c_void; +use std::future::Future; use std::mem::{align_of, offset_of, size_of}; use std::ptr::{self, NonNull}; use std::sync::{ Arc, Mutex, MutexGuard, - atomic::{AtomicUsize, Ordering}, + atomic::{AtomicBool, AtomicUsize, Ordering}, }; +use futures::{StreamExt, stream}; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, - Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmJsonStream, LlmNext, - LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, - NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, NemoRelayNativeAsyncCallbackState, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeEventSanitizeCb, - NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, + Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmContinuationFailureV2, + LlmContinuationInvocationV2, LlmContinuationTargetV2, LlmContinuationV2, LlmJsonStream, + LlmNext, LlmNonHttpFailureKindV2, LlmRequest, LlmRequestInterceptOutcome, LlmStream, + LlmStreamContinuationV2, LlmStreamExecutionOutcomeV2, LlmStreamNext, + NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, + NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, + NemoRelayNativeAsyncLlmResultCbV2, NemoRelayNativeAsyncLlmStreamForwardCbV2, + NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeAsyncTaskPollCbV2, + NemoRelayNativeAsyncTaskV2, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, + NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, - NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, - NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, - NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, - NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, - NemoRelayNativeWithScopeStackCb, NemoRelayStatus, PendingMarkSpec, PluginContext, - PluginRuntime, ScopeType, ToolExecutionInterceptOutcome, ToolNext, + NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativeLlmStreamV2, + NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, + NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, + NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, + NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, PendingMarkSpec, + PluginContext, PluginRuntime, ScopeType, ToolExecutionInterceptOutcome, ToolNext, }; +use serde::de::DeserializeOwned; use serde_json::{Map, json}; #[test] @@ -236,6 +247,38 @@ struct RegisteredLlmRequestIntercept { free_fn: NemoRelayNativeFreeFn, } +struct RegisteredAsyncV2 { + name: String, + priority: i32, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} + +impl RegisteredAsyncV2 { + unsafe fn free(self) { + if let Some(free_fn) = self.free_fn { + unsafe { free_fn(self.user_data as *mut c_void) }; + } + } +} + +struct RegisteredAsyncStreamV2 { + name: String, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} + +impl RegisteredAsyncStreamV2 { + unsafe fn free(self) { + if let Some(free_fn) = self.free_fn { + unsafe { free_fn(self.user_data as *mut c_void) }; + } + } +} + impl RegisteredLlmRequestIntercept { unsafe fn free(self) { if let Some(free_fn) = self.free_fn { @@ -272,6 +315,8 @@ impl_captured_registration!( RegisteredLlmExecution, RegisteredLlmStreamExecution, RegisteredLlmRequestIntercept, + RegisteredAsyncV2, + RegisteredAsyncStreamV2, ); fn replace_registration(slot: &Mutex>, registration: T) { @@ -331,10 +376,75 @@ static LLM_STREAM_EXECUTION_REGISTRATION: Mutex> = Mutex::new(None); +static ASYNC_V2_REGISTRATION: Mutex> = Mutex::new(None); +static ASYNC_STREAM_V2_REGISTRATION: Mutex> = Mutex::new(None); +static SAFE_V2_COMPLETION: Mutex>> = Mutex::new(None); +static SAFE_V2_COMPLETION_CANCELLED: AtomicBool = AtomicBool::new(false); +static SAFE_V2_OUTPUT: Mutex>> = Mutex::new(Vec::new()); +#[derive(Clone)] +enum SafeV2ProviderEvent { + Chunk(Json), + Done, + Failure(LlmContinuationFailureV2), +} + +static SAFE_V2_PROVIDER_EVENTS: Mutex> = Mutex::new(VecDeque::new()); +static SAFE_V2_OPEN_FAILURE: Mutex> = Mutex::new(None); +static SAFE_V2_OPEN_RETURNS_STREAM_AND_ERROR: AtomicBool = AtomicBool::new(false); +static SAFE_V2_HOLD_STREAM_OPEN_CALLBACK: AtomicBool = AtomicBool::new(false); +static SAFE_V2_HELD_STREAM_OPEN_CALLBACK: Mutex< + Option<(NemoRelayNativeAsyncLlmStreamOpenCbV2, usize)>, +> = Mutex::new(None); +static SAFE_V2_FORWARDED_REQUESTS: Mutex> = Mutex::new(Vec::new()); +static SAFE_V2_HOLD_TARGETED_CALLBACK: AtomicBool = AtomicBool::new(false); +static SAFE_V2_HELD_TARGETED_CALLBACK: Mutex> = + Mutex::new(None); +static SAFE_V2_NEXT_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_COMPLETION_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_OUTPUT_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_PROVIDER_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_OUTPUT_FINISHES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_OUTPUT_CANCELLED: AtomicBool = AtomicBool::new(false); +static SAFE_V2_REGISTRATION_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_TARGETED_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_TARGETED_FAILURE: Mutex> = Mutex::new(None); +static SAFE_V2_TARGETED_INVALID_OUTCOME: AtomicBool = AtomicBool::new(false); +static SAFE_V2_PASSTHROUGH_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_PASSTHROUGH_ERROR: Mutex> = Mutex::new(None); +static SAFE_V2_STREAM_OPEN_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_PROVIDER_NEXT_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_PROVIDER_EVENT_JSON: Mutex> = Mutex::new(None); +static SAFE_V2_PROVIDER_INVALID_EVENT: AtomicBool = AtomicBool::new(false); +static SAFE_V2_FORWARD_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_COMPLETION_RESOLVE_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_COMPLETION_REJECT_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_OUTPUT_PUSH_STATUSES: Mutex> = Mutex::new(VecDeque::new()); +static SAFE_V2_OUTPUT_FINISH_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_OUTPUT_REJECT_STATUSES: Mutex> = + Mutex::new(VecDeque::new()); +static SAFE_V2_TASKS: Mutex> = Mutex::new(Vec::new()); +static SAFE_V2_TASK_SPAWN_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_TASK_RETAINS: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_TASK_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_HELD_TASK_WAKER: Mutex> = Mutex::new(None); + +thread_local! { + static SAFE_V2_CURRENT_TASK: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +struct SafeV2Task { + refs: AtomicUsize, + woken: AtomicBool, + completed: AtomicBool, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} #[test] -fn native_abi_v3_struct_sizes_are_self_describing() { +fn native_abi_struct_sizes_are_self_describing() { assert_eq!(NEMO_RELAY_NATIVE_ABI_VERSION, 3); + assert_eq!(NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, 4); assert_eq!( size_of::(), test_host().struct_size @@ -348,62 +458,92 @@ fn native_abi_v3_struct_sizes_are_self_describing() { NemoRelayNativeLlmStreamV1::default().struct_size ); assert_eq!(NemoRelayStatus::StreamEnd as i32, 10); + assert_eq!(NemoRelayStatus::WouldBlock as i32, 11); - #[cfg(target_pointer_width = "64")] - { - assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 320); - assert_eq!( - host_api_offsets(), - [ - 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, - 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, - 280, 288, 296, 304, 312, - ] - ); - assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 440); - assert_eq!( - host_api_v3_offsets(), - [ - 0, 320, 328, 336, 344, 352, 360, 368, 376, 384, 392, 400, 408, 416, 424, 432 - ] - ); - assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 56); - assert_eq!(plugin_offsets(), [0, 8, 16, 24, 32, 40, 48]); - assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 40); - assert_eq!(stream_offsets(), [0, 8, 16, 24, 32]); - } + assert_native_abi_target_layout(); +} - #[cfg(target_pointer_width = "32")] - { - assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 160); - assert_eq!( - host_api_offsets(), - [ - 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, - 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, - 152, 156, - ] - ); - assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 216); - assert_eq!( - host_api_v3_offsets(), - [ - 0, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212 - ] - ); - assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 28); - assert_eq!(plugin_offsets(), [0, 4, 8, 12, 16, 20, 24]); - assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 20); - assert_eq!(stream_offsets(), [0, 4, 8, 12, 16]); - } +#[cfg(target_pointer_width = "64")] +fn assert_native_abi_target_layout() { + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 320); + assert_eq!( + host_api_offsets(), + [ + 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, + 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, 280, 288, + 296, 304, 312, + ] + ); + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 440); + assert_eq!( + host_api_v3_offsets(), + [ + 0, 320, 328, 336, 344, 352, 360, 368, 376, 384, 392, 400, 408, 416, 424, 432 + ] + ); + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 520); + assert_eq!( + host_api_v4_offsets(), + [0, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512] + ); + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 56); + assert_eq!(plugin_offsets(), [0, 8, 16, 24, 32, 40, 48]); + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 40); + assert_eq!(stream_offsets(), [0, 8, 16, 24, 32]); +} + +#[cfg(target_pointer_width = "32")] +fn assert_native_abi_target_layout() { + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 160); + assert_eq!( + host_api_offsets(), + [ + 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, + 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, + ] + ); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 220); + assert_eq!( + host_api_v3_offsets(), + [ + 0, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216 + ] + ); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 260); + assert_eq!( + host_api_v4_offsets(), + [0, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256] + ); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 28); + assert_eq!(plugin_offsets(), [0, 4, 8, 12, 16, 20, 24]); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 20); + assert_eq!(stream_offsets(), [0, 4, 8, 12, 16]); +} + +fn host_api_v4_offsets() -> [usize; 11] { + [ + offset_of!(NemoRelayNativeHostApiV4, v3), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_result_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_open_stream_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_next_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_release_v2), + offset_of!(NemoRelayNativeHostApiV4, async_completion_spawn_task_v2), + offset_of!(NemoRelayNativeHostApiV4, async_stream_spawn_task_v2), + offset_of!(NemoRelayNativeHostApiV4, async_task_retain_v2), + offset_of!(NemoRelayNativeHostApiV4, async_task_wake_v2), + offset_of!(NemoRelayNativeHostApiV4, async_task_release_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_forward_stream_v2), + ] } fn host_api_v3_offsets() -> [usize; 16] { @@ -1325,6 +1465,12 @@ fn reset_state() { clear_registration(&LLM_EXECUTION_REGISTRATION); clear_registration(&LLM_STREAM_EXECUTION_REGISTRATION); clear_registration(&LLM_REQUEST_INTERCEPT_REGISTRATION); + clear_registration(&ASYNC_V2_REGISTRATION); + clear_registration(&ASYNC_STREAM_V2_REGISTRATION); + assert!( + SAFE_V2_TASKS.lock().unwrap().is_empty(), + "previous test leaked cooperative host tasks" + ); assert_eq!( STRING_LIVE_COUNT.load(Ordering::SeqCst), 0, @@ -1352,6 +1498,43 @@ fn reset_state() { SCOPE_STACK_FREES.store(0, Ordering::SeqCst); SCOPE_STACK_BINDING_FREES.store(0, Ordering::SeqCst); SCOPE_STACK_BINDING_RESTORES.store(0, Ordering::SeqCst); + *SAFE_V2_COMPLETION.lock().unwrap() = None; + SAFE_V2_COMPLETION_CANCELLED.store(false, Ordering::SeqCst); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + SAFE_V2_PROVIDER_EVENTS.lock().unwrap().clear(); + *SAFE_V2_OPEN_FAILURE.lock().unwrap() = None; + SAFE_V2_OPEN_RETURNS_STREAM_AND_ERROR.store(false, Ordering::SeqCst); + SAFE_V2_HOLD_STREAM_OPEN_CALLBACK.store(false, Ordering::SeqCst); + assert!(SAFE_V2_HELD_STREAM_OPEN_CALLBACK.lock().unwrap().is_none()); + SAFE_V2_FORWARDED_REQUESTS.lock().unwrap().clear(); + SAFE_V2_HOLD_TARGETED_CALLBACK.store(false, Ordering::SeqCst); + assert!(SAFE_V2_HELD_TARGETED_CALLBACK.lock().unwrap().is_none()); + SAFE_V2_NEXT_RELEASES.store(0, Ordering::SeqCst); + SAFE_V2_COMPLETION_RELEASES.store(0, Ordering::SeqCst); + SAFE_V2_OUTPUT_RELEASES.store(0, Ordering::SeqCst); + SAFE_V2_PROVIDER_RELEASES.store(0, Ordering::SeqCst); + SAFE_V2_OUTPUT_FINISHES.store(0, Ordering::SeqCst); + SAFE_V2_OUTPUT_CANCELLED.store(false, Ordering::SeqCst); + *SAFE_V2_REGISTRATION_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_TARGETED_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_TARGETED_FAILURE.lock().unwrap() = None; + SAFE_V2_TARGETED_INVALID_OUTCOME.store(false, Ordering::SeqCst); + *SAFE_V2_PASSTHROUGH_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_PASSTHROUGH_ERROR.lock().unwrap() = None; + *SAFE_V2_STREAM_OPEN_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_PROVIDER_NEXT_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_PROVIDER_EVENT_JSON.lock().unwrap() = None; + SAFE_V2_PROVIDER_INVALID_EVENT.store(false, Ordering::SeqCst); + *SAFE_V2_FORWARD_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_COMPLETION_RESOLVE_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_COMPLETION_REJECT_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT_PUSH_STATUSES.lock().unwrap().clear(); + *SAFE_V2_OUTPUT_FINISH_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT_REJECT_STATUSES.lock().unwrap().clear(); + *SAFE_V2_TASK_SPAWN_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_TASK_RETAINS.store(0, Ordering::SeqCst); + SAFE_V2_TASK_RELEASES.store(0, Ordering::SeqCst); + assert!(SAFE_V2_HELD_TASK_WAKER.lock().unwrap().is_none()); } fn test_context(host: &NemoRelayNativeHostApiV1) -> PluginContext<'_> { @@ -1393,6 +1576,14 @@ fn required_host_string( read_host_string(host, value).ok_or(NemoRelayStatus::InvalidArg) } +fn required_host_json( + host: &NemoRelayNativeHostApiV1, + value: *const NemoRelayNativeString, +) -> std::result::Result { + let value = required_host_string(host, value)?; + serde_json::from_str(&value).map_err(|_| NemoRelayStatus::InvalidJson) +} + fn optional_host_string( host: &NemoRelayNativeHostApiV1, value: *const NemoRelayNativeString, @@ -5961,3 +6152,2145 @@ fn plugin_validate_and_register_panics_replace_last_error() { drop_exported_plugin(&host, register_plugin); } } + +unsafe extern "C" fn safe_v2_completion_resolve( + _completion: *const NemoRelayNativeAsyncCompletion, + value_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + let host = test_host(); + let value = match required_host_json(&host, value_json) { + Ok(value) => value, + Err(status) => return status, + }; + *SAFE_V2_COMPLETION.lock().unwrap() = Some(Ok(value)); + *SAFE_V2_COMPLETION_RESOLVE_STATUS.lock().unwrap() +} + +unsafe extern "C" fn safe_v2_completion_reject( + _completion: *const NemoRelayNativeAsyncCompletion, + message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + let host = test_host(); + let message = required_host_string(&host, message) + .unwrap_or_else(|status| format!("invalid rejection: {status:?}")); + *SAFE_V2_COMPLETION.lock().unwrap() = Some(Err(message)); + *SAFE_V2_COMPLETION_REJECT_STATUS.lock().unwrap() +} + +unsafe extern "C" fn safe_v2_completion_is_cancelled( + _completion: *const NemoRelayNativeAsyncCompletion, +) -> bool { + SAFE_V2_COMPLETION_CANCELLED.load(Ordering::SeqCst) +} + +unsafe extern "C" fn safe_v2_completion_release( + _completion: *const NemoRelayNativeAsyncCompletion, +) { + SAFE_V2_COMPLETION_RELEASES.fetch_add(1, Ordering::SeqCst); +} + +unsafe extern "C" fn safe_v2_async_next_invoke( + _next: *const NemoRelayNativeAsyncNext, + _invocation_json: *const NemoRelayNativeString, + _completion: *const NemoRelayNativeAsyncCompletion, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn safe_v2_next_release(_next: *const NemoRelayNativeAsyncNext) { + SAFE_V2_NEXT_RELEASES.fetch_add(1, Ordering::SeqCst); +} + +unsafe fn safe_v2_reject_registration( + status: NemoRelayStatus, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + status +} + +fn safe_v2_registration_name( + name: *const NemoRelayNativeString, +) -> std::result::Result { + let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return Err(status); + } + required_host_string(&test_host(), name) +} + +unsafe extern "C" fn safe_v2_register_generic_async( + _ctx: *mut NemoRelayNativePluginContext, + kind: u32, + name: *const NemoRelayNativeString, + priority: i32, + _break_chain: bool, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + if kind != NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept as u32 { + return unsafe { + safe_v2_reject_registration(NemoRelayStatus::InvalidArg, user_data, free_fn) + }; + } + let name = match safe_v2_registration_name(name) { + Ok(name) => name, + Err(status) => return unsafe { safe_v2_reject_registration(status, user_data, free_fn) }, + }; + replace_registration( + &ASYNC_V2_REGISTRATION, + RegisteredAsyncV2 { + name, + priority, + cb, + user_data: user_data as usize, + free_fn, + }, + ); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_stream_push( + _stream: *const NemoRelayNativeAsyncStream, + chunk_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + let host = test_host(); + let chunk = match required_host_json(&host, chunk_json) { + Ok(chunk) => chunk, + Err(status) => return status, + }; + let status = safe_v2_output_status(&SAFE_V2_OUTPUT_PUSH_STATUSES); + if status == NemoRelayStatus::Ok { + SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); + } + status +} + +fn safe_v2_output_status(statuses: &Mutex>) -> NemoRelayStatus { + let status = statuses + .lock() + .unwrap() + .pop_front() + .unwrap_or(NemoRelayStatus::Ok); + if status == NemoRelayStatus::WouldBlock { + SAFE_V2_CURRENT_TASK.with(|current| { + let task = current.get(); + if task != 0 { + unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + } + }); + } + status +} + +unsafe extern "C" fn safe_v2_stream_finish( + _stream: *const NemoRelayNativeAsyncStream, +) -> NemoRelayStatus { + SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); + *SAFE_V2_OUTPUT_FINISH_STATUS.lock().unwrap() +} + +unsafe extern "C" fn safe_v2_stream_reject( + _stream: *const NemoRelayNativeAsyncStream, + message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + let host = test_host(); + let message = required_host_string(&host, message) + .unwrap_or_else(|status| format!("invalid rejection: {status:?}")); + let status = safe_v2_output_status(&SAFE_V2_OUTPUT_REJECT_STATUSES); + if status == NemoRelayStatus::Ok { + SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); + } + status +} + +unsafe extern "C" fn safe_v2_stream_is_cancelled( + _stream: *const NemoRelayNativeAsyncStream, +) -> bool { + SAFE_V2_OUTPUT_CANCELLED.load(Ordering::SeqCst) +} + +unsafe extern "C" fn safe_v2_stream_release(_stream: *const NemoRelayNativeAsyncStream) { + SAFE_V2_OUTPUT_RELEASES.fetch_add(1, Ordering::SeqCst); +} + +unsafe extern "C" fn safe_v2_async_next_invoke_stream( + _next: *const NemoRelayNativeAsyncNext, + _invocation_json: *const NemoRelayNativeString, + _stream: *const NemoRelayNativeAsyncStream, + _cb: NemoRelayNativeAsyncNextStreamCb, + _user_data: *mut c_void, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn safe_v2_register_generic_stream( + _ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let name = match safe_v2_registration_name(name) { + Ok(name) => name, + Err(status) => return unsafe { safe_v2_reject_registration(status, user_data, free_fn) }, + }; + replace_registration( + &ASYNC_STREAM_V2_REGISTRATION, + RegisteredAsyncStreamV2 { + name, + priority, + cb, + user_data: user_data as usize, + free_fn, + }, + ); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_passthrough_result( + _next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncNextResultCb, + user_data: *mut c_void, +) -> NemoRelayStatus { + let status = *SAFE_V2_PASSTHROUGH_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let host = test_host(); + if let Err(status) = required_host_json::(&host, invocation_json) { + return status; + } + if let Some(error) = SAFE_V2_PASSTHROUGH_ERROR.lock().unwrap().clone() { + let error = host_string(&host, &error); + unsafe { cb(user_data, ptr::null(), error) }; + unsafe { (host.string_free)(error) }; + } else { + let value = json_host_string(&host, json!({ "passthrough": true })); + unsafe { cb(user_data, value, ptr::null()) }; + unsafe { (host.string_free)(value) }; + } + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_targeted_result( + _next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let status = *SAFE_V2_TARGETED_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let host = test_host(); + if let Err(status) = required_host_json::(&host, invocation_json) { + return status; + } + if SAFE_V2_HOLD_TARGETED_CALLBACK.load(Ordering::SeqCst) { + *SAFE_V2_HELD_TARGETED_CALLBACK.lock().unwrap() = Some((cb, user_data as usize)); + return NemoRelayStatus::Ok; + } + if SAFE_V2_TARGETED_INVALID_OUTCOME.swap(false, Ordering::SeqCst) { + unsafe { cb(user_data, ptr::null(), ptr::null()) }; + return NemoRelayStatus::Ok; + } + if let Some(error) = SAFE_V2_TARGETED_FAILURE.lock().unwrap().take() { + let error = json_host_string(&host, serde_json::to_value(error).unwrap()); + unsafe { cb(user_data, ptr::null(), error) }; + unsafe { (host.string_free)(error) }; + return NemoRelayStatus::Ok; + } + let response = json_host_string(&host, json!({ "targeted": true })); + unsafe { cb(user_data, response, ptr::null()) }; + unsafe { (host.string_free)(response) }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_stream_open( + _next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + _output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let status = *SAFE_V2_STREAM_OPEN_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let host = test_host(); + if let Err(status) = required_host_json::(&host, invocation_json) { + return status; + } + if SAFE_V2_HOLD_STREAM_OPEN_CALLBACK.load(Ordering::SeqCst) { + *SAFE_V2_HELD_STREAM_OPEN_CALLBACK.lock().unwrap() = Some((cb, user_data as usize)); + return NemoRelayStatus::Ok; + } + if let Some(error) = SAFE_V2_OPEN_FAILURE.lock().unwrap().clone() { + let error = json_host_string(&host, serde_json::to_value(error).unwrap()); + let stream = if SAFE_V2_OPEN_RETURNS_STREAM_AND_ERROR.load(Ordering::SeqCst) { + NonNull::::dangling().as_ptr() + } else { + ptr::null() + }; + unsafe { cb(user_data, stream, error) }; + unsafe { (host.string_free)(error) }; + return NemoRelayStatus::Ok; + } + unsafe { + cb( + user_data, + NonNull::::dangling().as_ptr(), + ptr::null(), + ) + }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_provider_next( + _stream: *const NemoRelayNativeLlmStreamV2, + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let status = *SAFE_V2_PROVIDER_NEXT_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let host = test_host(); + if let Some(event) = SAFE_V2_PROVIDER_EVENT_JSON.lock().unwrap().take() { + let chunk = host_string(&host, &event); + unsafe { cb(user_data, chunk, ptr::null(), false) }; + unsafe { (host.string_free)(chunk) }; + return NemoRelayStatus::Ok; + } + if SAFE_V2_PROVIDER_INVALID_EVENT.swap(false, Ordering::SeqCst) { + unsafe { cb(user_data, ptr::null(), ptr::null(), false) }; + return NemoRelayStatus::Ok; + } + let event = SAFE_V2_PROVIDER_EVENTS + .lock() + .unwrap() + .pop_front() + .unwrap_or(SafeV2ProviderEvent::Done); + match event { + SafeV2ProviderEvent::Chunk(chunk) => { + let chunk = json_host_string(&host, chunk); + unsafe { cb(user_data, chunk, ptr::null(), false) }; + unsafe { (host.string_free)(chunk) }; + } + SafeV2ProviderEvent::Done => unsafe { cb(user_data, ptr::null(), ptr::null(), true) }, + SafeV2ProviderEvent::Failure(error) => { + let error = json_host_string(&host, serde_json::to_value(error).unwrap()); + unsafe { cb(user_data, ptr::null(), error, true) }; + unsafe { (host.string_free)(error) }; + } + } + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_provider_release(_stream: *const NemoRelayNativeLlmStreamV2) { + SAFE_V2_PROVIDER_RELEASES.fetch_add(1, Ordering::SeqCst); +} + +unsafe extern "C" fn safe_v2_spawn_task( + _owner: *const c_void, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let status = *SAFE_V2_TASK_SPAWN_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let task = Box::new(SafeV2Task { + refs: AtomicUsize::new(1), + woken: AtomicBool::new(true), + completed: AtomicBool::new(false), + cb, + user_data: user_data as usize, + free_fn, + }); + let task = Box::into_raw(task) as usize; + SAFE_V2_TASKS.lock().unwrap().push(task); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_completion_spawn_task( + completion: *const NemoRelayNativeAsyncCompletion, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + unsafe { safe_v2_spawn_task(completion.cast(), cb, user_data, free_fn) } +} + +unsafe extern "C" fn safe_v2_stream_spawn_task( + stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncTaskPollCbV2, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + unsafe { safe_v2_spawn_task(stream.cast(), cb, user_data, free_fn) } +} + +unsafe extern "C" fn safe_v2_task_retain(task: *const NemoRelayNativeAsyncTaskV2) { + if let Some(task) = unsafe { task.cast::().as_ref() } { + task.refs.fetch_add(1, Ordering::Relaxed); + SAFE_V2_TASK_RETAINS.fetch_add(1, Ordering::SeqCst); + } +} + +unsafe extern "C" fn safe_v2_task_wake(task: *const NemoRelayNativeAsyncTaskV2) { + let Some(task) = (unsafe { task.cast::().as_ref() }) else { + return; + }; + if task.completed.load(Ordering::Acquire) { + return; + } + task.woken.store(true, Ordering::Release); +} + +unsafe extern "C" fn safe_v2_task_release(task: *const NemoRelayNativeAsyncTaskV2) { + let Some(task_ref) = (unsafe { task.cast::().as_ref() }) else { + return; + }; + SAFE_V2_TASK_RELEASES.fetch_add(1, Ordering::SeqCst); + if task_ref.refs.fetch_sub(1, Ordering::AcqRel) == 1 { + unsafe { drop(Box::from_raw(task.cast_mut().cast::())) }; + } +} + +fn wake_safe_v2_tasks() { + for task in SAFE_V2_TASKS.lock().unwrap().iter().copied() { + unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + } +} + +fn drive_safe_v2_tasks() { + loop { + let tasks = SAFE_V2_TASKS.lock().unwrap().clone(); + let mut made_progress = false; + for raw in tasks { + let task_ptr = raw as *const SafeV2Task; + let Some(task) = (unsafe { task_ptr.as_ref() }) else { + continue; + }; + if !task.woken.swap(false, Ordering::AcqRel) { + continue; + } + made_progress = true; + unsafe { + safe_v2_task_retain(task_ptr.cast()); + } + SAFE_V2_CURRENT_TASK.with(|current| current.set(raw)); + let state = unsafe { + (task.cb)( + task.user_data as *mut c_void, + task_ptr.cast::(), + ) + }; + SAFE_V2_CURRENT_TASK.with(|current| current.set(0)); + if NemoRelayNativeAsyncCallbackState::try_from(state) + == Ok(NemoRelayNativeAsyncCallbackState::Complete) + { + task.completed.store(true, Ordering::Release); + if let Some(free_fn) = task.free_fn { + unsafe { free_fn(task.user_data as *mut c_void) }; + } + SAFE_V2_TASKS.lock().unwrap().retain(|task| *task != raw); + unsafe { safe_v2_task_release(task_ptr.cast()) }; + } + unsafe { safe_v2_task_release(task_ptr.cast()) }; + } + if !made_progress { + return; + } + } +} + +unsafe extern "C" fn safe_v2_forward_stream( + _next: *const NemoRelayNativeAsyncNext, + request_json: *const NemoRelayNativeString, + _output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let status = *SAFE_V2_FORWARD_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + return status; + } + let host = test_host(); + let request = match required_host_json(&host, request_json) { + Ok(request) => request, + Err(status) => return status, + }; + SAFE_V2_FORWARDED_REQUESTS.lock().unwrap().push(request); + SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); + unsafe { cb(user_data) }; + NemoRelayStatus::Ok +} + +fn test_host_v4() -> NemoRelayNativeHostApiV4 { + let mut v1 = test_host(); + v1.abi_version = NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS; + v1.struct_size = size_of::(); + NemoRelayNativeHostApiV4 { + v3: NemoRelayNativeHostApiV3 { + v1, + async_completion_resolve_json: safe_v2_completion_resolve, + async_completion_reject: safe_v2_completion_reject, + async_completion_is_cancelled: safe_v2_completion_is_cancelled, + async_completion_release: safe_v2_completion_release, + async_next_invoke: safe_v2_async_next_invoke, + async_next_release: safe_v2_next_release, + plugin_context_register_async_middleware: safe_v2_register_generic_async, + async_stream_push_json: safe_v2_stream_push, + async_stream_finish: safe_v2_stream_finish, + async_stream_reject: safe_v2_stream_reject, + async_stream_is_cancelled: safe_v2_stream_is_cancelled, + async_stream_release: safe_v2_stream_release, + async_next_invoke_stream: safe_v2_async_next_invoke_stream, + plugin_context_register_async_stream_middleware: safe_v2_register_generic_stream, + async_next_invoke_result: safe_v2_passthrough_result, + }, + async_llm_next_invoke_result_v2: safe_v2_targeted_result, + async_llm_next_open_stream_v2: safe_v2_stream_open, + async_llm_stream_next_v2: safe_v2_provider_next, + async_llm_stream_release_v2: safe_v2_provider_release, + async_completion_spawn_task_v2: safe_v2_completion_spawn_task, + async_stream_spawn_task_v2: safe_v2_stream_spawn_task, + async_task_retain_v2: safe_v2_task_retain, + async_task_wake_v2: safe_v2_task_wake, + async_task_release_v2: safe_v2_task_release, + async_llm_next_forward_stream_v2: safe_v2_forward_stream, + } +} + +fn safe_v2_target() -> LlmContinuationTargetV2 { + LlmContinuationTargetV2 { + url: "https://provider.example/v1/chat/completions".into(), + headers: Default::default(), + } +} + +#[test] +fn native_v2_debug_output_redacts_requests_targets_and_credentials() { + let invocation = LlmContinuationInvocationV2 { + request: LlmRequest { + headers: Map::from_iter([("authorization".into(), json!("Bearer request-secret"))]), + content: json!({"prompt": "request-body-secret"}), + }, + target: LlmContinuationTargetV2 { + url: "https://provider.example/v1?api_key=url-secret".into(), + headers: BTreeMap::from([("authorization".into(), "Bearer target-secret".into())]), + }, + }; + + let debug = format!("{invocation:?}"); + assert!(debug.contains("LlmContinuationInvocationV2")); + assert!(debug.contains("authorization")); + for secret in [ + "request-secret", + "request-body-secret", + "target-secret", + "url-secret", + ] { + assert!(!debug.contains(secret)); + } +} + +fn take_safe_v2_buffered_registration() -> RegisteredAsyncV2 { + ASYNC_V2_REGISTRATION.lock().unwrap().take().unwrap() +} + +fn take_safe_v2_stream_registration() -> RegisteredAsyncStreamV2 { + ASYNC_STREAM_V2_REGISTRATION.lock().unwrap().take().unwrap() +} + +fn invoke_safe_v2_buffered( + host: &NemoRelayNativeHostApiV4, + registration: &RegisteredAsyncV2, + next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + next, + completion, + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + state +} + +fn invoke_safe_v2_streaming( + host: &NemoRelayNativeHostApiV4, + registration: &RegisteredAsyncStreamV2, + next: *const NemoRelayNativeAsyncNext, + output: *const NemoRelayNativeAsyncStream, +) -> u32 { + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + next, + output, + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + state +} + +fn run_safe_v2_buffered(host: &NemoRelayNativeHostApiV4, name: &str, callback: F) -> u32 +where + F: Fn(String, LlmRequest, LlmContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2(name, 0, callback) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + let state = invoke_safe_v2_buffered( + host, + ®istration, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + unsafe { registration.free() }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + state +} + +fn run_safe_v2_streaming(host: &NemoRelayNativeHostApiV4, name: &str, callback: F) -> u32 +where + F: Fn(String, LlmRequest, LlmStreamContinuationV2) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2(name, 0, callback) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let state = invoke_safe_v2_streaming( + host, + ®istration, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + unsafe { registration.free() }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + state +} + +async fn safe_v2_targeted_provider_stream( + request: LlmRequest, + next: LlmStreamContinuationV2, +) -> std::result::Result { + let provider = next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + provider.map(|item| item.map_err(|error| format!("{error:?}"))), + ))) +} + +fn run_safe_v2_targeted_stream(host: &NemoRelayNativeHostApiV4, name: &str) { + run_safe_v2_streaming(host, name, |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); +} + +fn safe_v2_one_chunk_stream() -> LlmStreamExecutionOutcomeV2 { + LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once(async { + Ok(json!({ "chunk": true })) + }))) +} + +unsafe extern "C" fn raw_v2_buffered_probe( + _user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + _next: *const NemoRelayNativeAsyncNext, + _completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +unsafe extern "C" fn raw_v2_streaming_probe( + _user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + _next: *const NemoRelayNativeAsyncNext, + _output: *const NemoRelayNativeAsyncStream, +) -> u32 { + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +#[test] +fn native_api_v2_uses_generic_raw_registration_as_advanced_escape_hatch() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + + assert_eq!( + unsafe { + ctx.register_async_middleware_raw( + NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept, + "raw-buffered", + 11, + false, + raw_v2_buffered_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::Ok + ); + let buffered = take_safe_v2_buffered_registration(); + assert_eq!( + (buffered.name.as_str(), buffered.priority), + ("raw-buffered", 11) + ); + unsafe { buffered.free() }; + + assert_eq!( + unsafe { + ctx.register_async_stream_middleware_raw( + "raw-streaming", + 12, + raw_v2_streaming_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::Ok + ); + let streaming = take_safe_v2_stream_registration(); + assert_eq!( + (streaming.name.as_str(), streaming.priority), + ("raw-streaming", 12) + ); + unsafe { streaming.free() }; + + let v1 = test_host(); + let mut v1_ctx = test_context(&v1); + assert_eq!( + unsafe { + v1_ctx.register_async_middleware_raw( + NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept, + "unsupported-buffered", + 0, + false, + raw_v2_buffered_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_eq!( + unsafe { + v1_ctx.register_async_stream_middleware_raw( + "unsupported-streaming", + 0, + raw_v2_streaming_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::InvalidArg + ); + + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_registration_wraps_targeted_and_passthrough_calls() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("safe-buffered", 7, |name, request, next| async move { + assert_eq!(name, "managed-llm"); + let targeted = next + .call(LlmContinuationInvocationV2 { + request: request.clone(), + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + let passthrough = next.call_passthrough(request).await?; + Ok(json!({ "targeted": targeted, "passthrough": passthrough })) + }) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + assert_eq!(registration.name, "safe-buffered"); + assert_eq!(registration.priority, 7); + + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed-llm", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + assert_eq!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Ok(json!({ + "targeted": { "targeted": true }, + "passthrough": { "passthrough": true } + }))) + ); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_COMPLETION_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_continuation_preserves_flattened_failure_and_rejects_invalid_outcome() { + let _guard = begin_test(); + let host = test_host_v4(); + let expected = LlmContinuationFailureV2::Http { + status: 429, + body: "bounded".into(), + headers: Default::default(), + }; + assert_eq!( + serde_json::to_value(&expected).unwrap(), + json!({ + "failure_type": "http", + "status": 429, + "body": "bounded", + "headers": {}, + }) + ); + *SAFE_V2_TARGETED_FAILURE.lock().unwrap() = Some(expected.clone()); + run_safe_v2_buffered(&host, "typed-failure", move |_, request, next| { + let expected = expected.clone(); + async move { + let error = next + .call(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .expect_err("the host returned a structured failure"); + assert_eq!(error, expected); + Ok(json!({ "observed": "typed failure" })) + } + }); + assert_eq!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Ok(json!({ "observed": "typed failure" }))) + ); + + SAFE_V2_TARGETED_INVALID_OUTCOME.store(true, Ordering::SeqCst); + run_safe_v2_buffered(&host, "invalid-outcome", |_, request, next| async move { + let error = next + .call(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .expect_err("two null callback values are not a valid outcome"); + Ok(json!({ "error": format!("{error:?}") })) + }); + let outcome = SAFE_V2_COMPLETION.lock().unwrap().take().unwrap().unwrap(); + assert!( + outcome["error"] + .as_str() + .is_some_and(|error| error.contains("invalid outcome")) + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_continuation_supports_repeated_concurrent_calls() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2( + "safe-buffered-concurrent", + 0, + |_, request, next| async move { + let calls = (0..8).map(|index| { + let next = next.clone(); + let mut request = request.clone(); + request.content["index"] = json!(index); + async move { + next.call(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}")) + } + }); + let results = futures::future::join_all(calls) + .await + .into_iter() + .collect::, _>>()?; + Ok(json!(results)) + }, + ) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + + let state = invoke_safe_v2_buffered( + &host, + ®istration, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + + let outcome = SAFE_V2_COMPLETION.lock().unwrap().take().unwrap().unwrap(); + let results = outcome + .as_array() + .expect("callback returns one result per call"); + assert_eq!(results.len(), 8); + assert!( + results + .iter() + .all(|result| result == &json!({ "targeted": true })) + ); + assert_eq!( + SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), + 1, + "all continuation clones share one retained host handle" + ); + assert_eq!(SAFE_V2_COMPLETION_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_callback_stops_when_the_caller_cancels() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("safe-buffered-cancel", 0, |_, _, _| async move { + std::future::pending::>().await + }) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + wake_safe_v2_tasks(); + drive_safe_v2_tasks(); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_COMPLETION_RELEASES.load(Ordering::SeqCst), 1); + assert!(SAFE_V2_COMPLETION.lock().unwrap().is_none()); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_cancelled_buffered_task_releases_completion_when_future_drop_panics() { + let _guard = begin_test(); + + struct PendingFutureWithPanickingDrop; + + impl Future for PendingFutureWithPanickingDrop { + type Output = std::result::Result; + + fn poll( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + std::task::Poll::Pending + } + } + + impl Drop for PendingFutureWithPanickingDrop { + fn drop(&mut self) { + panic!("safe callback future drop panic"); + } + } + + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("panic-on-cancel", 0, |_, _, _| { + PendingFutureWithPanickingDrop + }) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + let state = invoke_safe_v2_buffered( + &host, + ®istration, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + wake_safe_v2_tasks(); + drive_safe_v2_tasks(); + + assert!(SAFE_V2_TASKS.lock().unwrap().is_empty()); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_COMPLETION_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!( + LAST_ERROR.lock().unwrap().as_deref(), + Some("native API v2 host task state drop panicked") + ); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_cancelled_targeted_call_releases_next_before_the_host_callback() { + let _guard = begin_test(); + SAFE_V2_HOLD_TARGETED_CALLBACK.store(true, Ordering::SeqCst); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("safe-target-cancel", 0, |_, request, next| async move { + next.call(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}")) + }) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + while SAFE_V2_HELD_TARGETED_CALLBACK.lock().unwrap().is_none() { + std::thread::yield_now(); + } + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + wake_safe_v2_tasks(); + drive_safe_v2_tasks(); + assert_eq!( + SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), + 1, + "callback state must not retain the continuation until the host replies" + ); + let (targeted_callback, targeted_user_data) = SAFE_V2_HELD_TARGETED_CALLBACK + .lock() + .unwrap() + .take() + .unwrap(); + let error = json_host_string( + &host.v3.v1, + serde_json::to_value(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Cancelled, + message: "cancelled by host".into(), + }) + .unwrap(), + ); + unsafe { targeted_callback(targeted_user_data as *mut c_void, ptr::null(), error) }; + unsafe { + (host.v3.v1.string_free)(error); + registration.free(); + } + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_registration_pumps_provider_stream_and_releases_handles() { + let _guard = begin_test(); + SAFE_V2_PROVIDER_EVENTS.lock().unwrap().extend([ + SafeV2ProviderEvent::Chunk(json!({ "delta": "hello" })), + SafeV2ProviderEvent::Done, + ]); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2( + "safe-stream", + 3, + |_name, request, next| async move { + let mut provider = next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + let mut chunks = Vec::new(); + while let Some(item) = provider.next().await { + chunks.push(item.map_err(|error| format!("{error:?}"))?); + } + assert!( + provider.next().await.is_none(), + "completed streams stay fused" + ); + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::iter( + chunks.into_iter().map(Ok), + )))) + }, + ) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + assert_eq!(registration.name, "safe-stream"); + assert_eq!(registration.priority, 3); + + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed-llm", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + assert_eq!( + *SAFE_V2_OUTPUT.lock().unwrap(), + vec![Ok(json!({ "delta": "hello" }))] + ); + assert_eq!(SAFE_V2_OUTPUT_FINISHES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_passthrough_uses_host_owned_forwarding() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2( + "safe-passthrough", + 0, + |_name, request, _next| async move { + Ok(LlmStreamExecutionOutcomeV2::Passthrough(request)) + }, + ) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let request = test_llm_request(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "unmanaged", "request": request.clone() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + assert_eq!(*SAFE_V2_FORWARDED_REQUESTS.lock().unwrap(), vec![request]); + assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); + assert_eq!(SAFE_V2_OUTPUT_FINISHES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_provider_stream_drop_releases_unfinished_production() { + let _guard = begin_test(); + SAFE_V2_PROVIDER_EVENTS + .lock() + .unwrap() + .push_back(SafeV2ProviderEvent::Chunk(json!({ "unused": true }))); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2("safe-drop", 0, |_name, request, next| async move { + let provider = next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + drop(provider); + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::empty(), + ))) + }) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + drive_safe_v2_tasks(); + assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_open_result_releases_when_waiter_is_cancelled_before_consumption() { + let _guard = begin_test(); + SAFE_V2_HOLD_STREAM_OPEN_CALLBACK.store(true, Ordering::SeqCst); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2( + "safe-open-cancel", + 0, + |_name, request, next| async move { + let provider = next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + provider.map(|item| item.map_err(|error| format!("{error:?}"))), + ))) + }, + ) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let state = invoke_safe_v2_streaming( + &host, + ®istration, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + + let (callback, user_data) = SAFE_V2_HELD_STREAM_OPEN_CALLBACK + .lock() + .unwrap() + .take() + .expect("the stream-open continuation is pending"); + unsafe { + callback( + user_data as *mut c_void, + NonNull::::dangling().as_ptr(), + ptr::null(), + ) + }; + SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); + drive_safe_v2_tasks(); + + assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_callback_stops_when_the_caller_cancels() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2("safe-cancel", 0, |_, _, _| async move { + std::future::pending::>().await + }) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + let state = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); + SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); + wake_safe_v2_tasks(); + drive_safe_v2_tasks(); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_open_preserves_structured_failure() { + let _guard = begin_test(); + let expected = LlmContinuationFailureV2::Http { + status: 429, + body: "bounded".into(), + headers: Default::default(), + }; + *SAFE_V2_OPEN_FAILURE.lock().unwrap() = Some(expected.clone()); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2( + "typed-open-failure", + 0, + move |_name, request, next| { + let expected = expected.clone(); + async move { + let error = match next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + { + Ok(_) => panic!("stream setup should preserve the host failure"), + Err(error) => error, + }; + assert_eq!(error, expected); + Err("observed typed stream-open failure".into()) + } + }, + ) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + drive_safe_v2_tasks(); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error == "observed typed stream-open failure" + )); + assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 0); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_malformed_stream_open_releases_the_provider_stream() { + let _guard = begin_test(); + *SAFE_V2_OPEN_FAILURE.lock().unwrap() = Some(LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Internal, + message: "must not accompany a stream".into(), + }); + SAFE_V2_OPEN_RETURNS_STREAM_AND_ERROR.store(true, Ordering::SeqCst); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2( + "malformed-open", + 0, + |_, request, next| async move { + next.open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map(|provider| { + LlmStreamExecutionOutcomeV2::Stream(Box::pin( + provider.map(|item| item.map_err(|error| format!("{error:?}"))), + )) + }) + .map_err(|error| format!("{error:?}")) + }, + ) + .unwrap(); + let registration = take_safe_v2_stream_registration(); + let invocation = json_host_string( + &host.v3.v1, + json!({ "name": "managed", "request": test_llm_request() }), + ); + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invocation, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(invocation) }; + drive_safe_v2_tasks(); + + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("invalid outcome") + )); + assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { registration.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_registration_rejects_a_v1_host_without_leaking_callback_state() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + let dropped = Arc::new(AtomicUsize::new(0)); + struct DropCounter(Arc); + impl Drop for DropCounter { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + let guard = DropCounter(dropped.clone()); + let error = ctx + .register_async_llm_execution_v2("unsupported", 0, move |_, _, _| { + let _guard = &guard; + async { Ok(json!({})) } + }) + .unwrap_err(); + assert!(error.contains("ABI-v4")); + assert_eq!(dropped.load(Ordering::SeqCst), 1); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_failed_host_registration_frees_callback_state_exactly_once() { + let _guard = begin_test(); + *SAFE_V2_REGISTRATION_STATUS.lock().unwrap() = NemoRelayStatus::AlreadyExists; + let host = test_host_v4(); + + struct DropCounter(Arc); + impl Drop for DropCounter { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let buffered_drops = Arc::new(AtomicUsize::new(0)); + let buffered_guard = DropCounter(buffered_drops.clone()); + let mut ctx = test_context(&host.v3.v1); + assert!( + ctx.register_async_llm_execution_v2("duplicate", 0, move |_, _, _| { + let _guard = &buffered_guard; + async { Ok(json!({})) } + }) + .unwrap_err() + .contains("AlreadyExists") + ); + assert_eq!(buffered_drops.load(Ordering::SeqCst), 1); + + let stream_drops = Arc::new(AtomicUsize::new(0)); + let stream_guard = DropCounter(stream_drops.clone()); + let mut ctx = test_context(&host.v3.v1); + assert!( + ctx.register_async_llm_stream_execution_v2("duplicate", 0, move |_, _, _| { + let _guard = &stream_guard; + async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::empty(), + ))) + } + }) + .unwrap_err() + .contains("AlreadyExists") + ); + assert_eq!(stream_drops.load(Ordering::SeqCst), 1); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_task_spawn_failures_settle_synchronously_and_release_handles() { + let _guard = begin_test(); + *SAFE_V2_TASK_SPAWN_STATUS.lock().unwrap() = NemoRelayStatus::Internal; + let host = test_host_v4(); + + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("spawn-failure", 0, |_, _, _| async { + panic!("a failed spawn must not poll the buffered callback") + }) + .unwrap(); + let buffered = take_safe_v2_buffered_registration(); + let state = invoke_safe_v2_buffered( + &host, + &buffered, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Complete) + ); + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("task spawn failed: Internal") + )); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_COMPLETION_RELEASES.load(Ordering::SeqCst), 0); + unsafe { buffered.free() }; + + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2("stream-spawn-failure", 0, |_, _, _| async { + panic!("a failed spawn must not poll the streaming callback") + }) + .unwrap(); + let streaming = take_safe_v2_stream_registration(); + let state = invoke_safe_v2_streaming( + &host, + &streaming, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ); + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Complete) + ); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("task spawn failed: Internal") + )); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 2); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + assert!(SAFE_V2_TASKS.lock().unwrap().is_empty()); + unsafe { streaming.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_malformed_invocations_settle_and_release_callback_handles() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("malformed-buffered", 0, |_, _, _| async { + panic!("malformed invocation must not reach the callback") + }) + .unwrap(); + let buffered = take_safe_v2_buffered_registration(); + let malformed = host_string(&host.v3.v1, "not-json"); + unsafe { + (buffered.cb)( + buffered.user_data as *mut c_void, + malformed, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(malformed) }; + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("invalid native API v2 LLM invocation") + )); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { buffered.free() }; + + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2("malformed-stream", 0, |_, _, _| async { + panic!("malformed invocation must not reach the callback") + }) + .unwrap(); + let streaming = take_safe_v2_stream_registration(); + let malformed = host_string(&host.v3.v1, "not-json"); + unsafe { + (streaming.cb)( + streaming.user_data as *mut c_void, + malformed, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + unsafe { (host.v3.v1.string_free)(malformed) }; + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("invalid native API v2 stream invocation") + )); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 2); + assert_eq!(SAFE_V2_OUTPUT_RELEASES.load(Ordering::SeqCst), 1); + unsafe { streaming.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_host_task_honors_wakes_without_plugin_thread_local_state() { + let _guard = begin_test(); + let host = test_host_v4(); + run_safe_v2_buffered(&host, "self-waking", |_, _, _| async move { + let first_poll = Arc::new(AtomicBool::new(true)); + std::future::poll_fn(move |cx| { + if first_poll.swap(false, Ordering::SeqCst) { + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } else { + std::task::Poll::Ready(()) + } + }) + .await; + Ok(json!({ "woke": true })) + }); + assert_eq!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Ok(json!({ "woke": true }))) + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stale_waker_after_completion_is_a_silent_noop() { + let _guard = begin_test(); + let host = test_host_v4(); + run_safe_v2_buffered(&host, "retained-waker", |_, _, _| async { + std::future::poll_fn(|context| { + *SAFE_V2_HELD_TASK_WAKER.lock().unwrap() = Some(context.waker().clone()); + std::task::Poll::Ready(()) + }) + .await; + Ok(json!({ "done": true })) + }); + assert!(SAFE_V2_TASKS.lock().unwrap().is_empty()); + *LAST_ERROR.lock().unwrap() = None; + let waker = SAFE_V2_HELD_TASK_WAKER.lock().unwrap().take().unwrap(); + waker.wake_by_ref(); + assert!(LAST_ERROR.lock().unwrap().is_none()); + drop(waker); + assert_eq!(SAFE_V2_TASK_RETAINS.load(Ordering::SeqCst), 2); + assert_eq!(SAFE_V2_TASK_RELEASES.load(Ordering::SeqCst), 3); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_does_not_settle_a_completion_cancelled_during_callback_polling() { + let _guard = begin_test(); + let host = test_host_v4(); + run_safe_v2_buffered(&host, "cancel-before-settlement", |_, _, _| async { + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + Ok(json!({ "ignored": true })) + }); + assert!(SAFE_V2_COMPLETION.lock().unwrap().is_none()); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_trampolines_reject_invalid_callback_handles() { + let _guard = begin_test(); + let host = test_host_v4(); + + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("invalid-buffered-handles", 0, |_, _, _| async { + Ok(json!({})) + }) + .unwrap(); + let buffered = take_safe_v2_buffered_registration(); + assert_eq!( + unsafe { (buffered.cb)(ptr::null_mut(), ptr::null(), ptr::null(), ptr::null(),) }, + NemoRelayNativeAsyncCallbackState::Complete as u32 + ); + invoke_safe_v2_buffered( + &host, + &buffered, + ptr::null(), + NonNull::::dangling().as_ptr(), + ); + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("NullPointer") + )); + let releases = SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst); + invoke_safe_v2_buffered( + &host, + &buffered, + NonNull::::dangling().as_ptr(), + ptr::null(), + ); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), releases + 1); + unsafe { buffered.free() }; + + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_stream_execution_v2("invalid-stream-handles", 0, |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::empty(), + ))) + }) + .unwrap(); + let streaming = take_safe_v2_stream_registration(); + assert_eq!( + unsafe { (streaming.cb)(ptr::null_mut(), ptr::null(), ptr::null(), ptr::null(),) }, + NemoRelayNativeAsyncCallbackState::Complete as u32 + ); + invoke_safe_v2_streaming( + &host, + &streaming, + ptr::null(), + NonNull::::dangling().as_ptr(), + ); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("NullPointer") + )); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + let releases = SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst); + invoke_safe_v2_streaming( + &host, + &streaming, + NonNull::::dangling().as_ptr(), + ptr::null(), + ); + assert_eq!(SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), releases + 1); + unsafe { streaming.free() }; + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_registration_reports_allocation_failure_and_contains_drop_panics() { + let _guard = begin_test(); + let host = test_host_v4(); + + *STRING_NEW_REMAINING_SUCCESSES.lock().unwrap() = Some(0); + let mut ctx = test_context(&host.v3.v1); + assert!( + ctx.register_async_llm_execution_v2("cannot-allocate", 0, |_, _, _| async { + Ok(json!({})) + }) + .unwrap_err() + .contains("registration name") + ); + *STRING_NEW_REMAINING_SUCCESSES.lock().unwrap() = Some(0); + let mut ctx = test_context(&host.v3.v1); + assert!( + ctx.register_async_llm_stream_execution_v2("cannot-allocate", 0, |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::empty(), + ))) + }) + .unwrap_err() + .contains("registration name") + ); + *STRING_NEW_REMAINING_SUCCESSES.lock().unwrap() = None; + + struct PanicOnDrop; + impl Drop for PanicOnDrop { + fn drop(&mut self) { + panic!("safe callback state drop panic") + } + } + let panic_on_drop = PanicOnDrop; + let mut ctx = test_context(&host.v3.v1); + ctx.register_async_llm_execution_v2("drop-panic", 0, move |_, _, _| { + let _ = &panic_on_drop; + async { Ok(json!({})) } + }) + .unwrap(); + let registration = take_safe_v2_buffered_registration(); + unsafe { registration.free() }; + assert_eq!( + LAST_ERROR.lock().unwrap().as_deref(), + Some("native API v2 safe callback state drop panicked") + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_continuations_preserve_abi_and_provider_failures() { + let _guard = begin_test(); + let host = test_host_v4(); + + *SAFE_V2_TARGETED_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_buffered(&host, "target-status", |_, request, next| async move { + next.call(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}")) + }); + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("InvalidArg") + )); + *SAFE_V2_TARGETED_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + + *SAFE_V2_PASSTHROUGH_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_buffered(&host, "passthrough-status", |_, request, next| async move { + next.call_passthrough(request).await + }); + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("InvalidArg") + )); + *SAFE_V2_PASSTHROUGH_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + + *SAFE_V2_PASSTHROUGH_ERROR.lock().unwrap() = Some("provider rejected passthrough".into()); + run_safe_v2_buffered(&host, "passthrough-error", |_, request, next| async move { + next.call_passthrough(request).await + }); + assert_eq!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err("provider rejected passthrough".into())) + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_buffered_callback_settlement_is_bounded_and_panic_safe() { + let _guard = begin_test(); + let host = test_host_v4(); + + run_safe_v2_buffered(&host, "panic", |_, _, _| async move { + panic!("buffered callback panic") + }); + assert!(matches!( + SAFE_V2_COMPLETION.lock().unwrap().take(), + Some(Err(error)) if error.contains("callback panicked") + )); + + let long_error = "é".repeat(3_000); + run_safe_v2_buffered(&host, "bounded-error", move |_, _, _| { + let long_error = long_error.clone(); + async move { Err(long_error) } + }); + let error = SAFE_V2_COMPLETION + .lock() + .unwrap() + .take() + .unwrap() + .unwrap_err(); + assert_eq!(error.len(), 4 * 1024); + assert!(error.is_char_boundary(error.len())); + + *SAFE_V2_COMPLETION_RESOLVE_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_buffered(&host, "resolve-status", |_, _, _| async { + Ok(json!({ "ok": true })) + }); + assert!( + LAST_ERROR + .lock() + .unwrap() + .as_deref() + .is_some_and(|error| error.contains("completion failed")) + ); + *SAFE_V2_COMPLETION_RESOLVE_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + + *LAST_ERROR.lock().unwrap() = None; + *SAFE_V2_COMPLETION_REJECT_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_buffered(&host, "reject-status", |_, _, _| async { + Err("rejected".into()) + }); + assert!( + LAST_ERROR + .lock() + .unwrap() + .as_deref() + .is_some_and(|error| error.contains("rejection failed")) + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { + let _guard = begin_test(); + let host = test_host_v4(); + + *SAFE_V2_STREAM_OPEN_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_targeted_stream(&host, "open-status"); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("InvalidArg") + )); + *SAFE_V2_STREAM_OPEN_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + *SAFE_V2_PROVIDER_NEXT_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_targeted_stream(&host, "poll-status"); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("InvalidArg") + )); + *SAFE_V2_PROVIDER_NEXT_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + SAFE_V2_PROVIDER_EVENTS + .lock() + .unwrap() + .push_back(SafeV2ProviderEvent::Failure( + LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Transport, + message: "late provider failure".into(), + }, + )); + run_safe_v2_targeted_stream(&host, "late-failure"); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("late provider failure") + )); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + *SAFE_V2_PROVIDER_EVENT_JSON.lock().unwrap() = Some("not-json".into()); + run_safe_v2_targeted_stream(&host, "malformed-event"); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("stream chunk") + )); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + SAFE_V2_PROVIDER_INVALID_EVENT.store(true, Ordering::SeqCst); + run_safe_v2_targeted_stream(&host, "invalid-event"); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("invalid event") + )); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors() { + let _guard = begin_test(); + let host = test_host_v4(); + + SAFE_V2_OUTPUT_PUSH_STATUSES + .lock() + .unwrap() + .extend([NemoRelayStatus::WouldBlock, NemoRelayStatus::Ok]); + run_safe_v2_streaming(&host, "push-backpressure", |_, _, _| async { + Ok(safe_v2_one_chunk_stream()) + }); + assert_eq!( + *SAFE_V2_OUTPUT.lock().unwrap(), + vec![Ok(json!({ "chunk": true }))] + ); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + for (name, status, expected) in [ + ( + "push-internal", + NemoRelayStatus::Internal, + "output push failed: Internal", + ), + ( + "push-failure", + NemoRelayStatus::InvalidArg, + "output push failed", + ), + ] { + SAFE_V2_OUTPUT_PUSH_STATUSES + .lock() + .unwrap() + .push_back(status); + run_safe_v2_streaming(&host, name, |_, _, _| async { + Ok(safe_v2_one_chunk_stream()) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains(expected) + )); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + } + + *SAFE_V2_OUTPUT_FINISH_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_streaming(&host, "finish-failure", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::empty(), + ))) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("output finish failed") + )); + *SAFE_V2_OUTPUT_FINISH_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + SAFE_V2_OUTPUT_REJECT_STATUSES + .lock() + .unwrap() + .extend([NemoRelayStatus::WouldBlock, NemoRelayStatus::Ok]); + run_safe_v2_streaming(&host, "reject-backpressure", |_, _, _| async { + Err("stream rejected".into()) + }); + assert_eq!( + *SAFE_V2_OUTPUT.lock().unwrap(), + vec![Err("stream rejected".into())] + ); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + SAFE_V2_OUTPUT_REJECT_STATUSES + .lock() + .unwrap() + .push_back(NemoRelayStatus::InvalidArg); + *LAST_ERROR.lock().unwrap() = None; + run_safe_v2_streaming(&host, "reject-failure", |_, _, _| async { + Err("stream rejected".into()) + }); + assert!( + LAST_ERROR + .lock() + .unwrap() + .as_deref() + .is_some_and(|error| error.contains("output rejection failed")) + ); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_stream_cancellation_and_pass_through_failures_settle_once() { + let _guard = begin_test(); + let host = test_host_v4(); + + run_safe_v2_streaming(&host, "cancel-before-push", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( + async { + SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); + Ok(json!({ "ignored": true })) + }, + )))) + }); + assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); + SAFE_V2_OUTPUT_CANCELLED.store(false, Ordering::SeqCst); + + run_safe_v2_streaming(&host, "cancel-before-finish", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin( + stream::poll_fn(|_| { + SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); + std::task::Poll::Ready(None) + }), + ))) + }); + assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); + SAFE_V2_OUTPUT_CANCELLED.store(false, Ordering::SeqCst); + + *SAFE_V2_FORWARD_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; + run_safe_v2_streaming(&host, "forward-status", |_, request, _| async move { + Ok(LlmStreamExecutionOutcomeV2::Passthrough(request)) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("InvalidArg") + )); + *SAFE_V2_FORWARD_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + run_safe_v2_streaming(&host, "stream-panic", |_, _, _| async move { + panic!("stream callback panic") + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("callback panicked") + )); + assert_eq!(live_host_strings(), 0); +} + +#[test] +fn safe_v2_host_string_failures_do_not_leave_unsettled_handles() { + let _guard = begin_test(); + let host = test_host_v4(); + + run_safe_v2_buffered(&host, "resolve-allocation", |_, _, _| async { + *STRING_NEW_RETURNS_NULL.lock().unwrap() = true; + Ok(json!({ "cannot": "allocate" })) + }); + *STRING_NEW_RETURNS_NULL.lock().unwrap() = false; + assert!(LAST_ERROR.lock().unwrap().is_none()); + + *LAST_ERROR.lock().unwrap() = None; + run_safe_v2_streaming(&host, "reject-allocation", |_, _, _| async { + *STRING_NEW_RETURNS_NULL.lock().unwrap() = true; + Err("cannot allocate rejection".into()) + }); + *STRING_NEW_RETURNS_NULL.lock().unwrap() = false; + assert!(LAST_ERROR.lock().unwrap().is_none()); + assert_eq!(live_host_strings(), 0); +} diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 772a0b827..b380e14d2 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -13,7 +13,7 @@ two execution lanes: | Lane | Use when | Stable boundary | | --- | --- | --- | -| `rust_dynamic` | Behavior must run in the Relay process. | Native ABI v2 | +| `rust_dynamic` | Behavior must run in the Relay process. | Versioned native plugin C ABI | | `worker` | Behavior should run in a separate local process. | `grpc-v1` | The manifest describes compatibility, capabilities, artifact integrity, and the @@ -76,7 +76,7 @@ The following requirements vary by execution lane: | Manifest area | Native dynamic plugin | Worker plugin | | --- | --- | --- | | `plugin.kind` | `rust_dynamic` | `worker` | -| `compat` | `native_api = "1"` | `worker_protocol = "grpc-v1"` | +| `compat` | `native_api = "1"` for the established surface. For targeted LLM continuations, use `native_api = "2"` with `relay = ">=0.8,<1.0"`. | `worker_protocol = "grpc-v1"` | | `capabilities.items` | Includes `plugin_native` | Includes `plugin_worker` | | `load` | `library` and `symbol` | `runtime` and `entrypoint` | | `source.manifest_root` | Optional | Required for `runtime = "python"`; `nemo-relay plugins add` uses it to create and retain the managed worker environment. | diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 82167c28e..c62476f65 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -1,6 +1,6 @@ --- title: "Native Dynamic Plugins (Rust)" -description: "Build in-process Rust shared-library plugins against the NeMo Relay Native ABI v3." +description: "Build in-process Rust shared-library plugins against NeMo Relay native API v1 or v2." position: 10 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. @@ -13,7 +13,7 @@ contract: a stable kind, JSON component configuration, validation diagnostics, and registration through a component-scoped context. -Native plugins are not sandboxed. They run in the gateway process, must not +Native plugins are not sandboxed. They run in the Relay host process, must not unwind across ABI callbacks, and remain loaded until Relay removes their registered callbacks. @@ -116,11 +116,161 @@ path, then replace `` with that library's SHA-256 digest. Use Native Plugin](/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example) for a complete example with validation, middleware, scopes, and configuration schema support. -## Native ABI v3 +## Select a Native API + +`compat.native_api` is the operator-facing native plugin C API version. It is +separate from the internal host-table `abi_version` field: + +| Manifest API | Rust export macro | Host table | Use case | +|---|---|---|---| +| `"1"` | `nemo_relay_plugin!` | `NemoRelayNativeHostApiV3` | Existing subscribers, guardrails, scopes, and generic middleware | +| `"2"` | `nemo_relay_plugin_v2!` | `NemoRelayNativeHostApiV4` | Host-dispatched LLM calls with typed HTTP targets and provider streams | + +Relay preserves native API v1 symbols, layouts, numeric values, and loader +behavior. Existing binaries remain compatible. Rust plugins that exhaustively +match `NemoRelayStatus` must handle the new `WouldBlock` variant when they +recompile against the 0.8 SDK. A v2-only plugin is rejected clearly by a v1 +host and is not retried against an older table. Validation, inspection, and +doctor output report the selected manifest API. + +Native API v2 is intended for in-process orchestrators that decide which LLM +call to make while Relay still owns provider transport. The plugin supplies a +replacement `LlmRequest` and a target containing the absolute HTTP(S) URL and +explicit outbound headers. Relay sends targeted LLM requests with HTTP `POST`. +Target headers may contain provider credentials. Relay validates and +transports them but never places their values in diagnostics or observability. +Protocol selection and request translation remain plugin concerns; Relay sends +the supplied JSON to the selected target. + +The typed target is invocation-scoped continuation context. It is not encoded +into `LlmRequest.headers`, and target credentials are not visible to downstream +request middleware. Relay returns either: + +- buffered response JSON, bounded to 16 MiB; +- an HTTP failure with status, a body bounded to 16 KiB, and a conservative + safe-header allowlist; +- a non-HTTP failure classified as transport, timeout, cancelled, invalid + request, guardrail, or internal; or +- a host-owned provider stream pulled one JSON event at a time, including typed + setup and late failures. + +Targeted streaming endpoints must return SSE with a JSON value in each `data` +frame. Relay removes the SSE framing and exposes those JSON events to the +plugin; raw byte streams and non-SSE streaming protocols are outside this +contract. + +Relay core owns the terminal HTTP transport for these continuations. The same +targeted plugin therefore works in the CLI gateway and in an SDK-embedded Relay +host that calls `llm_call_execute` or `llm_stream_call_execute` directly. +Remaining LLM execution intercepts still run before core dispatches the target; +the host's original provider callback is used only when no target is bound. + +Relay reports neutral HTTP status and non-HTTP failure data. Each plugin owns +its retry, reselection, and fallback policy; the SDK does not classify failures +as retryable. Relay does not inspect provider bodies to infer context-window or +model-availability errors. + +Targeted dispatch does not follow redirects. Relay rejects embedded URL +credentials, hop-by-hop headers, host-owned framing headers, and +`x-nemo-relay-internal-*` headers. Provider streams are pull-based and permit +one pending poll; plugin output and direct pass-through use bounded queues. +Dropping or cancelling a stream stops provider production, and the +library stays loaded until all callbacks and streams release their handles. + +A clean activation unloads its native library normally. If activation teardown +finds an opaque callback, task, continuation, or stream handle that still owns +plugin code, Relay releases the descriptor and handle state but conservatively +keeps that library mapping loaded for the rest of the process. This prevents a +final release made from plugin code from unmapping its own caller before the FFI +operation returns. + +Safe v2 callbacks register through the generic V3 asynchronous-middleware APIs +and return `Pending`. Buffered callbacks use one completion; streaming +callbacks produce incrementally through an output stream. Relay polls their +Rust futures cooperatively on its existing Tokio runtime, restoring the +captured continuation and scope context for every poll. A pending callback does +not occupy a blocking worker or create an OS thread. Separate callback +invocations can run concurrently and have no stable OS-thread affinity. + +Rust authors use the safe SDK facade. Buffered callbacks receive +`LlmContinuationV2`; streaming callbacks receive `LlmStreamContinuationV2` +and return either a boxed Rust stream or an explicit pass-through request. +For this native API v2 example, use `nemo-relay-plugin = "0.8.0"` and add +`futures = "0.3"` to the plugin's `Cargo.toml` dependencies. + +```rust +use futures::StreamExt; +use nemo_relay_plugin::{ + LlmContinuationInvocationV2, LlmStreamExecutionOutcomeV2, +}; + +let buffered_target = target.clone(); +context.register_async_llm_execution_v2("route", 0, move |_name, request, next| { + let target = buffered_target.clone(); + async move { + next.call(LlmContinuationInvocationV2 { request, target }) + .await + .map_err(|failure| format!("{failure:?}")) + } +})?; + +let stream_target = target.clone(); +context.register_async_llm_stream_execution_v2( + "route-stream", + 0, + move |_name, request, next| { + let target = stream_target.clone(); + async move { + let provider = next + .open_stream(LlmContinuationInvocationV2 { request, target }) + .await + .map_err(|failure| format!("{failure:?}"))?; + let stream = provider.map(|item| { + item.map_err(|failure| format!("{failure:?}")) + }); + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream))) + } + }, +)?; +``` + +Use `LlmContinuationV2::call_passthrough` to invoke the ordinary untargeted +downstream continuation for a buffered call. For streaming, return +`LlmStreamExecutionOutcomeV2::Passthrough(request)`. The call remains inside +its managed Relay LLM lifecycle, but Relay connects the downstream stream +directly to the caller through its bounded queue, so pass-through events do not +cross the plugin boundary. + +Continuation operations follow the outer callback mode. A streaming callback +can open targeted streams or pass through to the downstream stream, but it +cannot invoke a buffered continuation. A streaming policy that needs a judge +or side call must consume that call as a stream and aggregate it itself. + +The SDK returns `Pending` to the host, which drives safe callback futures and +returned streams as cooperative tasks. `LlmProviderStreamV2` implements +`Stream`, enforces one pending pull, and cancels unfinished provider production +on drop. If the bounded plugin-output queue fills, Relay wakes the task when +the consumer makes room instead of tying up a worker while waiting. + +Export a v2-only plugin with: + +```rust +nemo_relay_plugin::nemo_relay_plugin_v2!( + nemo_relay_register_plugin, + || NativePolicy +); +``` + +Set `compat.native_api = "2"` in its manifest. Use +the safe registration methods above for normal Rust plugins. + +## Native C ABI Table Versioning The entry symbol receives a `*const NemoRelayNativeHostApiV1` pointer. It -points at the v1 prefix of a v3 `NemoRelayNativeHostApiV3` table; check -`abi_version` and `struct_size` before casting. The plugin returns a +points at the frozen v1 prefix of the negotiated host table; check +`abi_version` and `struct_size` before casting. A manifest native API v1 plugin +receives `NemoRelayNativeHostApiV3`, while native API v2 receives the V4 +extension. The plugin returns a `NemoRelayNativePluginV1` descriptor: ```rust @@ -130,7 +280,18 @@ extern "C" fn nemo_relay_register_plugin( ) -> NemoRelayStatus ``` -The v3 host table retains the frozen legacy prefix and appends a +`PluginContext::host_api_v4` and the generic V3 raw registration methods are +the advanced escape hatch for raw ABI consumers and non-Rust bindings. The +safe Rust facade uses V3's completion-based buffered registration and +incremental async-stream registration. V4 contributes the targeted LLM +continuation operations and general opaque host-task hooks needed to poll +Rust-side futures cooperatively; it does not add an LLM-specific execution +lane. Raw callers must own host strings, completion settlement, task and stream +backpressure, cancellation, panic fencing, and every release operation. Rust +futures, streams, trait objects, and allocator-owned strings never cross the +shared-library boundary. + +The V3 host table used by manifest native API v1 retains the frozen legacy prefix and appends a completion-based asynchronous middleware extension. An entry that rejects the v3 table with `InvalidArg` is retried with the legacy table. Rust plugins using the typed `NativePlugin` APIs continue to work unchanged. Raw ABI plugins can use diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index ed76d4bb1..ab170287c 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -335,21 +335,25 @@ wrong-direction capability IDs. ### Rebuild Native Plugins and Raw FFI Consumers -NeMo Relay 0.7 uses native ABI v3. Recompile native plugins against the 0.7 +NeMo Relay 0.7 uses the V3 native C host table for manifest native API v1. +Recompile native plugins against the 0.7 `nemo-relay-plugin` crate and rebuild raw FFI consumers against the generated 0.7 header. The v3 table preserves the v2 prefix, and Relay retries a legacy v2 table when loading a plugin that rejects v3. That fallback supports loading, not compatibility with changed middleware, LLM sanitizer, ABI, or schema contracts. -Rebuild native plugins that use the raw plugin ABI callbacks: native ABI v3 +Rebuild native plugins that use the raw plugin ABI callbacks: the V3 host table adds completion-based async middleware registration, async execution continuations, and explicit cancellation/late-settlement behavior. This is separate from the synchronous `nemo-relay-ffi` middleware registration API. -The plugin manifest value remains `compat.native_api = "1"`. This manifest -contract version is separate from the host ABI version; do not change it to -`"2"`. +Existing plugins keep `compat.native_api = "1"`. This manifest contract +version is separate from the host-table revision. Select `"2"` only when a +plugin adopts the targeted LLM continuation surface; rebuilding an existing +native API v1 plugin does not require that migration. Rust plugins with an +exhaustive `NemoRelayStatus` match must add an arm for `WouldBlock` when +recompiling against the 0.8 SDK; existing plugin binaries remain compatible. Request and response callbacks now receive distinct context structures. Each structure contains structured codec identity and a borrowed directional codec @@ -362,7 +366,7 @@ after the callback returns. Release host-owned output strings with the standard host string release operation. For the complete ABI contract, refer to -[Native ABI v3](/build-plugins/dynamic-plugins/native-dynamic/about#native-abi-v3). +[Native C ABI Table Versioning](/build-plugins/dynamic-plugins/native-dynamic/about#native-c-abi-table-versioning). ### Update PII Redaction Configuration