From 8d58cce9d6e1bb571f6ef95401e150c74ca66fb1 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 30 Jul 2026 14:46:31 -0600 Subject: [PATCH 01/32] feat(plugin): add native API v2 LLM dispatch Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 649 +++++++++++++++++- .../tests/fixtures/native_plugin/src/lib.rs | 29 + .../tests/integration/native_plugin_tests.rs | 57 +- crates/core/tests/unit/native_plugin_tests.rs | 309 ++++++++- crates/plugin/src/lib.rs | 401 ++++++++++- crates/plugin/tests/typed_callbacks.rs | 52 +- 6 files changed, 1440 insertions(+), 57 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 3b0124b5b..750535fea 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -32,7 +32,8 @@ use crate::api::runtime::{ }; 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 +42,7 @@ 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}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, deregister_plugin_registration_checked, register_plugin_tracked, @@ -49,14 +50,17 @@ use crate::plugin::{ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ - NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + LlmCallErrorV2, LlmCallOutcomeV2, LlmDispatchRequestV2, LlmStreamEventV2, + LlmUpstreamFailureClassV2, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, + NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, + NemoRelayNativeAsyncLlmStreamCbV2, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, + NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, @@ -327,11 +331,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 +385,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); @@ -797,16 +811,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 +886,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 +910,21 @@ 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_TYPED_LLM_DISPATCH; + v3.v1.struct_size = std::mem::size_of::(); + NemoRelayNativeHostApiV4 { + v3, + async_llm_next_invoke_result_v2: native_async_llm_next_invoke_result_v2, + async_llm_next_invoke_stream_v2: native_async_llm_next_invoke_stream_v2, + plugin_context_register_async_llm_execution_v2: + native_plugin_context_register_async_llm_execution_v2, + plugin_context_register_async_llm_stream_execution_v2: + native_plugin_context_register_async_llm_stream_execution_v2, + } +} + fn read_native_string(value: *const NemoRelayNativeString) -> crate::plugin::Result { if value.is_null() { return Ok(String::new()); @@ -1392,6 +1426,7 @@ fn make_user_data( } const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY: usize = 64; +const NATIVE_API_V2_STREAM_CHANNEL_CAPACITY: usize = 32; struct NativeAsyncCompletion { sender: Mutex>>>, @@ -1493,6 +1528,56 @@ struct NativeAsyncStreamCallbackGuard { active: bool, } +struct NativeAsyncLlmStreamCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamCbV2, + user_data: usize, + stream: Arc, + _library_guard: Option>, + active: bool, +} + +impl NativeAsyncLlmStreamCallbackGuardV2 { + fn emit(&self, event: &LlmStreamEventV2) -> bool { + let Some(event) = native_string_from_json( + &serde_json::to_value(event) + .expect("native API v2 stream events contain serializable Relay DTOs"), + ) else { + return false; + }; + let keep_going = unsafe { (self.cb)(self.user_data as *mut c_void, event) }; + unsafe { native_string_free(event) }; + keep_going + } + + fn finish(&mut self, event: &LlmStreamEventV2) { + if self.active { + let _ = self.emit(event); + self.active = false; + } + } +} + +impl Drop for NativeAsyncLlmStreamCallbackGuardV2 { + fn drop(&mut self) { + if !self.active { + return; + } + let message = if self.stream.cancelled.load(Ordering::Acquire) { + "typed native LLM stream continuation was cancelled" + } else if self.stream.settled.load(Ordering::Acquire) { + "typed native LLM stream output settled" + } else { + "typed native LLM stream continuation ended without a terminal event" + }; + let event = LlmStreamEventV2::Failure { + error: LlmCallErrorV2::Cancelled { + message: message.into(), + }, + }; + let _ = self.emit(&event); + } +} + impl NativeAsyncStreamCallbackGuard { fn finish(&mut self) { self.active = false; @@ -1591,6 +1676,25 @@ async fn invoke_native_async_callback( user_data: Arc, invocation: Json, next: Option, +) -> FlowResult { + invoke_native_async_callback_with_lane(cb, user_data, invocation, next, false).await +} + +async fn invoke_native_async_callback_blocking( + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: Arc, + invocation: Json, + next: Option, +) -> FlowResult { + invoke_native_async_callback_with_lane(cb, user_data, invocation, next, true).await +} + +async fn invoke_native_async_callback_with_lane( + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: Arc, + invocation: Json, + next: Option, + blocking: bool, ) -> FlowResult { let runtime = if next.is_some() { Some(tokio::runtime::Handle::try_current().map_err(|error| { @@ -1629,16 +1733,32 @@ async fn invoke_native_async_callback( (None, None) => None, _ => unreachable!("runtime is present exactly for native async intercepts"), }; - 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, - ) - })) { + let callback_user_data = user_data.ptr as usize; + let callback_scope_stack = current_scope_stack(); + let invoke = move || { + with_scope_stack(callback_scope_stack, || { + catch_unwind(AssertUnwindSafe(|| 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, + ) + })) + }) + }; + let callback_result = if blocking { + tokio::task::spawn_blocking(invoke).await.map_err(|error| { + FlowError::Internal(format!( + "native API v2 blocking callback task failed: {error}" + )) + })? + } else { + invoke() + }; + let state = match callback_result { Ok(state) => state, Err(_) => { unsafe { @@ -2151,6 +2271,140 @@ unsafe extern "C" fn native_async_next_invoke_result( NemoRelayStatus::Ok } +const NATIVE_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; +const NATIVE_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; +const NATIVE_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; + +fn prepare_typed_llm_dispatch( + mut dispatch: LlmDispatchRequestV2, +) -> std::result::Result { + let url = match reqwest::Url::parse(&dispatch.target.url) { + Ok(url) if matches!(url.scheme(), "http" | "https") && url.has_host() => url, + _ => { + set_native_last_error("typed LLM dispatch target must be an absolute HTTP(S) URL"); + return Err(NemoRelayStatus::InvalidArg); + } + }; + dispatch.request.headers.insert( + NATIVE_DISPATCH_URL_HEADER.into(), + Json::String(url.to_string()), + ); + dispatch.request.headers.insert( + NATIVE_DISPATCH_ROUTE_HEADER.into(), + Json::String(dispatch.target.route.as_str().into()), + ); + dispatch.request.headers.insert( + NATIVE_RETRY_AWARE_HEADER.into(), + Json::String("true".into()), + ); + Ok(dispatch.request) +} + +fn typed_llm_failure(error: FlowError) -> LlmCallErrorV2 { + match error { + FlowError::Upstream(failure) => { + let class = match failure.class { + UpstreamFailureClass::Connection => LlmUpstreamFailureClassV2::Connection, + UpstreamFailureClass::Timeout => LlmUpstreamFailureClassV2::Timeout, + UpstreamFailureClass::RetryableStatus => LlmUpstreamFailureClassV2::RetryableStatus, + UpstreamFailureClass::ContextWindow => LlmUpstreamFailureClassV2::ContextWindow, + UpstreamFailureClass::ModelUnavailable => { + LlmUpstreamFailureClassV2::ModelUnavailable + } + UpstreamFailureClass::Authentication => LlmUpstreamFailureClassV2::Authentication, + UpstreamFailureClass::InvalidRequest => LlmUpstreamFailureClassV2::InvalidRequest, + UpstreamFailureClass::Other => LlmUpstreamFailureClassV2::Other, + }; + LlmCallErrorV2::Upstream { + class, + retryable: failure.is_retryable(), + status: failure.status, + body: failure.body, + headers: failure.headers, + } + } + FlowError::GuardrailRejected(message) => LlmCallErrorV2::GuardrailRejected { message }, + FlowError::InvalidArgument(message) => LlmCallErrorV2::InvalidRequest { message }, + other => LlmCallErrorV2::Internal { + message: other.to_string(), + }, + } +} + +unsafe fn invoke_typed_llm_result_callback( + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: *mut c_void, + outcome: &LlmCallOutcomeV2, +) { + if let Some(outcome) = native_string_from_json( + &serde_json::to_value(outcome) + .expect("native API v2 LLM outcomes contain only serializable Relay DTOs"), + ) { + unsafe { + cb(user_data, outcome); + native_string_free(outcome); + } + } +} + +/// Invokes a unary LLM continuation through native API v2. +unsafe extern "C" fn native_async_llm_next_invoke_result_v2( + next: *const NemoRelayNativeAsyncNext, + dispatch_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("typed unary LLM dispatch requires an LLM execution continuation"); + return NemoRelayStatus::InvalidArg; + }; + let dispatch = match parse_json_arg(dispatch_json, "typed LLM dispatch").and_then(|value| { + serde_json::from_value(value).map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidJson + }) + }) { + Ok(dispatch) => dispatch, + Err(status) => return status, + }; + let request = match prepare_typed_llm_dispatch(dispatch) { + 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 next_fn = next_fn.clone(); + 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(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()) + ))) + }); + let outcome = match result { + Ok(response) => LlmCallOutcomeV2::Success { response }, + Err(error) => LlmCallOutcomeV2::Failure { + error: typed_llm_failure(error), + }, + }; + unsafe { + invoke_typed_llm_result_callback(cb, user_data as *mut c_void, &outcome); + } + drop(library_guard); + }); + NemoRelayStatus::Ok +} + unsafe extern "C" fn native_async_next_invoke_stream( next: *const NemoRelayNativeAsyncNext, invocation_json: *const NemoRelayNativeString, @@ -2307,6 +2561,135 @@ async fn forward_native_async_next_stream_with( ); } callback_guard.finish(); +/// Invokes a streaming LLM continuation through native API v2. +unsafe extern "C" fn native_async_llm_next_invoke_stream_v2( + next: *const NemoRelayNativeAsyncNext, + dispatch_json: *const NemoRelayNativeString, + output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamCbV2, + 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( + "typed streaming LLM dispatch requires an LLM stream execution continuation", + ); + return NemoRelayStatus::InvalidArg; + }; + let dispatch = + match parse_json_arg(dispatch_json, "typed streaming LLM dispatch").and_then(|value| { + serde_json::from_value(value).map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidJson + }) + }) { + Ok(dispatch) => dispatch, + Err(status) => return status, + }; + let request = match prepare_typed_llm_dispatch(dispatch) { + 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 = NativeAsyncLlmStreamCallbackGuardV2 { + 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(async move { + let mut callback_guard = callback_guard; + let result = AssertUnwindSafe(async { + match next_fn(request).await { + Ok(mut stream) => { + while let Some(item) = stream.next().await { + let event = match item { + Ok(chunk) => LlmStreamEventV2::Chunk { chunk }, + Err(error) => { + callback_guard.finish(&LlmStreamEventV2::Failure { + error: typed_llm_failure(error), + }); + return; + } + }; + if !callback_guard.emit(&event) { + callback_guard.active = false; + return; + } + } + callback_guard.finish(&LlmStreamEventV2::Done); + } + Err(error) => { + callback_guard.finish(&LlmStreamEventV2::Failure { + error: typed_llm_failure(error), + }); + } + } + }) + .catch_unwind() + .await; + if let Err(payload) = result { + callback_guard.finish(&LlmStreamEventV2::Failure { + error: LlmCallErrorV2::Internal { + message: format!( + "typed native LLM stream continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ), + }, + }); + } + }) + .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); + drop(settlement); + let _ = start_tx.send(()); + NemoRelayStatus::Ok } fn wrap_native_async_tool_json( @@ -2563,6 +2946,28 @@ fn wrap_native_async_llm_execution( }) } +fn wrap_native_async_llm_execution_v2( + instance: Arc, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> LlmExecutionFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |name, request, next| { + let user_data = user_data.clone(); + let name = name.to_owned(); + Box::pin(async move { + invoke_native_async_callback_blocking( + cb, + user_data, + serde_json::json!({"name": name, "request": request}), + Some(NativeAsyncNextInner::Llm(next)), + ) + .await + }) + }) +} + fn wrap_native_incremental_llm_stream_execution( instance: Arc, cb: NemoRelayNativeAsyncStreamMiddlewareCb, @@ -2651,6 +3056,126 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data( }) } +fn wrap_native_incremental_llm_stream_execution_v2( + instance: Arc, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> LlmStreamExecutionFn { + let user_data = make_user_data(instance, user_data, free_fn); + 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_API_V2_STREAM_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(()), + #[cfg(test)] + before_settlement_lock: None, + _callback_user_data: Some(user_data.clone()), + }); + let output = NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }; + let invocation = + native_string_from_json(&serde_json::json!({"name": name, "request": request})) + .ok_or_else(|| { + FlowError::Internal( + "failed to allocate native API v2 stream invocation".into(), + ) + })? as usize; + let runtime = tokio::runtime::Handle::try_current().map_err(|error| { + FlowError::Internal(format!( + "native API v2 stream intercept requires a Tokio runtime: {error}" + )) + })?; + let next_ref = Arc::into_raw(Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(next), + runtime, + Some(user_data.clone()), + ))) as usize; + let stream_ref = Arc::into_raw(stream.clone()) as usize; + let callback_user_data = user_data.ptr as usize; + let callback_scope_stack = current_scope_stack(); + let blocking_task = tokio::task::spawn_blocking(move || { + with_scope_stack(callback_scope_stack, || { + catch_unwind(AssertUnwindSafe(|| unsafe { + cb( + callback_user_data as *mut c_void, + invocation as *const NemoRelayNativeString, + next_ref as *const NemoRelayNativeAsyncNext, + stream_ref as *const NemoRelayNativeAsyncStream, + ) + })) + }) + }); + let stream_for_monitor = Arc::clone(&stream); + tokio::spawn(async move { + let state = blocking_task + .await + .ok() + .and_then(std::result::Result::ok) + .and_then(|state| NemoRelayNativeAsyncCallbackState::try_from(state).ok()); + unsafe { native_string_free(invocation as *mut NemoRelayNativeString) }; + let error = match state { + Some(NemoRelayNativeAsyncCallbackState::Pending) => None, + Some(NemoRelayNativeAsyncCallbackState::Complete) + if stream_for_monitor.settled.load(Ordering::Acquire) + || stream_for_monitor.cancelled.load(Ordering::Acquire) => + { + None + } + Some(NemoRelayNativeAsyncCallbackState::Complete) => Some( + "native API v2 stream callback returned Complete without finishing" + .to_string(), + ), + None => Some( + "native API v2 stream callback panicked or returned an invalid state" + .to_string(), + ), + }; + if let Some(error) = error { + settle_native_api_v2_stream_error(stream_for_monitor, error).await; + } + }); + Ok(LlmJsonStream::new(output)) + }) + }) +} + +async fn settle_native_api_v2_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(); + } +} + unsafe extern "C" fn native_plugin_context_register_async_stream_middleware( ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -2682,6 +3207,68 @@ unsafe extern "C" fn native_plugin_context_register_async_stream_middleware( } } +unsafe extern "C" fn native_plugin_context_register_async_llm_stream_execution_v2( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + clear_native_last_error(); + let user_data_guard = NativeCallbackUserDataGuard::new(user_data, free_fn); + let host_ctx = match host_ctx_mut(ctx) { + Ok(ctx) => ctx, + Err(status) => return status, + }; + let instance = host_ctx.instance.clone(); + let name = match read_name(name) { + Ok(name) => name, + Err(status) => return status, + }; + let (user_data, free_fn) = user_data_guard.transfer(); + let context = unsafe { &mut *host_ctx.ctx }; + match context.register_llm_stream_execution_intercept( + &name, + priority, + wrap_native_incremental_llm_stream_execution_v2(instance, cb, user_data, free_fn), + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(error), + } +} + +unsafe extern "C" fn native_plugin_context_register_async_llm_execution_v2( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + clear_native_last_error(); + let user_data_guard = NativeCallbackUserDataGuard::new(user_data, free_fn); + let host_ctx = match host_ctx_mut(ctx) { + Ok(ctx) => ctx, + Err(status) => return status, + }; + let instance = host_ctx.instance.clone(); + let name = match read_name(name) { + Ok(name) => name, + Err(status) => return status, + }; + let (user_data, free_fn) = user_data_guard.transfer(); + let context = unsafe { &mut *host_ctx.ctx }; + match context.register_llm_execution_intercept( + &name, + priority, + wrap_native_async_llm_execution_v2(instance, cb, user_data, free_fn), + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(error), + } +} + unsafe extern "C" fn native_plugin_context_register_async_middleware( ctx: *mut NemoRelayNativePluginContext, kind: u32, diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index a8c47a1bc..f8c077123 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -297,6 +297,35 @@ 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 +); + +#[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 != 3 + || 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..7b5c491ab 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -1198,7 +1198,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 +1246,61 @@ 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_manifest_writer_escapes_toml_strings() { let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index f0978fc6d..56c24bc95 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; @@ -111,6 +111,54 @@ 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, + outcome_json: *const NemoRelayNativeString, +) { + let sender = + unsafe { Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) }; + let outcome = parse_json_arg(outcome_json, "typed LLM result") + .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) + .expect("host emitted a valid typed LLM outcome"); + let _ = sender.send(outcome); +} + +unsafe extern "C" fn reject_unexpected_typed_llm_result( + _user_data: *mut c_void, + _outcome_json: *const NemoRelayNativeString, +) { + panic!("invalid dispatch unexpectedly invoked its callback"); +} + +#[derive(Default)] +struct TypedLlmStreamCallbackState { + events: Mutex>, + terminal: tokio::sync::Notify, +} + +unsafe extern "C" fn record_typed_llm_stream_result( + user_data: *mut c_void, + event_json: *const NemoRelayNativeString, +) -> bool { + let state = unsafe { &*(user_data as *const TypedLlmStreamCallbackState) }; + let event: LlmStreamEventV2 = parse_json_arg(event_json, "typed LLM stream result") + .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) + .expect("host emitted a valid typed LLM stream event"); + let terminal = matches!( + event, + LlmStreamEventV2::Done | LlmStreamEventV2::Failure { .. } + ); + state + .events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(event); + if terminal { + state.terminal.notify_one(); + } + true +} + #[derive(Default)] struct NativeStreamCallbackState { error: Mutex>, @@ -1207,10 +1255,13 @@ fn assert_native_json_output_and_host_api() { ); 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_TYPED_LLM_DISPATCH + ); assert_eq!( host_api.struct_size, - std::mem::size_of::() + std::mem::size_of::() ); } @@ -1467,6 +1518,147 @@ fn native_async_next_result_supports_repeated_concurrent_calls() { } } +#[test] +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 next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|request| { + Box::pin(async move { + assert_eq!( + request + .headers + .get(NATIVE_DISPATCH_URL_HEADER) + .and_then(Json::as_str), + Some("https://provider.example/v1/chat/completions") + ); + assert_eq!( + request + .headers + .get(NATIVE_DISPATCH_ROUTE_HEADER) + .and_then(Json::as_str), + Some("openai_chat") + ); + assert_eq!( + request + .headers + .get(NATIVE_RETRY_AWARE_HEADER) + .and_then(Json::as_str), + Some("true") + ); + 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(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model"}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + url: "https://provider.example/v1/chat/completions".into(), + route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + }, + }) + .unwrap(), + ) + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + let outcome = runtime.block_on(receiver).unwrap(); + assert_eq!( + outcome, + LlmCallOutcomeV2::Failure { + error: LlmCallErrorV2::Upstream { + class: LlmUpstreamFailureClassV2::RetryableStatus, + retryable: true, + status: Some(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_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(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + url: "/v1/chat/completions".into(), + route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + }, + }) + .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_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { let runtime = tokio::runtime::Builder::new_current_thread() @@ -1528,6 +1720,117 @@ fn native_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { } } +#[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 next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|request| { + assert_eq!( + request + .headers + .get(NATIVE_DISPATCH_ROUTE_HEADER) + .and_then(Json::as_str), + Some("anthropic_messages") + ); + Box::pin(async move { + Ok(LlmJsonStream::new(tokio_stream::iter(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, + })), + ]))) + }) + })), + runtime.handle().clone(), + 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_ref = + Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; + let callback_state = Arc::new(TypedLlmStreamCallbackState::default()); + let dispatch = native_string_from_json( + &serde_json::to_value(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model", "stream": true}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + url: "https://provider.example/v1/messages".into(), + route: nemo_relay_plugin::LlmDispatchRouteV2::AnthropicMessages, + }, + }) + .unwrap(), + ) + .unwrap(); + + assert_eq!( + unsafe { + native_async_llm_next_invoke_stream_v2( + next_ref, + dispatch, + output_stream_ref, + record_typed_llm_stream_result, + Arc::as_ptr(&callback_state).cast_mut().cast(), + ) + }, + NemoRelayStatus::Ok + ); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), callback_state.terminal.notified()) + .await + .expect("typed LLM stream should emit a terminal event"); + }); + assert_eq!( + *callback_state + .events + .lock() + .unwrap_or_else(|error| error.into_inner()), + vec![ + LlmStreamEventV2::Chunk { + chunk: json!({ + "type": "content_block_delta", + "delta": {"text": "hi"}, + }), + }, + LlmStreamEventV2::Failure { + error: LlmCallErrorV2::Upstream { + class: LlmUpstreamFailureClassV2::ModelUnavailable, + retryable: true, + status: Some(503), + body: "unavailable".into(), + headers: BTreeMap::new(), + }, + }, + ] + ); + + drop(NativeAsyncStreamReceiver { + 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 native_continuation_context_observation( expected_stack: &ScopeStackHandle, expected_event_uuid: uuid::Uuid, diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index b6bf3da7d..0fe9102d3 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -9,6 +9,7 @@ //! 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::marker::{PhantomData, PhantomPinned}; use std::panic::{AssertUnwindSafe, catch_unwind}; @@ -35,13 +36,15 @@ 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. +/// 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 typed LLM target dispatch and outcomes. +pub const NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH: u32 = 4; /// Legacy native plugin ABI accepted by Relay hosts for compatibility. pub const NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY: u32 = 2; @@ -897,6 +900,170 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( error: *const NemoRelayNativeString, ); +/// Provider protocol selected for one typed LLM dispatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LlmDispatchRouteV2 { + /// OpenAI Chat Completions. + OpenaiChat, + /// OpenAI Responses. + OpenaiResponses, + /// Anthropic Messages. + AnthropicMessages, +} + +impl LlmDispatchRouteV2 { + /// Returns the stable Relay gateway route identifier. + pub const fn as_str(self) -> &'static str { + match self { + Self::OpenaiChat => "openai_chat", + Self::OpenaiResponses => "openai_responses", + Self::AnthropicMessages => "anthropic_messages", + } + } +} + +/// Explicit provider target supplied to Relay through native API v2. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +pub struct LlmDispatchTargetV2 { + /// Absolute HTTP(S) provider URL including the selected endpoint. + pub url: String, + /// Provider protocol used by the selected endpoint. + pub route: LlmDispatchRouteV2, +} + +/// Typed LLM continuation invocation supplied through native API v2. +#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] +pub struct LlmDispatchRequestV2 { + /// Replacement request passed to the Relay execution continuation. + pub request: LlmRequest, + /// Explicit provider target selected by the plugin. + pub target: LlmDispatchTargetV2, +} + +/// Stable provider-failure classification exposed to native plugins. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LlmUpstreamFailureClassV2 { + /// Provider connection could not be established or was interrupted. + Connection, + /// Provider request timed out. + Timeout, + /// Retryable HTTP status without a more specific provider classification. + RetryableStatus, + /// Provider rejected the request because its context window was exceeded. + ContextWindow, + /// Requested provider model is temporarily unavailable. + ModelUnavailable, + /// Provider authentication or authorization failed. + Authentication, + /// Provider rejected an invalid request. + InvalidRequest, + /// Other non-retryable provider failure. + Other, +} + +/// Structured LLM continuation failure exposed through native API v2. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum LlmCallErrorV2 { + /// Relay received a provider or provider-transport failure. + Upstream { + /// Stable failure classification. + class: LlmUpstreamFailureClassV2, + /// Whether Relay considers the failure safe to retry. + retryable: bool, + /// Provider HTTP status when a response was received. + status: Option, + /// Bounded provider response body or transport error. + body: String, + /// Safe response headers with credential-bearing fields removed. + headers: BTreeMap, + }, + /// A guardrail rejected the provider call. + GuardrailRejected { + /// Human-readable rejection reason. + message: String, + }, + /// Relay rejected an invalid dispatch request. + InvalidRequest { + /// Human-readable validation failure. + message: String, + }, + /// The caller cancelled the operation. + Cancelled { + /// Human-readable cancellation context. + message: String, + }, + /// Relay could not complete the operation because of an internal failure. + Internal { + /// Human-readable internal failure. + message: String, + }, +} + +impl LlmCallErrorV2 { + /// Returns whether Relay classified this failure as retryable. + pub const fn is_retryable(&self) -> bool { + matches!( + self, + Self::Upstream { + retryable: true, + .. + } + ) + } +} + +/// Unary LLM continuation outcome delivered through native API v2. +#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum LlmCallOutcomeV2 { + /// Provider call completed successfully. + Success { + /// Provider response JSON. + response: Json, + }, + /// Provider call failed before producing a response. + Failure { + /// Structured Relay/provider failure. + error: LlmCallErrorV2, + }, +} + +/// Streaming LLM continuation event delivered through native API v2. +#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum LlmStreamEventV2 { + /// One provider stream event. + Chunk { + /// Provider event JSON. + chunk: Json, + }, + /// Provider stream completed successfully. + Done, + /// Provider stream failed before clean completion. + Failure { + /// Structured Relay/provider failure. + error: LlmCallErrorV2, + }, +} + +/// Receives one typed unary LLM continuation outcome. +/// +/// `outcome_json` contains one serialized [`LlmCallOutcomeV2`] and is borrowed +/// for the callback. +pub type NemoRelayNativeAsyncLlmResultCbV2 = + unsafe extern "C" fn(user_data: *mut c_void, outcome_json: *const NemoRelayNativeString); + +/// Receives one typed streaming LLM continuation event. +/// +/// `event_json` contains one serialized [`LlmStreamEventV2`] and is borrowed +/// for the callback. Return `false` to cancel downstream production after the +/// current callback. +pub type NemoRelayNativeAsyncLlmStreamCbV2 = + unsafe extern "C" fn(user_data: *mut c_void, event_json: *const NemoRelayNativeString) -> bool; + /// Incremental native LLM stream intercept callback. /// /// The callback owns `next` and `stream` and must release each exactly once. @@ -1062,6 +1229,64 @@ pub struct NemoRelayNativeHostApiV3 { ) -> NemoRelayStatus, } +/// ABI-v4 host extension implementing native API v2 typed LLM dispatch. +/// +/// 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. + pub async_llm_next_invoke_result_v2: unsafe extern "C" fn( + next: *const NemoRelayNativeAsyncNext, + dispatch_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncLlmResultCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Invokes a streaming LLM continuation with an explicit target and + /// structured incremental outcomes. + pub async_llm_next_invoke_stream_v2: unsafe extern "C" fn( + next: *const NemoRelayNativeAsyncNext, + dispatch_json: *const NemoRelayNativeString, + stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Registers a native API v2 unary LLM execution callback. + /// + /// Relay invokes the callback on its reusable blocking executor with the + /// active scope stack bound to that worker. Provider continuations invoked + /// through this table continue to execute on Relay's Tokio runtime. + pub plugin_context_register_async_llm_execution_v2: unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) + -> NemoRelayStatus, + /// Registers a native API v2 incremental LLM stream execution callback. + /// + /// Relay starts the callback on its reusable blocking executor and exposes + /// the bounded output stream to the caller concurrently. + pub plugin_context_register_async_llm_stream_execution_v2: + unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, +} + +unsafe impl Send for NemoRelayNativeHostApiV4 {} +unsafe impl Sync for NemoRelayNativeHostApiV4 {} + unsafe impl Send for NemoRelayNativeHostApiV3 {} unsafe impl Sync for NemoRelayNativeHostApiV3 {} @@ -1789,6 +2014,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_TYPED_LLM_DISPATCH + && 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) @@ -2617,6 +2850,54 @@ impl<'a> PluginContext<'a> { }) } + /// Registers a native API v2 unary LLM execution callback. + /// + /// # Safety + /// The callback and user data must remain valid until deregistration or + /// `free_fn`. The callback owns its completion and `next` handles and must + /// settle/release them exactly once. + pub unsafe fn register_async_llm_execution_v2_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + let Some(host) = self.host_api_v4() else { + return NemoRelayStatus::InvalidArg; + }; + self.with_name(name, |_, name| unsafe { + (host.plugin_context_register_async_llm_execution_v2)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + + /// Registers a native API v2 incremental LLM stream execution callback. + /// + /// # 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. + pub unsafe fn register_async_llm_stream_execution_v2_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + let Some(host) = self.host_api_v4() else { + return NemoRelayStatus::InvalidArg; + }; + self.with_name(name, |_, name| unsafe { + (host.plugin_context_register_async_llm_stream_execution_v2)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + fn with_name( &self, name: &str, @@ -3288,11 +3569,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_TYPED_LLM_DISPATCH + && 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 +3591,7 @@ impl OwnedHostApi { match self { Self::V1(host) => host, Self::V3(host) => &host.v1, + Self::V4(host) => &host.v3.v1, } } } @@ -3582,7 +3869,38 @@ 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 API v2 plugin descriptor for a Rust SDK plugin value. +/// +/// # 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. +pub unsafe fn export_plugin_v2( + host: *const NemoRelayNativeHostApiV1, + out: *mut NemoRelayNativePluginV1, + plugin: P, +) -> NemoRelayStatus { + 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_TYPED_LLM_DISPATCH, + std::mem::size_of::(), + || plugin, + ) } /// Initializes a native plugin descriptor from a constructor callback. @@ -3606,22 +3924,58 @@ 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` and `out` must satisfy [`export_plugin_v2`]'s requirements. +#[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_TYPED_LLM_DISPATCH, + 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 +4031,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 typed LLM dispatch 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/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 0161400fb..f113f1ab1 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -16,20 +16,21 @@ 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, - 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, + NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NativePlugin, + NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareKind, + 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, }; use serde_json::{Map, json}; @@ -335,6 +336,7 @@ static LLM_REQUEST_INTERCEPT_REGISTRATION: Mutex(), test_host().struct_size @@ -369,6 +371,9 @@ fn native_abi_v3_struct_sizes_are_self_describing() { 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::(), 472); + assert_eq!(host_api_v4_offsets(), [0, 440, 448, 456, 464]); assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 56); assert_eq!(plugin_offsets(), [0, 8, 16, 24, 32, 40, 48]); @@ -397,6 +402,9 @@ fn native_abi_v3_struct_sizes_are_self_describing() { 0, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212 ] ); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 232); + assert_eq!(host_api_v4_offsets(), [0, 216, 220, 224, 228]); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); assert_eq!(plugin_offsets(), [0, 4, 8, 12, 16, 20, 24]); @@ -406,6 +414,22 @@ fn native_abi_v3_struct_sizes_are_self_describing() { } } +fn host_api_v4_offsets() -> [usize; 5] { + [ + offset_of!(NemoRelayNativeHostApiV4, v3), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_result_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_stream_v2), + offset_of!( + NemoRelayNativeHostApiV4, + plugin_context_register_async_llm_execution_v2 + ), + offset_of!( + NemoRelayNativeHostApiV4, + plugin_context_register_async_llm_stream_execution_v2 + ), + ] +} + fn host_api_v3_offsets() -> [usize; 16] { [ offset_of!(NemoRelayNativeHostApiV3, v1), From 934d2b2e7554d239966605d82485edb786615873 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 30 Jul 2026 15:08:41 -0600 Subject: [PATCH 02/32] fix(plugin): bound native API v2 provider streams Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 287 +++++++++++++----- crates/core/tests/unit/native_plugin_tests.rs | 98 ++++-- crates/plugin/src/lib.rs | 55 +++- crates/plugin/tests/typed_callbacks.rs | 21 +- 4 files changed, 342 insertions(+), 119 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 750535fea..65be62230 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -54,9 +54,9 @@ use nemo_relay_plugin::{ LlmUpstreamFailureClassV2, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, - NemoRelayNativeAsyncLlmStreamCbV2, NemoRelayNativeAsyncMiddlewareCb, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, - NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, @@ -64,11 +64,12 @@ use nemo_relay_plugin::{ 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}; @@ -917,7 +918,10 @@ fn build_native_host_api_v4() -> NemoRelayNativeHostApiV4 { NemoRelayNativeHostApiV4 { v3, async_llm_next_invoke_result_v2: native_async_llm_next_invoke_result_v2, - async_llm_next_invoke_stream_v2: native_async_llm_next_invoke_stream_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_cancel_v2: native_async_llm_stream_cancel_v2, + async_llm_stream_release_v2: native_async_llm_stream_release_v2, plugin_context_register_async_llm_execution_v2: native_plugin_context_register_async_llm_execution_v2, plugin_context_register_async_llm_stream_execution_v2: @@ -1528,53 +1532,70 @@ struct NativeAsyncStreamCallbackGuard { active: bool, } -struct NativeAsyncLlmStreamCallbackGuardV2 { - cb: NemoRelayNativeAsyncLlmStreamCbV2, - user_data: usize, - stream: Arc, +struct NativeLlmProviderStreamV2 { + receiver: tokio::sync::Mutex>, + producer_abort: Mutex>, + runtime: tokio::runtime::Handle, + next_in_flight: AtomicBool, + cancelled: AtomicBool, _library_guard: Option>, +} + +impl Drop for NativeLlmProviderStreamV2 { + fn drop(&mut self) { + self.cancelled.store(true, Ordering::Release); + if let Some(abort) = self + .producer_abort + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + abort.abort(); + } + } +} + +struct NativeLlmStreamOpenCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, + user_data: usize, active: bool, + _library_guard: Option>, } -impl NativeAsyncLlmStreamCallbackGuardV2 { - fn emit(&self, event: &LlmStreamEventV2) -> bool { - let Some(event) = native_string_from_json( - &serde_json::to_value(event) - .expect("native API v2 stream events contain serializable Relay DTOs"), - ) else { - return false; - }; - let keep_going = unsafe { (self.cb)(self.user_data as *mut c_void, event) }; - unsafe { native_string_free(event) }; - keep_going +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 finish(&mut self, event: &LlmStreamEventV2) { - if self.active { - let _ = self.emit(event); - self.active = false; + fn failure(&mut self, error: &LlmCallErrorV2) { + if !self.active { + return; } + if let Some(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); + native_string_free(error); + } + } + self.active = false; } } -impl Drop for NativeAsyncLlmStreamCallbackGuardV2 { +impl Drop for NativeLlmStreamOpenCallbackGuardV2 { fn drop(&mut self) { - if !self.active { - return; + if self.active { + self.failure(&LlmCallErrorV2::Cancelled { + message: "typed native LLM stream setup was cancelled".into(), + }); } - let message = if self.stream.cancelled.load(Ordering::Acquire) { - "typed native LLM stream continuation was cancelled" - } else if self.stream.settled.load(Ordering::Acquire) { - "typed native LLM stream output settled" - } else { - "typed native LLM stream continuation ended without a terminal event" - }; - let event = LlmStreamEventV2::Failure { - error: LlmCallErrorV2::Cancelled { - message: message.into(), - }, - }; - let _ = self.emit(&event); } } @@ -2561,12 +2582,14 @@ async fn forward_native_async_next_stream_with( ); } callback_guard.finish(); -/// Invokes a streaming LLM continuation through native API v2. -unsafe extern "C" fn native_async_llm_next_invoke_stream_v2( +} + +/// Opens a streaming LLM continuation through native API v2. +unsafe extern "C" fn native_async_llm_next_open_stream_v2( next: *const NemoRelayNativeAsyncNext, dispatch_json: *const NemoRelayNativeString, output_stream: *const NemoRelayNativeAsyncStream, - cb: NemoRelayNativeAsyncLlmStreamCbV2, + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, user_data: *mut c_void, ) -> NemoRelayStatus { let Some(next) = (unsafe { (next as *const NativeAsyncNext).as_ref() }) else { @@ -2620,13 +2643,14 @@ unsafe extern "C" fn native_async_llm_next_invoke_stream_v2( .lock() .unwrap_or_else(|error| error.into_inner()); let next_fn = next_fn.clone(); + let provider_runtime = next.runtime.clone(); + let provider_library_guard = next._callback_user_data.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 = NativeAsyncLlmStreamCallbackGuardV2 { + let callback_guard = NativeLlmStreamOpenCallbackGuardV2 { cb, user_data, - stream: output_stream_for_task, _library_guard: next._callback_user_data.clone(), active: true, }; @@ -2640,41 +2664,89 @@ unsafe extern "C" fn native_async_llm_next_invoke_stream_v2( let mut callback_guard = callback_guard; let result = AssertUnwindSafe(async { match next_fn(request).await { - Ok(mut stream) => { - while let Some(item) = stream.next().await { - let event = match item { - Ok(chunk) => LlmStreamEventV2::Chunk { chunk }, - Err(error) => { - callback_guard.finish(&LlmStreamEventV2::Failure { + Ok(mut provider_stream) => { + let (sender, receiver) = + tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + let provider = Arc::new(NativeLlmProviderStreamV2 { + receiver: tokio::sync::Mutex::new(receiver), + producer_abort: Mutex::new(None), + runtime: provider_runtime.clone(), + next_in_flight: AtomicBool::new(false), + cancelled: AtomicBool::new(false), + _library_guard: provider_library_guard.clone(), + }); + let output_for_producer = Arc::clone(&output_stream_for_task); + let producer = provider_runtime.spawn(async move { + let mut failed = false; + while let Some(item) = provider_stream.next().await { + let terminal = item.is_err(); + let event = match item { + Ok(chunk) => LlmStreamEventV2::Chunk { chunk }, + Err(error) => LlmStreamEventV2::Failure { error: typed_llm_failure(error), - }); - return; + }, + }; + if sender.send(event).await.is_err() || terminal { + failed = terminal; + break; } - }; - if !callback_guard.emit(&event) { - callback_guard.active = false; + } + if !failed + && !output_for_producer.cancelled.load(Ordering::Acquire) + && !output_for_producer.settled.load(Ordering::Acquire) + { + let _ = sender.send(LlmStreamEventV2::Done).await; + } + output_for_producer + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&tokio::task::id()); + }); + let producer_abort = producer.abort_handle(); + let producer_id = producer.id(); + *provider + .producer_abort + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some(producer_abort.clone()); + { + 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) + { + producer_abort.abort(); + callback_guard.failure(&LlmCallErrorV2::Cancelled { + message: + "typed native LLM stream output settled during setup" + .into(), + }); return; } + output_stream_for_task + .downstream_aborts + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert(producer_id, producer_abort); } - callback_guard.finish(&LlmStreamEventV2::Done); + callback_guard.success(provider); } Err(error) => { - callback_guard.finish(&LlmStreamEventV2::Failure { - error: typed_llm_failure(error), - }); + callback_guard.failure(&typed_llm_failure(error)); } } }) .catch_unwind() .await; if let Err(payload) = result { - callback_guard.finish(&LlmStreamEventV2::Failure { - error: LlmCallErrorV2::Internal { - message: format!( - "typed native LLM stream continuation panicked: {}", - panic_payload_message(payload.as_ref()) - ), - }, + callback_guard.failure(&LlmCallErrorV2::Internal { + message: format!( + "typed native LLM stream continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ), }); } }) @@ -2692,6 +2764,79 @@ unsafe extern "C" fn native_async_llm_next_invoke_stream_v2( NemoRelayStatus::Ok } +/// 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; + }; + if stream.cancelled.load(Ordering::Acquire) { + set_native_last_error("native API v2 provider stream is cancelled"); + return NemoRelayStatus::InvalidArg; + } + if stream.next_in_flight.swap(true, Ordering::AcqRel) { + set_native_last_error("native API v2 provider stream already has a pending next operation"); + return NemoRelayStatus::InvalidArg; + } + unsafe { Arc::increment_strong_count(stream as *const NativeLlmProviderStreamV2) }; + let stream = unsafe { Arc::from_raw(stream as *const NativeLlmProviderStreamV2) }; + let runtime = stream.runtime.clone(); + let user_data = user_data as usize; + runtime.spawn(async move { + let event = stream + .receiver + .lock() + .await + .recv() + .await + .unwrap_or_else(|| LlmStreamEventV2::Failure { + error: LlmCallErrorV2::Cancelled { + message: "native API v2 provider stream closed without a terminal event".into(), + }, + }); + if let Some(event) = native_string_from_json( + &serde_json::to_value(&event) + .expect("native API v2 stream events contain serializable Relay DTOs"), + ) { + unsafe { + cb(user_data as *mut c_void, event); + native_string_free(event); + } + } + stream.next_in_flight.store(false, Ordering::Release); + }); + NemoRelayStatus::Ok +} + +/// Cancels a native API v2 provider stream. +unsafe extern "C" fn native_async_llm_stream_cancel_v2( + stream: *const NemoRelayNativeLlmStreamV2, +) -> NemoRelayStatus { + let Some(stream) = (unsafe { (stream as *const NativeLlmProviderStreamV2).as_ref() }) else { + return NemoRelayStatus::NullPointer; + }; + stream.cancelled.store(true, Ordering::Release); + if let Some(abort) = stream + .producer_abort + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + abort.abort(); + } + 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() { + unsafe { drop(Arc::from_raw(stream as *const NativeLlmProviderStreamV2)) }; + } +} + fn wrap_native_async_tool_json( instance: Arc, cb: NemoRelayNativeAsyncMiddlewareCb, diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 56c24bc95..ff3f8bab7 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -130,33 +130,38 @@ unsafe extern "C" fn reject_unexpected_typed_llm_result( panic!("invalid dispatch unexpectedly invoked its callback"); } -#[derive(Default)] -struct TypedLlmStreamCallbackState { - events: Mutex>, - terminal: tokio::sync::Notify, +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 = if error_json.is_null() { + Ok(stream as usize) + } else { + parse_json_arg(error_json, "typed LLM stream open error") + .and_then(|value| { + serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) + }) + .map_err(|status| LlmCallErrorV2::Internal { + message: format!("invalid typed stream open error: {status:?}"), + }) + }; + let _ = sender.send(result); } -unsafe extern "C" fn record_typed_llm_stream_result( +unsafe extern "C" fn record_typed_llm_stream_next( user_data: *mut c_void, event_json: *const NemoRelayNativeString, -) -> bool { - let state = unsafe { &*(user_data as *const TypedLlmStreamCallbackState) }; +) { + let sender = + unsafe { Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) }; let event: LlmStreamEventV2 = parse_json_arg(event_json, "typed LLM stream result") .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) .expect("host emitted a valid typed LLM stream event"); - let terminal = matches!( - event, - LlmStreamEventV2::Done | LlmStreamEventV2::Failure { .. } - ); - state - .events - .lock() - .unwrap_or_else(|error| error.into_inner()) - .push(event); - if terminal { - state.terminal.notify_one(); - } - true + let _ = sender.send(event); } #[derive(Default)] @@ -1763,7 +1768,6 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { }); let output_stream_ref = Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; - let callback_state = Arc::new(TypedLlmStreamCallbackState::default()); let dispatch = native_string_from_json( &serde_json::to_value(LlmDispatchRequestV2 { request: LlmRequest { @@ -1779,28 +1783,57 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { ) .unwrap(); + let (open_sender, open_receiver) = + tokio::sync::oneshot::channel::>(); assert_eq!( unsafe { - native_async_llm_next_invoke_stream_v2( + native_async_llm_next_open_stream_v2( next_ref, dispatch, output_stream_ref, - record_typed_llm_stream_result, - Arc::as_ptr(&callback_state).cast_mut().cast(), + record_typed_llm_stream_open, + Box::into_raw(Box::new(open_sender)).cast(), ) }, NemoRelayStatus::Ok ); - runtime.block_on(async { - tokio::time::timeout(Duration::from_secs(1), callback_state.terminal.notified()) + let provider_stream = runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), open_receiver) .await - .expect("typed LLM stream should emit a terminal event"); - }); + .expect("typed LLM stream should open") + .expect("open callback should be delivered") + .expect("provider stream should open") + }) as *const NemoRelayNativeLlmStreamV2; + let mut events = Vec::new(); + loop { + let (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_stream, + record_typed_llm_stream_next, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + 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, + LlmStreamEventV2::Done | LlmStreamEventV2::Failure { .. } + ); + events.push(event); + if terminal { + break; + } + } assert_eq!( - *callback_state - .events - .lock() - .unwrap_or_else(|error| error.into_inner()), + events, vec![ LlmStreamEventV2::Chunk { chunk: json!({ @@ -1825,6 +1858,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { 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); diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 0fe9102d3..04d92921a 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -874,6 +874,18 @@ pub struct NemoRelayNativeAsyncStream { _marker: PhantomData<(*mut u8, PhantomPinned)>, } +/// Opaque host-owned provider stream returned by native API v2 dispatch. +/// +/// The plugin requests one item at a time with the v2 host table, then cancels +/// or releases the handle exactly once. Relay pumps provider output into a +/// bounded queue so a plugin that consumes slowly applies backpressure without +/// blocking a runtime worker. +#[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 @@ -1056,13 +1068,22 @@ pub enum LlmStreamEventV2 { pub type NemoRelayNativeAsyncLlmResultCbV2 = unsafe extern "C" fn(user_data: *mut c_void, outcome_json: *const NemoRelayNativeString); -/// Receives one typed streaming LLM continuation event. +/// 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. /// /// `event_json` contains one serialized [`LlmStreamEventV2`] and is borrowed -/// for the callback. Return `false` to cancel downstream production after the -/// current callback. -pub type NemoRelayNativeAsyncLlmStreamCbV2 = - unsafe extern "C" fn(user_data: *mut c_void, event_json: *const NemoRelayNativeString) -> bool; +/// for the callback. Only one `next` operation may be active per stream. +pub type NemoRelayNativeAsyncLlmStreamNextCbV2 = + unsafe extern "C" fn(user_data: *mut c_void, event_json: *const NemoRelayNativeString); /// Incremental native LLM stream intercept callback. /// @@ -1246,15 +1267,29 @@ pub struct NemoRelayNativeHostApiV4 { cb: NemoRelayNativeAsyncLlmResultCbV2, user_data: *mut c_void, ) -> NemoRelayStatus, - /// Invokes a streaming LLM continuation with an explicit target and - /// structured incremental outcomes. - pub async_llm_next_invoke_stream_v2: unsafe extern "C" fn( + /// 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`. + pub async_llm_next_open_stream_v2: unsafe extern "C" fn( next: *const NemoRelayNativeAsyncNext, dispatch_json: *const NemoRelayNativeString, - stream: *const NemoRelayNativeAsyncStream, - cb: NemoRelayNativeAsyncLlmStreamCbV2, + output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, + user_data: *mut c_void, + ) -> NemoRelayStatus, + /// Requests one provider event from a native API v2 stream. + pub async_llm_stream_next_v2: unsafe extern "C" fn( + stream: *const NemoRelayNativeLlmStreamV2, + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, user_data: *mut c_void, ) -> NemoRelayStatus, + /// Cancels provider production for a native API v2 stream. + pub async_llm_stream_cancel_v2: + unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2) -> NemoRelayStatus, + /// Releases the plugin-owned provider-stream reference. + pub async_llm_stream_release_v2: + unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2), /// Registers a native API v2 unary LLM execution callback. /// /// Relay invokes the callback on its reusable blocking executor with the diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index f113f1ab1..3ca498710 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -372,8 +372,11 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 472); - assert_eq!(host_api_v4_offsets(), [0, 440, 448, 456, 464]); + assert_eq!(size_of::(), 496); + assert_eq!( + host_api_v4_offsets(), + [0, 440, 448, 456, 464, 472, 480, 488] + ); assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 56); assert_eq!(plugin_offsets(), [0, 8, 16, 24, 32, 40, 48]); @@ -403,8 +406,11 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 232); - assert_eq!(host_api_v4_offsets(), [0, 216, 220, 224, 228]); + assert_eq!(size_of::(), 244); + assert_eq!( + host_api_v4_offsets(), + [0, 216, 220, 224, 228, 232, 236, 240] + ); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); assert_eq!(plugin_offsets(), [0, 4, 8, 12, 16, 20, 24]); @@ -414,11 +420,14 @@ fn native_abi_v3_struct_sizes_are_self_describing() { } } -fn host_api_v4_offsets() -> [usize; 5] { +fn host_api_v4_offsets() -> [usize; 8] { [ offset_of!(NemoRelayNativeHostApiV4, v3), offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_result_v2), - offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_stream_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_open_stream_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_next_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_cancel_v2), + offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_release_v2), offset_of!( NemoRelayNativeHostApiV4, plugin_context_register_async_llm_execution_v2 From 1e1545176c92e7fe7c88edf8809c27b45bed1986 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 30 Jul 2026 16:38:52 -0600 Subject: [PATCH 03/32] fix(plugin): harden native API v2 streams Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 54 ++- crates/core/tests/unit/native_plugin_tests.rs | 310 ++++++++++++++++++ crates/plugin/README.md | 31 +- .../dynamic-plugins/native-dynamic/about.mdx | 62 +++- 4 files changed, 447 insertions(+), 10 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 65be62230..429b4db9e 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -6,9 +6,9 @@ #[cfg(test)] use std::cell::Cell; use std::cell::RefCell; -use std::collections::HashMap; #[cfg(test)] use std::collections::HashSet; +use std::collections::{BTreeMap, HashMap}; use std::ffi::c_void; use std::future::Future; use std::panic::{AssertUnwindSafe, catch_unwind}; @@ -2295,6 +2295,8 @@ unsafe extern "C" fn native_async_next_invoke_result( const NATIVE_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; const NATIVE_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; const NATIVE_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; +const NATIVE_API_V2_MAX_FAILURE_BODY_BYTES: usize = 16 * 1024; +const NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES: usize = 1024; fn prepare_typed_llm_dispatch( mut dispatch: LlmDispatchRequestV2, @@ -2340,8 +2342,8 @@ fn typed_llm_failure(error: FlowError) -> LlmCallErrorV2 { class, retryable: failure.is_retryable(), status: failure.status, - body: failure.body, - headers: failure.headers, + body: bounded_utf8(failure.body, NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), + headers: safe_native_api_v2_failure_headers(failure.headers), } } FlowError::GuardrailRejected(message) => LlmCallErrorV2::GuardrailRejected { message }, @@ -2352,6 +2354,47 @@ fn typed_llm_failure(error: FlowError) -> LlmCallErrorV2 { } } +fn safe_native_api_v2_failure_headers( + headers: BTreeMap, +) -> BTreeMap { + headers + .into_iter() + .filter_map(|(name, value)| { + let normalized = name.to_ascii_lowercase(); + let safe = 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" + ); + safe.then(|| { + ( + normalized, + bounded_utf8(value, NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES), + ) + }) + }) + .collect() +} + +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_string() +} + unsafe fn invoke_typed_llm_result_callback( cb: NemoRelayNativeAsyncLlmResultCbV2, user_data: *mut c_void, @@ -2797,6 +2840,10 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( message: "native API v2 provider stream closed without a terminal event".into(), }, }); + // The pull operation is complete before callback delivery. Clearing + // the guard first lets a callback wake plugin code that immediately + // requests the next event without racing this task's epilogue. + stream.next_in_flight.store(false, Ordering::Release); if let Some(event) = native_string_from_json( &serde_json::to_value(&event) .expect("native API v2 stream events contain serializable Relay DTOs"), @@ -2806,7 +2853,6 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( native_string_free(event); } } - stream.next_in_flight.store(false, Ordering::Release); }); NemoRelayStatus::Ok } diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index ff3f8bab7..d31846949 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -1664,6 +1664,37 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { } } +#[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(NATIVE_API_V2_MAX_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(), + "x".repeat(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES + 50), + ), + ]), + class: UpstreamFailureClass::RetryableStatus, + })); + + let LlmCallErrorV2::Upstream { body, headers, .. } = error else { + panic!("expected an upstream failure"); + }; + assert_eq!(body.len(), NATIVE_API_V2_MAX_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(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES) + ); + 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() @@ -1865,6 +1896,285 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { } } +#[test] +fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (provider_sender, provider_receiver) = + tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + let provider = Arc::new(NativeLlmProviderStreamV2 { + receiver: tokio::sync::Mutex::new(provider_receiver), + producer_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + next_in_flight: AtomicBool::new(false), + cancelled: AtomicBool::new(false), + _library_guard: None, + }); + let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; + + let (first_sender, first_receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next, + Box::into_raw(Box::new(first_sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + + 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)) }; + + assert_eq!( + unsafe { native_async_llm_stream_cancel_v2(provider_ref) }, + NemoRelayStatus::Ok + ); + assert_eq!( + unsafe { native_async_llm_stream_cancel_v2(provider_ref) }, + NemoRelayStatus::Ok + ); + drop(provider_sender); + // Releasing the plugin's reference while `next` is pending is safe because + // the callback task retains its own provider-stream reference. + unsafe { native_async_llm_stream_release_v2(provider_ref) }; + + let event = runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(1), first_receiver).await }) + .expect("cancelled provider next should settle") + .expect("provider next callback should be delivered"); + assert!(matches!( + event, + LlmStreamEventV2::Failure { + error: LlmCallErrorV2::Cancelled { .. } + } + )); +} + +#[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(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model"}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + url: "https://provider.example/v1/chat/completions".into(), + route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + }, + }) + .unwrap(), + ) + .unwrap(); + + let mut receivers = Vec::with_capacity(DISPATCH_COUNT); + for _ in 0..DISPATCH_COUNT { + let (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + receivers.push(receiver); + } + + 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(LlmCallOutcomeV2::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_handles_64_concurrent_100_event_provider_streams() { + const STREAM_COUNT: usize = 64; + const EVENT_COUNT: usize = 100; + assert_eq!(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY, 32); + + 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_sender, output_receiver) = + tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + let output_stream = Arc::new(NativeAsyncStream { + sender: Mutex::new(Some(output_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_ref = + Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; + let dispatch = native_string_from_json( + &serde_json::to_value(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"model": "provider/model", "stream": true}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + url: "https://provider.example/v1/chat/completions".into(), + route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + }, + }) + .unwrap(), + ) + .unwrap(); + + let mut open_receivers = Vec::with_capacity(STREAM_COUNT); + for _ in 0..STREAM_COUNT { + let (sender, receiver) = + tokio::sync::oneshot::channel::>(); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + next_ref, + dispatch, + output_stream_ref, + record_typed_llm_stream_open, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + open_receivers.push(receiver); + } + + 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 (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_stream, + record_typed_llm_stream_next, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + match receiver.await.expect("next callback should be delivered") { + LlmStreamEventV2::Chunk { .. } => chunks += 1, + LlmStreamEventV2::Done => break, + LlmStreamEventV2::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 native_continuation_context_observation( expected_stack: &ScopeStackHandle, expected_event_uuid: uuid::Uuid, diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 8df78823f..e05f8ce2e 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. @@ -49,6 +49,9 @@ the dynamic-library boundary on the stable C-compatible ABI. - **Raw async middleware**: Completion-based raw registrations for plugins that need asynchronous guardrails, intercepts, or event sanitizers. Typed Rust callbacks remain synchronous convenience APIs. +- **Native API v2 LLM dispatch**: Register v2-only execution callbacks that + send an explicit provider target through Relay and receive buffered JSON, + structured failures, or a bounded host-owned provider stream. ## Installation @@ -94,6 +97,32 @@ 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`. During registration, +`PluginContext::host_api_v4` exposes the C-safe typed LLM dispatch table and +the raw v2 registration helpers. + +The plugin provides JSON plus an absolute target URL and protocol route. Relay +returns response JSON or a structured provider failure. Streaming dispatch +returns an opaque host-owned stream; request one JSON event at a time, then +cancel and release it exactly once. No Rust future, trait object, +`serde_json::Value`, or allocator-owned Rust string crosses the ABI boundary. + +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. + ## Documentation - [NeMo Relay documentation](https://docs.nvidia.com/nemo/relay) diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 82167c28e..af24486e4 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. @@ -116,11 +116,63 @@ 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 targets, failures, and provider streams | + +Relay retains native API v1 unchanged. 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`, absolute HTTP(S) target URL, and target route +(`openai_chat`, `openai_responses`, or `anthropic_messages`). Relay returns +either: + +- buffered response JSON; +- a structured provider failure with stable class, retryability, optional HTTP + status, a body bounded to 16 KiB, and a conservative safe-header allowlist; or +- a host-owned provider stream pulled one JSON event at a time, including typed + setup and late failures. + +The host converts the public target into its private dispatch metadata. Plugins +do not construct private Relay headers. Provider stream production and plugin +output each use bounded 32-event queues. Dropping or cancelling a stream stops +provider production, and the library stays loaded until all callbacks and +streams release their handles. + +V2 callbacks run on Relay's blocking executor with the active scope stack +restored. Continuations run on Relay's existing Tokio runtime; Relay does not +create an OS thread per continuation. Separate callback invocations can run +concurrently and have no stable OS-thread affinity. + +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 +`PluginContext::host_api_v4` and the raw v2 LLM registration methods when the +plugin needs the typed dispatch contract. + +## 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 +182,7 @@ extern "C" fn nemo_relay_register_plugin( ) -> NemoRelayStatus ``` -The v3 host table retains the frozen legacy prefix and appends a +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 From 59f59d71919c21c205af2fc3381168fddecfce9e Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 31 Jul 2026 08:46:43 -0600 Subject: [PATCH 04/32] refactor(plugin): use targeted HTTP LLM dispatch Signed-off-by: Bryan Bednarski --- crates/cli/src/gateway/mod.rs | 84 ++++- crates/cli/src/server/mod.rs | 9 + .../tests/coverage/shared/gateway_tests.rs | 68 ++++ crates/core/src/api/runtime.rs | 3 + .../src/api/runtime/continuation_context.rs | 17 + .../src/api/runtime/llm_dispatch_context.rs | 87 ++++++ crates/core/src/plugin/dynamic/native.rs | 225 ++++++++----- crates/core/tests/unit/native_plugin_tests.rs | 295 ++++++++++++++---- crates/plugin/README.md | 19 +- crates/plugin/src/lib.rs | 123 ++++---- crates/plugin/tests/typed_callbacks.rs | 65 +++- .../dynamic-plugins/native-dynamic/about.mdx | 37 ++- 12 files changed, 801 insertions(+), 231 deletions(-) create mode 100644 crates/core/src/api/runtime/llm_dispatch_context.rs diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 76feccd4c..c9a39942f 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -26,7 +26,8 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::runtime::{ - LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, + LlmCodecIdentity, LlmDispatchTargetContext, LlmExecutionNextFn, LlmJsonStream, + LlmStreamExecutionNextFn, TASK_SCOPE_STACK, current_llm_dispatch_target, }; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ @@ -176,6 +177,7 @@ async fn run_unmanaged_gateway( &prepared.body_bytes, &prepared.headers, None, + None, ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config), ) .await?; @@ -318,6 +320,7 @@ fn build_buffered_func( upstream_failures: CapturedUpstreamFailuresRef, ) -> LlmExecutionNextFn { let http = state.http.clone(); + let targeted_http = state.targeted_http.clone(); let method = prepared.method.clone(); let url = prepared.upstream_url.clone(); let body_bytes = prepared.body_bytes.clone(); @@ -326,6 +329,7 @@ fn build_buffered_func( ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config); Arc::new(move |request| { let http = http.clone(); + let targeted_http = targeted_http.clone(); let forwarding = forwarding.clone(); let method = method.clone(); let url = url.clone(); @@ -333,14 +337,17 @@ fn build_buffered_func( let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); Box::pin(async move { - let retry_aware = retry_aware_dispatch(&request); + let typed_target = current_llm_dispatch_target(); + let retry_aware = typed_target.is_some() || retry_aware_dispatch(&request); + let http = typed_target.as_ref().map_or(&http, |_| &targeted_http); let response = match forward_upstream_request( - &http, + http, &method, &url, &body_bytes, &headers, Some(&request), + typed_target.as_ref(), forwarding, ) .await @@ -512,6 +519,7 @@ fn build_streaming_func( upstream_failures: CapturedUpstreamFailuresRef, ) -> LlmStreamExecutionNextFn { let http = state.http.clone(); + let targeted_http = state.targeted_http.clone(); let method = prepared.method.clone(); let url = prepared.upstream_url.clone(); let body_bytes = prepared.body_bytes.clone(); @@ -520,6 +528,7 @@ fn build_streaming_func( ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config); Arc::new(move |request| { let http = http.clone(); + let targeted_http = targeted_http.clone(); let forwarding = forwarding.clone(); let method = method.clone(); let url = url.clone(); @@ -527,14 +536,17 @@ fn build_streaming_func( let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); Box::pin(async move { - let retry_aware = retry_aware_dispatch(&request); + let typed_target = current_llm_dispatch_target(); + let retry_aware = typed_target.is_some() || retry_aware_dispatch(&request); + let http = typed_target.as_ref().map_or(&http, |_| &targeted_http); let response = match forward_upstream_request( - &http, + http, &method, &url, &body_bytes, &headers, Some(&request), + typed_target.as_ref(), forwarding, ) .await @@ -788,6 +800,7 @@ async fn forward_upstream_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, + typed_target: Option<&LlmDispatchTargetContext>, forwarding: ProviderForwarding, ) -> Result { debug_assert_eq!( @@ -797,16 +810,18 @@ async fn forward_upstream_request( .provider_credential_present(), crate::provider_auth::has_provider_credential(headers) ); - let effective = effective_dispatch_request( + let effective = effective_dispatch_request_with_target( body_bytes, headers, effective_request, + typed_target, url, + method, forwarding.source_route, ); let configured_auth_header = forwarding.configured_auth_header(effective.target_route); let mut upstream = http - .request(method.clone(), &effective.url) + .request(effective.method.clone(), &effective.url) .body(effective.body_bytes.clone()); for (name, value) in &effective.headers { if should_forward_request_header(name, &effective.headers) { @@ -830,6 +845,7 @@ async fn forward_upstream_request( struct EffectiveUpstreamRequest { body_bytes: Bytes, headers: HeaderMap, + method: Method, url: String, target_route: ProviderRoute, credential_policy: TargetCredentialPolicy, @@ -857,21 +873,67 @@ fn effective_upstream_request( (effective.body_bytes, effective.headers) } +#[cfg(test)] fn effective_dispatch_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, url: &str, route: ProviderRoute, +) -> EffectiveUpstreamRequest { + effective_dispatch_request_with_target( + body_bytes, + headers, + effective_request, + None, + url, + &Method::POST, + route, + ) +} + +fn effective_dispatch_request_with_target( + body_bytes: &Bytes, + headers: &HeaderMap, + effective_request: Option<&LlmRequest>, + typed_target: Option<&LlmDispatchTargetContext>, + url: &str, + method: &Method, + route: ProviderRoute, ) -> EffectiveUpstreamRequest { let mut headers = headers.clone(); strip_internal_dispatch_headers(&mut headers); let Some(request) = effective_request else { - return source_request(body_bytes, headers, url, route); + return source_request(body_bytes, headers, method, url, route); }; let Some((body_bytes, body_reencoded)) = reencode_request_body(request, body_bytes) else { - return source_request(body_bytes, headers, url, route); + return source_request(body_bytes, headers, method, url, route); }; + if let Some(target) = typed_target { + let mut target_headers = HeaderMap::new(); + for (name, value) in target.headers() { + let name = HeaderName::from_bytes(name.as_bytes()) + .expect("native API v2 target header names were validated"); + let value = HeaderValue::from_str(value) + .expect("native API v2 target header values were validated"); + target_headers.insert(name, value); + } + target_headers + .entry(header::CONTENT_TYPE) + .or_insert(HeaderValue::from_static("application/json")); + let target_route = ProviderRoute::from_dispatch_override(target.route()) + .expect("native API v2 target routes originate from a typed route enum"); + let method = Method::from_bytes(target.method().as_bytes()) + .expect("native API v2 target methods were validated"); + return EffectiveUpstreamRequest { + body_bytes, + headers: target_headers, + method, + url: target.url().to_owned(), + target_route, + credential_policy: TargetCredentialPolicy::ExplicitTarget, + }; + } let overrides = dispatch_overrides(&request.headers); let credential_policy = if overrides.is_explicit_target() { crate::provider_auth::remove_provider_credentials(&mut headers); @@ -889,6 +951,7 @@ fn effective_dispatch_request( EffectiveUpstreamRequest { body_bytes, headers, + method: method.clone(), url: overrides.resolve_url(url), target_route: overrides.route.unwrap_or(route), credential_policy, @@ -898,12 +961,14 @@ fn effective_dispatch_request( fn source_request( body_bytes: &Bytes, headers: HeaderMap, + method: &Method, url: &str, route: ProviderRoute, ) -> EffectiveUpstreamRequest { EffectiveUpstreamRequest { body_bytes: body_bytes.clone(), headers, + method: method.clone(), url: url.to_string(), target_route: route, credential_policy: TargetCredentialPolicy::SourceOrEnvironment, @@ -1106,6 +1171,7 @@ async fn passthrough_streaming( &prepared.body_bytes, &prepared.headers, None, + None, ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config), ) .await?; diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index c4ca3ea62..b52a07e36 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -63,6 +63,7 @@ pub(crate) struct AppState { pub(crate) transparent_proxy_credential: Option, pub(crate) http: Client, + pub(crate) targeted_http: Client, pub(crate) sessions: SessionManager, pub(crate) last_activity: Arc>, pub(crate) bootstrap_shutdown: Option, @@ -512,6 +513,13 @@ impl AppState { .read_timeout(HTTP_READ_TIMEOUT) .build() .expect("gateway HTTP client configuration is valid"); + let targeted_http = Client::builder() + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .timeout(HTTP_REQUEST_TIMEOUT) + .read_timeout(HTTP_READ_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("targeted gateway HTTP client configuration is valid"); Self { config, bootstrap_fingerprint, @@ -519,6 +527,7 @@ impl AppState { require_provider_client_token, transparent_proxy_credential, http, + targeted_http, sessions, last_activity: Arc::new(Mutex::new(Instant::now())), bootstrap_shutdown, diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index a775ddf93..d71998f13 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -13,6 +13,7 @@ use axum::response::IntoResponse; use http_body_util::BodyExt; use reqwest::Client; use serde_json::{Map, json}; +use std::collections::BTreeMap; use tokio::io::{AsyncReadExt, AsyncWriteExt}; fn test_http_client() -> Client { @@ -702,6 +703,71 @@ fn internal_dispatch_controls_are_consumed_and_never_forwarded() { assert!(retry_aware_dispatch(&request)); } +#[test] +fn typed_dispatch_target_overrides_legacy_headers_without_exposing_transport_data() { + let original_body = Bytes::from_static(br#"{"model":"original"}"#); + let mut original_headers = HeaderMap::new(); + original_headers.insert( + header::AUTHORIZATION, + HeaderValue::from_static("Bearer source-secret"), + ); + original_headers.insert("x-source", HeaderValue::from_static("source")); + let request = LlmRequest { + headers: Map::from_iter([ + ( + INTERNAL_DISPATCH_URL_HEADER.to_string(), + json!("http://attacker.invalid/v1/chat/completions"), + ), + ("authorization".into(), json!("Bearer request-secret")), + ("x-request".into(), json!("request")), + ]), + content: json!({"model": "selected"}), + }; + let target = LlmDispatchTargetContext::new( + "POST".into(), + "http://selected.invalid/v1/responses".into(), + "openai_responses".into(), + BTreeMap::from([ + ("authorization".into(), "Bearer target-secret".into()), + ("x-target".into(), "selected".into()), + ]), + ); + + let effective = effective_dispatch_request_with_target( + &original_body, + &original_headers, + Some(&request), + Some(&target), + "http://default.invalid/v1/chat/completions", + &Method::GET, + ProviderRoute::OpenAiChatCompletions, + ); + + assert_eq!(effective.method, Method::POST); + assert_eq!( + effective.url, + "http://selected.invalid/v1/responses".to_string() + ); + assert_eq!(effective.target_route, ProviderRoute::OpenAiResponses); + assert_eq!( + effective.headers.get(header::AUTHORIZATION).unwrap(), + "Bearer target-secret" + ); + assert_eq!(effective.headers.get("x-target").unwrap(), "selected"); + assert_eq!( + effective.headers.get(header::CONTENT_TYPE).unwrap(), + "application/json" + ); + assert!(effective.headers.get("x-source").is_none()); + assert!(effective.headers.get("x-request").is_none()); + assert!( + effective + .headers + .get(INTERNAL_DISPATCH_URL_HEADER) + .is_none() + ); +} + #[test] fn explicit_keyless_target_drops_source_credentials() { let mut source_headers = HeaderMap::new(); @@ -1944,6 +2010,7 @@ async fn passthrough_rejects_unsupported_provider_path_directly() { require_provider_client_token: false, transparent_proxy_credential: None, http: test_http_client(), + targeted_http: test_http_client(), sessions: SessionManager::new(config), last_activity: std::sync::Arc::new(std::sync::Mutex::new(std::time::Instant::now())), bootstrap_shutdown: None, @@ -1982,6 +2049,7 @@ async fn models_rejects_non_get_requests_directly() { require_provider_client_token: false, transparent_proxy_credential: None, http: test_http_client(), + targeted_http: test_http_client(), sessions: SessionManager::new(config), last_activity: std::sync::Arc::new(std::sync::Mutex::new(std::time::Instant::now())), bootstrap_shutdown: None, diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 12d6612c3..b4e9165c6 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,8 @@ pub use continuation_context::MiddlewareContinuationContext; #[cfg(test)] pub(crate) use continuation_context::MiddlewareContinuationLease; pub use global::global_context; +#[doc(hidden)] +pub use llm_dispatch_context::{LlmDispatchTargetContext, current_llm_dispatch_target}; 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..71666a145 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, 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, @@ -103,6 +106,20 @@ impl MiddlewareContinuationContext { } } + /// Invoke a callback and poll its future with the captured Relay context and typed LLM target. + #[doc(hidden)] + pub async fn invoke_with_llm_dispatch_target( + &self, + target: LlmDispatchTargetContext, + callback: C, + ) -> F::Output + where + C: FnOnce() -> F, + F: Future, + { + scope_llm_dispatch_target(target, self.run(async move { callback().await })).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..3b24cc619 --- /dev/null +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Invocation-scoped transport target for managed LLM continuations. + +use std::collections::BTreeMap; +use std::future::Future; + +tokio::task_local! { + static TASK_LLM_DISPATCH_TARGET: LlmDispatchTargetContext; +} + +/// Validated provider transport target bound to one LLM continuation invocation. +/// +/// This internal bridge keeps transport data out of [`crate::api::llm::LlmRequest`] +/// while allowing the terminal gateway callback to dispatch a plugin-selected +/// provider request. +#[doc(hidden)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LlmDispatchTargetContext { + method: String, + url: String, + route: String, + headers: BTreeMap, +} + +impl LlmDispatchTargetContext { + /// Construct a validated target from the native-plugin adapter. + #[doc(hidden)] + #[must_use] + pub fn new( + method: String, + url: String, + route: String, + headers: BTreeMap, + ) -> Self { + Self { + method, + url, + route, + headers, + } + } + + /// HTTP method selected for this invocation. + #[doc(hidden)] + #[must_use] + pub fn method(&self) -> &str { + &self.method + } + + /// Absolute provider URL selected for this invocation. + #[doc(hidden)] + #[must_use] + pub fn url(&self) -> &str { + &self.url + } + + /// Relay provider-route identifier selected for this invocation. + #[doc(hidden)] + #[must_use] + pub fn route(&self) -> &str { + &self.route + } + + /// Explicit provider headers selected for this invocation. + #[doc(hidden)] + #[must_use] + pub fn headers(&self) -> &BTreeMap { + &self.headers + } +} + +/// Return the typed target bound to the current continuation invocation. +#[doc(hidden)] +#[must_use] +pub fn current_llm_dispatch_target() -> Option { + TASK_LLM_DISPATCH_TARGET.try_with(Clone::clone).ok() +} + +/// Poll a future with one typed target bound to its continuation invocation. +pub(crate) async fn scope_llm_dispatch_target( + target: LlmDispatchTargetContext, + future: F, +) -> F::Output { + TASK_LLM_DISPATCH_TARGET.scope(target, future).await +} diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 429b4db9e..cf078fbf6 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -24,9 +24,10 @@ 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, }; @@ -50,13 +51,14 @@ use crate::plugin::{ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ - LlmCallErrorV2, LlmCallOutcomeV2, LlmDispatchRequestV2, LlmStreamEventV2, - LlmUpstreamFailureClassV2, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, - NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NemoRelayNativeAsyncCallbackState, - NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, - NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + LlmCallFailureV2, LlmCallOutcomeV2, LlmDispatchRequestV2, LlmHttpFailureV2, + LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmStreamEventV2, NEMO_RELAY_NATIVE_ABI_VERSION, + NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, + NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, + NemoRelayNativeAsyncLlmResultCbV2, NemoRelayNativeAsyncLlmStreamNextCbV2, + NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, @@ -1572,7 +1574,7 @@ impl NativeLlmStreamOpenCallbackGuardV2 { self.active = false; } - fn failure(&mut self, error: &LlmCallErrorV2) { + fn failure(&mut self, error: &LlmCallFailureV2) { if !self.active { return; } @@ -1592,9 +1594,10 @@ impl NativeLlmStreamOpenCallbackGuardV2 { impl Drop for NativeLlmStreamOpenCallbackGuardV2 { fn drop(&mut self) { if self.active { - self.failure(&LlmCallErrorV2::Cancelled { - message: "typed native LLM stream setup was cancelled".into(), - }); + self.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM stream setup was cancelled".into(), + )); } } } @@ -2292,65 +2295,117 @@ unsafe extern "C" fn native_async_next_invoke_result( NemoRelayStatus::Ok } -const NATIVE_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; -const NATIVE_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; -const NATIVE_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; const NATIVE_API_V2_MAX_FAILURE_BODY_BYTES: usize = 16 * 1024; const NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES: usize = 1024; +const NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES: usize = 4 * 1024; fn prepare_typed_llm_dispatch( - mut dispatch: LlmDispatchRequestV2, -) -> std::result::Result { + dispatch: LlmDispatchRequestV2, +) -> std::result::Result<(LlmRequest, LlmDispatchTargetContext), NemoRelayStatus> { let url = match reqwest::Url::parse(&dispatch.target.url) { - Ok(url) if matches!(url.scheme(), "http" | "https") && url.has_host() => url, + Ok(url) + if matches!(url.scheme(), "http" | "https") + && url.has_host() + && url.username().is_empty() + && url.password().is_none() => + { + url + } _ => { - set_native_last_error("typed LLM dispatch target must be an absolute HTTP(S) URL"); + set_native_last_error( + "typed LLM dispatch target must be an absolute HTTP(S) URL without user info", + ); return Err(NemoRelayStatus::InvalidArg); } }; - dispatch.request.headers.insert( - NATIVE_DISPATCH_URL_HEADER.into(), - Json::String(url.to_string()), - ); - dispatch.request.headers.insert( - NATIVE_DISPATCH_ROUTE_HEADER.into(), - Json::String(dispatch.target.route.as_str().into()), - ); - dispatch.request.headers.insert( - NATIVE_RETRY_AWARE_HEADER.into(), - Json::String("true".into()), + let method = match reqwest::Method::from_bytes(dispatch.target.method.as_bytes()) { + Ok(method) if method != reqwest::Method::CONNECT && method != reqwest::Method::TRACE => { + method + } + _ => { + set_native_last_error("typed LLM dispatch method was invalid or prohibited"); + return Err(NemoRelayStatus::InvalidArg); + } + }; + for (name, value) in &dispatch.target.headers { + let Ok(parsed_name) = reqwest::header::HeaderName::from_bytes(name.as_bytes()) else { + set_native_last_error("typed LLM dispatch contained an invalid target header name"); + return Err(NemoRelayStatus::InvalidArg); + }; + if prohibited_target_header(&parsed_name) { + set_native_last_error(format!( + "typed LLM dispatch target header {parsed_name} is host-owned or prohibited" + )); + return Err(NemoRelayStatus::InvalidArg); + } + if reqwest::header::HeaderValue::from_str(value).is_err() { + set_native_last_error(format!( + "typed LLM dispatch target header {parsed_name} had an invalid value" + )); + return Err(NemoRelayStatus::InvalidArg); + } + } + let target = LlmDispatchTargetContext::new( + method.as_str().to_owned(), + url.to_string(), + dispatch.target.route.as_str().into(), + dispatch.target.headers, ); - Ok(dispatch.request) + Ok((dispatch.request, target)) +} + +fn prohibited_target_header(name: &reqwest::header::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" + ) } -fn typed_llm_failure(error: FlowError) -> LlmCallErrorV2 { +fn non_http_llm_failure(kind: LlmNonHttpFailureKindV2, message: String) -> LlmCallFailureV2 { + LlmCallFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind, + message: bounded_utf8(message, NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES), + }, + } +} + +fn typed_llm_failure(error: FlowError) -> LlmCallFailureV2 { match error { - FlowError::Upstream(failure) => { - let class = match failure.class { - UpstreamFailureClass::Connection => LlmUpstreamFailureClassV2::Connection, - UpstreamFailureClass::Timeout => LlmUpstreamFailureClassV2::Timeout, - UpstreamFailureClass::RetryableStatus => LlmUpstreamFailureClassV2::RetryableStatus, - UpstreamFailureClass::ContextWindow => LlmUpstreamFailureClassV2::ContextWindow, - UpstreamFailureClass::ModelUnavailable => { - LlmUpstreamFailureClassV2::ModelUnavailable - } - UpstreamFailureClass::Authentication => LlmUpstreamFailureClassV2::Authentication, - UpstreamFailureClass::InvalidRequest => LlmUpstreamFailureClassV2::InvalidRequest, - UpstreamFailureClass::Other => LlmUpstreamFailureClassV2::Other, - }; - LlmCallErrorV2::Upstream { - class, - retryable: failure.is_retryable(), - status: failure.status, - body: bounded_utf8(failure.body, NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), - headers: safe_native_api_v2_failure_headers(failure.headers), + FlowError::Upstream(failure) => match failure.status { + Some(status) => LlmCallFailureV2::Http { + failure: LlmHttpFailureV2 { + status, + body: bounded_utf8(failure.body, NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), + headers: safe_native_api_v2_failure_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) => LlmCallErrorV2::GuardrailRejected { message }, - FlowError::InvalidArgument(message) => LlmCallErrorV2::InvalidRequest { message }, - other => LlmCallErrorV2::Internal { - message: other.to_string(), }, + 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()), } } @@ -2434,8 +2489,8 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( Ok(dispatch) => dispatch, Err(status) => return status, }; - let request = match prepare_typed_llm_dispatch(dispatch) { - Ok(request) => request, + let (request, target) = match prepare_typed_llm_dispatch(dispatch) { + Ok(prepared) => prepared, Err(status) => return status, }; let continuation_context = match next.context.isolated_for_current_invocation() { @@ -2446,15 +2501,17 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( 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(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()) - ))) - }); + 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()) + ))) + }); let outcome = match result { Ok(response) => LlmCallOutcomeV2::Success { response }, Err(error) => LlmCallOutcomeV2::Failure { @@ -2659,8 +2716,8 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( Ok(dispatch) => dispatch, Err(status) => return status, }; - let request = match prepare_typed_llm_dispatch(dispatch) { - Ok(request) => request, + let (request, target) = match prepare_typed_llm_dispatch(dispatch) { + Ok(prepared) => prepared, Err(status) => return status, }; let continuation_context = match next.context.isolated_for_current_invocation() { @@ -2703,7 +2760,7 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( return; } continuation_context - .run(async move { + .invoke_with_llm_dispatch_target(target, move || async move { let mut callback_guard = callback_guard; let result = AssertUnwindSafe(async { match next_fn(request).await { @@ -2762,11 +2819,11 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( || output_stream_for_task.settled.load(Ordering::Acquire) { producer_abort.abort(); - callback_guard.failure(&LlmCallErrorV2::Cancelled { - message: - "typed native LLM stream output settled during setup" - .into(), - }); + callback_guard.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM stream output settled during setup" + .into(), + )); return; } output_stream_for_task @@ -2785,12 +2842,13 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( .catch_unwind() .await; if let Err(payload) = result { - callback_guard.failure(&LlmCallErrorV2::Internal { - message: format!( + callback_guard.failure(&non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!( "typed native LLM stream continuation panicked: {}", panic_payload_message(payload.as_ref()) ), - }); + )); } }) .await; @@ -2836,9 +2894,10 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( .recv() .await .unwrap_or_else(|| LlmStreamEventV2::Failure { - error: LlmCallErrorV2::Cancelled { - message: "native API v2 provider stream closed without a terminal event".into(), - }, + error: non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "native API v2 provider stream closed without a terminal event".into(), + ), }); // The pull operation is complete before callback delivery. Clearing // the guard first lets a callback wake plugin code that immediately diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index d31846949..58710a028 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -28,8 +28,8 @@ 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; @@ -136,7 +136,9 @@ unsafe extern "C" fn record_typed_llm_stream_open( error_json: *const NemoRelayNativeString, ) { let sender = unsafe { - Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender>) + Box::from_raw( + user_data as *mut tokio::sync::oneshot::Sender>, + ) }; let result = if error_json.is_null() { Ok(stream as usize) @@ -145,13 +147,28 @@ unsafe extern "C" fn record_typed_llm_stream_open( .and_then(|value| { serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) }) - .map_err(|status| LlmCallErrorV2::Internal { - message: format!("invalid typed stream open error: {status:?}"), + .map_err(|status| { + non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!("invalid typed stream open error: {status:?}"), + ) }) }; let _ = sender.send(result); } +fn test_dispatch_target( + url: &str, + route: nemo_relay_plugin::LlmDispatchRouteV2, +) -> nemo_relay_plugin::LlmDispatchTargetV2 { + nemo_relay_plugin::LlmDispatchTargetV2 { + method: "POST".into(), + url: url.into(), + route, + headers: BTreeMap::new(), + } +} + unsafe extern "C" fn record_typed_llm_stream_next( user_data: *mut c_void, event_json: *const NemoRelayNativeString, @@ -1531,27 +1548,15 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { .unwrap(); 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.method(), "POST"); + assert_eq!(target.url(), "https://provider.example/v1/chat/completions"); + assert_eq!(target.route(), "openai_chat"); assert_eq!( - request - .headers - .get(NATIVE_DISPATCH_URL_HEADER) - .and_then(Json::as_str), - Some("https://provider.example/v1/chat/completions") - ); - assert_eq!( - request - .headers - .get(NATIVE_DISPATCH_ROUTE_HEADER) - .and_then(Json::as_str), - Some("openai_chat") - ); - assert_eq!( - request - .headers - .get(NATIVE_RETRY_AWARE_HEADER) - .and_then(Json::as_str), - Some("true") + target.headers().get("authorization").map(String::as_str), + Some("Bearer target-secret") ); Err(FlowError::Upstream(crate::error::UpstreamFailure { status: Some(429), @@ -1572,8 +1577,11 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { content: json!({"model": "provider/model"}), }, target: nemo_relay_plugin::LlmDispatchTargetV2 { - url: "https://provider.example/v1/chat/completions".into(), - route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + headers: BTreeMap::from([("authorization".into(), "Bearer target-secret".into())]), + ..test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ) }, }) .unwrap(), @@ -1595,12 +1603,12 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { assert_eq!( outcome, LlmCallOutcomeV2::Failure { - error: LlmCallErrorV2::Upstream { - class: LlmUpstreamFailureClassV2::RetryableStatus, - retryable: true, - status: Some(429), - body: "rate limited".into(), - headers: BTreeMap::from([("retry-after".into(), "1".into())]), + error: LlmCallFailureV2::Http { + failure: LlmHttpFailureV2 { + status: 429, + body: "rate limited".into(), + headers: BTreeMap::from([("retry-after".into(), "1".into())]), + }, }, } ); @@ -1636,10 +1644,10 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { headers: Map::new(), content: json!({}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { - url: "/v1/chat/completions".into(), - route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, - }, + target: test_dispatch_target( + "/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ), }) .unwrap(), ) @@ -1664,6 +1672,74 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { } } +#[test] +fn native_api_v2_rejects_prohibited_target_methods_and_headers() { + let request = LlmRequest { + headers: Map::new(), + content: json!({}), + }; + let mut target = test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ); + target.method = "CONNECT".into(); + assert_eq!( + prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + request: request.clone(), + target, + }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); + + let mut target = test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ); + target + .headers + .insert("x-nemo-relay-internal-dispatch-url".into(), "secret".into()); + assert_eq!( + prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + request: request.clone(), + target, + }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); + + let mut target = test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ); + target + .headers + .insert("host".into(), "attacker.invalid".into()); + assert_eq!( + prepare_typed_llm_dispatch(LlmDispatchRequestV2 { request, target }).unwrap_err(), + NemoRelayStatus::InvalidArg + ); +} + +#[test] +fn native_api_v2_rejects_target_url_credentials() { + let target = test_dispatch_target( + "https://user:secret@provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ); + assert_eq!( + prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({}), + }, + target, + }) + .unwrap_err(), + NemoRelayStatus::InvalidArg + ); +} + #[test] fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { let error = typed_llm_failure(FlowError::Upstream(crate::error::UpstreamFailure { @@ -1681,9 +1757,10 @@ fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { class: UpstreamFailureClass::RetryableStatus, })); - let LlmCallErrorV2::Upstream { body, headers, .. } = error else { + let LlmCallFailureV2::Http { failure } = error else { panic!("expected an upstream failure"); }; + let LlmHttpFailureV2 { body, headers, .. } = failure; assert_eq!(body.len(), NATIVE_API_V2_MAX_FAILURE_BODY_BYTES); assert!(body.is_char_boundary(body.len())); assert_eq!(headers.get("retry-after").map(String::as_str), Some("1")); @@ -1764,14 +1841,14 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { .unwrap(); let next = Arc::new(NativeAsyncNext::new( NativeAsyncNextInner::LlmStream(Arc::new(|request| { - assert_eq!( - request - .headers - .get(NATIVE_DISPATCH_ROUTE_HEADER) - .and_then(Json::as_str), - Some("anthropic_messages") - ); + assert!(request.headers.is_empty()); Box::pin(async move { + assert_eq!( + current_llm_dispatch_target() + .expect("typed target is bound") + .route(), + "anthropic_messages" + ); Ok(LlmJsonStream::new(tokio_stream::iter(vec![ Ok(json!({"type": "content_block_delta", "delta": {"text": "hi"}})), Err(FlowError::Upstream(crate::error::UpstreamFailure { @@ -1805,17 +1882,17 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { headers: Map::new(), content: json!({"model": "provider/model", "stream": true}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { - url: "https://provider.example/v1/messages".into(), - route: nemo_relay_plugin::LlmDispatchRouteV2::AnthropicMessages, - }, + target: test_dispatch_target( + "https://provider.example/v1/messages", + nemo_relay_plugin::LlmDispatchRouteV2::AnthropicMessages, + ), }) .unwrap(), ) .unwrap(); let (open_sender, open_receiver) = - tokio::sync::oneshot::channel::>(); + tokio::sync::oneshot::channel::>(); assert_eq!( unsafe { native_async_llm_next_open_stream_v2( @@ -1873,12 +1950,12 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { }), }, LlmStreamEventV2::Failure { - error: LlmCallErrorV2::Upstream { - class: LlmUpstreamFailureClassV2::ModelUnavailable, - retryable: true, - status: Some(503), - body: "unavailable".into(), - headers: BTreeMap::new(), + error: LlmCallFailureV2::Http { + failure: LlmHttpFailureV2 { + status: 503, + body: "unavailable".into(), + headers: BTreeMap::new(), + }, }, }, ] @@ -1961,7 +2038,12 @@ fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight assert!(matches!( event, LlmStreamEventV2::Failure { - error: LlmCallErrorV2::Cancelled { .. } + error: LlmCallFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + } } )); } @@ -1996,10 +2078,10 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { headers: Map::new(), content: json!({"model": "provider/model"}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { - url: "https://provider.example/v1/chat/completions".into(), - route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, - }, + target: test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ), }) .unwrap(), ) @@ -2045,6 +2127,89 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { } } +#[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"); + Ok(json!({ + "request": request.content, + "url": target.url(), + "authorization": target.headers().get("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(LlmDispatchRequestV2 { + request: LlmRequest { + headers: Map::new(), + content: json!({"index": index}), + }, + target: nemo_relay_plugin::LlmDispatchTargetV2 { + headers: BTreeMap::from([( + "authorization".into(), + format!("Bearer target-{index}"), + )]), + ..test_dispatch_target( + &format!("https://provider-{index}.example/v1/chat/completions"), + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ) + }, + }) + .unwrap(), + ) + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + 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"), + LlmCallOutcomeV2::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_handles_64_concurrent_100_event_provider_streams() { const STREAM_COUNT: usize = 64; @@ -2087,10 +2252,10 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { headers: Map::new(), content: json!({"model": "provider/model", "stream": true}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { - url: "https://provider.example/v1/chat/completions".into(), - route: nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, - }, + target: test_dispatch_target( + "https://provider.example/v1/chat/completions", + nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + ), }) .unwrap(), ) @@ -2099,7 +2264,7 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { let mut open_receivers = Vec::with_capacity(STREAM_COUNT); for _ in 0..STREAM_COUNT { let (sender, receiver) = - tokio::sync::oneshot::channel::>(); + tokio::sync::oneshot::channel::>(); assert_eq!( unsafe { native_async_llm_next_open_stream_v2( diff --git a/crates/plugin/README.md b/crates/plugin/README.md index e05f8ce2e..887e05d88 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -114,14 +114,23 @@ Set `compat.native_api = "2"` in `relay-plugin.toml`. During registration, `PluginContext::host_api_v4` exposes the C-safe typed LLM dispatch table and the raw v2 registration helpers. -The plugin provides JSON plus an absolute target URL and protocol route. Relay -returns response JSON or a structured provider failure. Streaming dispatch -returns an opaque host-owned stream; request one JSON event at a time, then -cancel and release it exactly once. No Rust future, trait object, -`serde_json::Value`, or allocator-owned Rust string crosses the ABI boundary. +The plugin provides JSON plus an HTTP method, absolute target URL, protocol +route, and explicit target headers. Relay binds that transport target to the +current LLM continuation without storing it in `LlmRequest.headers`. Successful +calls return provider JSON. Provider rejections return an HTTP status, bounded +body, and safe response headers; failures without an HTTP response use a small +transport-oriented kind. + +Streaming dispatch returns an opaque host-owned stream; request one JSON event +at a time, then cancel and release it exactly once. No Rust future, trait +object, `serde_json::Value`, or allocator-owned Rust string crosses the ABI +boundary. 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. ## Documentation diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 04d92921a..c6972ef1a 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -938,10 +938,17 @@ impl LlmDispatchRouteV2 { /// Explicit provider target supplied to Relay through native API v2. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] pub struct LlmDispatchTargetV2 { + /// HTTP method used for the provider call. + pub method: String, /// Absolute HTTP(S) provider URL including the selected endpoint. pub url: String, /// Provider protocol used by the selected endpoint. pub route: LlmDispatchRouteV2, + /// 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, } /// Typed LLM continuation invocation supplied through native API v2. @@ -953,80 +960,82 @@ pub struct LlmDispatchRequestV2 { pub target: LlmDispatchTargetV2, } -/// Stable provider-failure classification exposed to native plugins. +/// Bounded HTTP failure returned by a provider. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +pub struct LlmHttpFailureV2 { + /// Provider HTTP status. + pub status: u16, + /// Bounded provider response body. + pub body: String, + /// Safe response headers with credential-bearing fields removed. + pub headers: BTreeMap, +} + +/// 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 LlmUpstreamFailureClassV2 { +pub enum LlmNonHttpFailureKindV2 { /// Provider connection could not be established or was interrupted. - Connection, + Transport, /// Provider request timed out. Timeout, - /// Retryable HTTP status without a more specific provider classification. - RetryableStatus, - /// Provider rejected the request because its context window was exceeded. - ContextWindow, - /// Requested provider model is temporarily unavailable. - ModelUnavailable, - /// Provider authentication or authorization failed. - Authentication, - /// Provider rejected an invalid request. + /// The caller cancelled the operation. + Cancelled, + /// Relay rejected an invalid dispatch request. InvalidRequest, - /// Other non-retryable provider failure. - Other, + /// A guardrail rejected the provider call. + Guardrail, + /// Relay could not complete the operation. + Internal, +} + +/// Bounded failure for an operation that produced no provider HTTP response. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +pub struct LlmNonHttpFailureV2 { + /// Stable failure kind. + pub kind: LlmNonHttpFailureKindV2, + /// Bounded human-readable context. + pub message: String, } /// Structured LLM continuation failure exposed through native API v2. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmCallErrorV2 { - /// Relay received a provider or provider-transport failure. - Upstream { - /// Stable failure classification. - class: LlmUpstreamFailureClassV2, - /// Whether Relay considers the failure safe to retry. - retryable: bool, - /// Provider HTTP status when a response was received. - status: Option, - /// Bounded provider response body or transport error. - body: String, - /// Safe response headers with credential-bearing fields removed. - headers: BTreeMap, +pub enum LlmCallFailureV2 { + /// A provider returned a non-success HTTP response. + Http { + /// Bounded HTTP failure details. + failure: LlmHttpFailureV2, }, - /// A guardrail rejected the provider call. - GuardrailRejected { - /// Human-readable rejection reason. - message: String, - }, - /// Relay rejected an invalid dispatch request. - InvalidRequest { - /// Human-readable validation failure. - message: String, - }, - /// The caller cancelled the operation. - Cancelled { - /// Human-readable cancellation context. - message: String, - }, - /// Relay could not complete the operation because of an internal failure. - Internal { - /// Human-readable internal failure. - message: String, + /// No provider HTTP response was available. + NonHttp { + /// Bounded non-HTTP failure details. + failure: LlmNonHttpFailureV2, }, } -impl LlmCallErrorV2 { - /// Returns whether Relay classified this failure as retryable. +impl LlmCallFailureV2 { + /// Return Relay's provider-neutral retry disposition. + /// + /// The disposition is derived rather than serialized so the wire contract + /// contains only HTTP semantics and the minimal non-HTTP failure kind. pub const fn is_retryable(&self) -> bool { - matches!( - self, - Self::Upstream { - retryable: true, - .. - } - ) + match self { + Self::Http { failure } => is_retryable_http_status_v2(failure.status), + Self::NonHttp { failure } => matches!( + failure.kind, + LlmNonHttpFailureKindV2::Transport | LlmNonHttpFailureKindV2::Timeout + ), + } } } +/// Return Relay's provider-neutral retry disposition for an HTTP status. +#[must_use] +pub const fn is_retryable_http_status_v2(status: u16) -> bool { + matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504) +} + /// Unary LLM continuation outcome delivered through native API v2. #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] @@ -1039,7 +1048,7 @@ pub enum LlmCallOutcomeV2 { /// Provider call failed before producing a response. Failure { /// Structured Relay/provider failure. - error: LlmCallErrorV2, + error: LlmCallFailureV2, }, } @@ -1057,7 +1066,7 @@ pub enum LlmStreamEventV2 { /// Provider stream failed before clean completion. Failure { /// Structured Relay/provider failure. - error: LlmCallErrorV2, + error: LlmCallFailureV2, }, } diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 3ca498710..8e4c9e361 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -14,7 +14,8 @@ use std::sync::{ use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, - Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmJsonStream, LlmNext, + Event, EventCategory, EventSanitizeFields, Json, LlmCallFailureV2, LlmCodecIdentity, + LlmHttpFailureV2, LlmJsonStream, LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareKind, @@ -34,6 +35,68 @@ use nemo_relay_plugin::{ }; use serde_json::{Map, json}; +#[test] +fn native_api_v2_retry_policy_is_derived_from_http_semantics() { + for status in [408, 425, 429, 500, 502, 503, 504] { + assert!( + LlmCallFailureV2::Http { + failure: LlmHttpFailureV2 { + status, + body: String::new(), + headers: Default::default(), + }, + } + .is_retryable(), + "status={status}" + ); + } + for status in [400, 401, 404, 409, 422, 501] { + assert!( + !LlmCallFailureV2::Http { + failure: LlmHttpFailureV2 { + status, + body: String::new(), + headers: Default::default(), + }, + } + .is_retryable(), + "status={status}" + ); + } + for kind in [ + LlmNonHttpFailureKindV2::Transport, + LlmNonHttpFailureKindV2::Timeout, + ] { + assert!( + LlmCallFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind, + message: String::new(), + }, + } + .is_retryable(), + "kind={kind:?}" + ); + } + for kind in [ + LlmNonHttpFailureKindV2::Cancelled, + LlmNonHttpFailureKindV2::InvalidRequest, + LlmNonHttpFailureKindV2::Guardrail, + LlmNonHttpFailureKindV2::Internal, + ] { + assert!( + !LlmCallFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind, + message: String::new(), + }, + } + .is_retryable(), + "kind={kind:?}" + ); + } +} + #[test] fn async_abi_discriminants_reject_unknown_values() { use NemoRelayNativeAsyncMiddlewareKind as Kind; diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index af24486e4..652687ab1 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -124,7 +124,7 @@ 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 targets, failures, and provider streams | +| `"2"` | `nemo_relay_plugin_v2!` | `NemoRelayNativeHostApiV4` | Host-dispatched LLM calls with typed HTTP targets and provider streams | Relay retains native API v1 unchanged. A v2-only plugin is rejected clearly by a v1 host and is not retried against an older table. Validation, inspection, @@ -132,21 +132,36 @@ 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`, absolute HTTP(S) target URL, and target route -(`openai_chat`, `openai_responses`, or `anthropic_messages`). Relay returns -either: +replacement `LlmRequest` and a target containing the HTTP method, absolute +HTTP(S) URL, provider route (`openai_chat`, `openai_responses`, or +`anthropic_messages`), and explicit outbound headers. Target headers may +contain provider credentials. Relay validates and transports them but never +places their values in diagnostics or observability. + +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; -- a structured provider failure with stable class, retryability, optional HTTP - status, a body bounded to 16 KiB, and a conservative safe-header allowlist; or +- 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. -The host converts the public target into its private dispatch metadata. Plugins -do not construct private Relay headers. Provider stream production and plugin -output each use bounded 32-event queues. Dropping or cancelling a stream stops -provider production, and the library stays loaded until all callbacks and -streams release their handles. +HTTP retryability is derived by the SDK from status alone: `408`, `425`, `429`, +`500`, `502`, `503`, and `504` are retryable. Transport and timeout failures +are retryable; other non-HTTP failures are not. Relay does not inspect provider +bodies to infer context-window or model-availability errors, so ordinary `400` +and `404` responses do not trigger automatic reselection. + +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 stream production and plugin output +each use bounded 32-event queues. Dropping or cancelling a stream stops provider +production, and the library stays loaded until all callbacks and streams +release their handles. V2 callbacks run on Relay's blocking executor with the active scope stack restored. Continuations run on Relay's existing Tokio runtime; Relay does not From de3fd45b4470679afa01eda34a4aebdaebf81d71 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 31 Jul 2026 10:56:43 -0600 Subject: [PATCH 05/32] refactor(plugin): clarify LLM continuation types Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 114 +++++++++--------- crates/core/tests/unit/native_plugin_tests.rs | 108 +++++++++-------- crates/plugin/README.md | 10 +- crates/plugin/src/lib.rs | 56 ++++----- crates/plugin/tests/typed_callbacks.rs | 16 +-- 5 files changed, 156 insertions(+), 148 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index cf078fbf6..08e56f7d1 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -51,14 +51,14 @@ use crate::plugin::{ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ - LlmCallFailureV2, LlmCallOutcomeV2, LlmDispatchRequestV2, LlmHttpFailureV2, - LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmStreamEventV2, NEMO_RELAY_NATIVE_ABI_VERSION, - NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, - NemoRelayNativeAsyncLlmResultCbV2, NemoRelayNativeAsyncLlmStreamNextCbV2, - NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, - NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, + LlmContinuationStreamEventV2, LlmHttpFailureV2, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, + NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, + NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, @@ -915,7 +915,7 @@ 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_TYPED_LLM_DISPATCH; + v3.v1.abi_version = NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS; v3.v1.struct_size = std::mem::size_of::(); NemoRelayNativeHostApiV4 { v3, @@ -1535,7 +1535,7 @@ struct NativeAsyncStreamCallbackGuard { } struct NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex>, + receiver: tokio::sync::Mutex>, producer_abort: Mutex>, runtime: tokio::runtime::Handle, next_in_flight: AtomicBool, @@ -1574,7 +1574,7 @@ impl NativeLlmStreamOpenCallbackGuardV2 { self.active = false; } - fn failure(&mut self, error: &LlmCallFailureV2) { + fn failure(&mut self, error: &LlmContinuationFailureV2) { if !self.active { return; } @@ -2299,10 +2299,10 @@ const NATIVE_API_V2_MAX_FAILURE_BODY_BYTES: usize = 16 * 1024; const NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES: usize = 1024; const NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES: usize = 4 * 1024; -fn prepare_typed_llm_dispatch( - dispatch: LlmDispatchRequestV2, +fn prepare_llm_continuation_invocation( + invocation: LlmContinuationInvocationV2, ) -> std::result::Result<(LlmRequest, LlmDispatchTargetContext), NemoRelayStatus> { - let url = match reqwest::Url::parse(&dispatch.target.url) { + let url = match reqwest::Url::parse(&invocation.target.url) { Ok(url) if matches!(url.scheme(), "http" | "https") && url.has_host() @@ -2313,34 +2313,34 @@ fn prepare_typed_llm_dispatch( } _ => { set_native_last_error( - "typed LLM dispatch target must be an absolute HTTP(S) URL without user info", + "LLM continuation target must be an absolute HTTP(S) URL without user info", ); return Err(NemoRelayStatus::InvalidArg); } }; - let method = match reqwest::Method::from_bytes(dispatch.target.method.as_bytes()) { + let method = match reqwest::Method::from_bytes(invocation.target.method.as_bytes()) { Ok(method) if method != reqwest::Method::CONNECT && method != reqwest::Method::TRACE => { method } _ => { - set_native_last_error("typed LLM dispatch method was invalid or prohibited"); + set_native_last_error("LLM continuation method was invalid or prohibited"); return Err(NemoRelayStatus::InvalidArg); } }; - for (name, value) in &dispatch.target.headers { + for (name, value) in &invocation.target.headers { let Ok(parsed_name) = reqwest::header::HeaderName::from_bytes(name.as_bytes()) else { - set_native_last_error("typed LLM dispatch contained an invalid target header name"); + set_native_last_error("LLM continuation contained an invalid target header name"); return Err(NemoRelayStatus::InvalidArg); }; if prohibited_target_header(&parsed_name) { set_native_last_error(format!( - "typed LLM dispatch target header {parsed_name} is host-owned or prohibited" + "LLM continuation target header {parsed_name} is host-owned or prohibited" )); return Err(NemoRelayStatus::InvalidArg); } if reqwest::header::HeaderValue::from_str(value).is_err() { set_native_last_error(format!( - "typed LLM dispatch target header {parsed_name} had an invalid value" + "LLM continuation target header {parsed_name} had an invalid value" )); return Err(NemoRelayStatus::InvalidArg); } @@ -2348,10 +2348,10 @@ fn prepare_typed_llm_dispatch( let target = LlmDispatchTargetContext::new( method.as_str().to_owned(), url.to_string(), - dispatch.target.route.as_str().into(), - dispatch.target.headers, + invocation.target.route.as_str().into(), + invocation.target.headers, ); - Ok((dispatch.request, target)) + Ok((invocation.request, target)) } fn prohibited_target_header(name: &reqwest::header::HeaderName) -> bool { @@ -2371,8 +2371,11 @@ fn prohibited_target_header(name: &reqwest::header::HeaderName) -> bool { ) } -fn non_http_llm_failure(kind: LlmNonHttpFailureKindV2, message: String) -> LlmCallFailureV2 { - LlmCallFailureV2::NonHttp { +fn non_http_llm_failure( + kind: LlmNonHttpFailureKindV2, + message: String, +) -> LlmContinuationFailureV2 { + LlmContinuationFailureV2::NonHttp { failure: LlmNonHttpFailureV2 { kind, message: bounded_utf8(message, NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES), @@ -2380,10 +2383,10 @@ fn non_http_llm_failure(kind: LlmNonHttpFailureKindV2, message: String) -> LlmCa } } -fn typed_llm_failure(error: FlowError) -> LlmCallFailureV2 { +fn typed_llm_failure(error: FlowError) -> LlmContinuationFailureV2 { match error { FlowError::Upstream(failure) => match failure.status { - Some(status) => LlmCallFailureV2::Http { + Some(status) => LlmContinuationFailureV2::Http { failure: LlmHttpFailureV2 { status, body: bounded_utf8(failure.body, NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), @@ -2453,7 +2456,7 @@ fn bounded_utf8(value: String, max_bytes: usize) -> String { unsafe fn invoke_typed_llm_result_callback( cb: NemoRelayNativeAsyncLlmResultCbV2, user_data: *mut c_void, - outcome: &LlmCallOutcomeV2, + outcome: &LlmContinuationOutcomeV2, ) { if let Some(outcome) = native_string_from_json( &serde_json::to_value(outcome) @@ -2469,7 +2472,7 @@ unsafe fn invoke_typed_llm_result_callback( /// Invokes a unary LLM continuation through native API v2. unsafe extern "C" fn native_async_llm_next_invoke_result_v2( next: *const NemoRelayNativeAsyncNext, - dispatch_json: *const NemoRelayNativeString, + invocation_json: *const NemoRelayNativeString, cb: NemoRelayNativeAsyncLlmResultCbV2, user_data: *mut c_void, ) -> NemoRelayStatus { @@ -2477,19 +2480,20 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( return NemoRelayStatus::NullPointer; }; let NativeAsyncNextInner::Llm(next_fn) = &next.inner else { - set_native_last_error("typed unary LLM dispatch requires an LLM execution continuation"); + set_native_last_error("targeted LLM continuation requires an LLM execution continuation"); return NemoRelayStatus::InvalidArg; }; - let dispatch = match parse_json_arg(dispatch_json, "typed LLM dispatch").and_then(|value| { - serde_json::from_value(value).map_err(|error| { - set_native_last_error(error.to_string()); - NemoRelayStatus::InvalidJson - }) - }) { - Ok(dispatch) => dispatch, - Err(status) => return status, - }; - let (request, target) = match prepare_typed_llm_dispatch(dispatch) { + 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, }; @@ -2513,8 +2517,8 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( ))) }); let outcome = match result { - Ok(response) => LlmCallOutcomeV2::Success { response }, - Err(error) => LlmCallOutcomeV2::Failure { + Ok(response) => LlmContinuationOutcomeV2::Success { response }, + Err(error) => LlmContinuationOutcomeV2::Failure { error: typed_llm_failure(error), }, }; @@ -2687,7 +2691,7 @@ async fn forward_native_async_next_stream_with( /// Opens a streaming LLM continuation through native API v2. unsafe extern "C" fn native_async_llm_next_open_stream_v2( next: *const NemoRelayNativeAsyncNext, - dispatch_json: *const NemoRelayNativeString, + invocation_json: *const NemoRelayNativeString, output_stream: *const NemoRelayNativeAsyncStream, cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, user_data: *mut c_void, @@ -2702,21 +2706,21 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( let output_stream = unsafe { Arc::from_raw(output_stream as *const NativeAsyncStream) }; let NativeAsyncNextInner::LlmStream(next_fn) = &next.inner else { set_native_last_error( - "typed streaming LLM dispatch requires an LLM stream execution continuation", + "targeted LLM continuation requires an LLM stream execution continuation", ); return NemoRelayStatus::InvalidArg; }; - let dispatch = - match parse_json_arg(dispatch_json, "typed streaming LLM dispatch").and_then(|value| { + 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(dispatch) => dispatch, - Err(status) => return status, - }; - let (request, target) = match prepare_typed_llm_dispatch(dispatch) { + Ok(invocation) => invocation, + Err(status) => return status, + }; + let (request, target) = match prepare_llm_continuation_invocation(invocation) { Ok(prepared) => prepared, Err(status) => return status, }; @@ -2781,8 +2785,8 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( while let Some(item) = provider_stream.next().await { let terminal = item.is_err(); let event = match item { - Ok(chunk) => LlmStreamEventV2::Chunk { chunk }, - Err(error) => LlmStreamEventV2::Failure { + Ok(chunk) => LlmContinuationStreamEventV2::Chunk { chunk }, + Err(error) => LlmContinuationStreamEventV2::Failure { error: typed_llm_failure(error), }, }; @@ -2795,7 +2799,7 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( && !output_for_producer.cancelled.load(Ordering::Acquire) && !output_for_producer.settled.load(Ordering::Acquire) { - let _ = sender.send(LlmStreamEventV2::Done).await; + let _ = sender.send(LlmContinuationStreamEventV2::Done).await; } output_for_producer .downstream_aborts @@ -2893,7 +2897,7 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( .await .recv() .await - .unwrap_or_else(|| LlmStreamEventV2::Failure { + .unwrap_or_else(|| LlmContinuationStreamEventV2::Failure { error: non_http_llm_failure( LlmNonHttpFailureKindV2::Cancelled, "native API v2 provider stream closed without a terminal event".into(), diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 58710a028..cf12e6593 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -115,8 +115,9 @@ unsafe extern "C" fn complete_typed_llm_result( user_data: *mut c_void, outcome_json: *const NemoRelayNativeString, ) { - let sender = - unsafe { Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) }; + let sender = unsafe { + Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) + }; let outcome = parse_json_arg(outcome_json, "typed LLM result") .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) .expect("host emitted a valid typed LLM outcome"); @@ -137,7 +138,7 @@ unsafe extern "C" fn record_typed_llm_stream_open( ) { let sender = unsafe { Box::from_raw( - user_data as *mut tokio::sync::oneshot::Sender>, + user_data as *mut tokio::sync::oneshot::Sender>, ) }; let result = if error_json.is_null() { @@ -159,9 +160,9 @@ unsafe extern "C" fn record_typed_llm_stream_open( fn test_dispatch_target( url: &str, - route: nemo_relay_plugin::LlmDispatchRouteV2, -) -> nemo_relay_plugin::LlmDispatchTargetV2 { - nemo_relay_plugin::LlmDispatchTargetV2 { + route: nemo_relay_plugin::LlmContinuationRouteV2, +) -> nemo_relay_plugin::LlmContinuationTargetV2 { + nemo_relay_plugin::LlmContinuationTargetV2 { method: "POST".into(), url: url.into(), route, @@ -173,9 +174,10 @@ unsafe extern "C" fn record_typed_llm_stream_next( user_data: *mut c_void, event_json: *const NemoRelayNativeString, ) { - let sender = - unsafe { Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) }; - let event: LlmStreamEventV2 = parse_json_arg(event_json, "typed LLM stream result") + let sender = unsafe { + Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) + }; + let event: LlmContinuationStreamEventV2 = parse_json_arg(event_json, "typed LLM stream result") .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) .expect("host emitted a valid typed LLM stream event"); let _ = sender.send(event); @@ -1279,7 +1281,7 @@ fn assert_native_json_output_and_host_api() { let host_api = unsafe { &*native_host_api() }; assert_eq!( host_api.abi_version, - NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS ); assert_eq!( host_api.struct_size, @@ -1571,23 +1573,23 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &serde_json::to_value(LlmContinuationInvocationV2 { request: LlmRequest { headers: Map::new(), content: json!({"model": "provider/model"}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { + target: nemo_relay_plugin::LlmContinuationTargetV2 { headers: BTreeMap::from([("authorization".into(), "Bearer target-secret".into())]), ..test_dispatch_target( "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ) }, }) .unwrap(), ) .unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); assert_eq!( unsafe { native_async_llm_next_invoke_result_v2( @@ -1602,8 +1604,8 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { let outcome = runtime.block_on(receiver).unwrap(); assert_eq!( outcome, - LlmCallOutcomeV2::Failure { - error: LlmCallFailureV2::Http { + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::Http { failure: LlmHttpFailureV2 { status: 429, body: "rate limited".into(), @@ -1639,14 +1641,14 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &serde_json::to_value(LlmContinuationInvocationV2 { request: LlmRequest { headers: Map::new(), content: json!({}), }, target: test_dispatch_target( "/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ), }) .unwrap(), @@ -1680,11 +1682,11 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { }; let mut target = test_dispatch_target( "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ); target.method = "CONNECT".into(); assert_eq!( - prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request: request.clone(), target, }) @@ -1694,13 +1696,13 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { let mut target = test_dispatch_target( "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ); target .headers .insert("x-nemo-relay-internal-dispatch-url".into(), "secret".into()); assert_eq!( - prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request: request.clone(), target, }) @@ -1710,13 +1712,14 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { let mut target = test_dispatch_target( "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ); target .headers .insert("host".into(), "attacker.invalid".into()); assert_eq!( - prepare_typed_llm_dispatch(LlmDispatchRequestV2 { request, target }).unwrap_err(), + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request, target }) + .unwrap_err(), NemoRelayStatus::InvalidArg ); } @@ -1725,10 +1728,10 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { fn native_api_v2_rejects_target_url_credentials() { let target = test_dispatch_target( "https://user:secret@provider.example/v1/chat/completions", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ); assert_eq!( - prepare_typed_llm_dispatch(LlmDispatchRequestV2 { + prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request: LlmRequest { headers: Map::new(), content: json!({}), @@ -1757,7 +1760,7 @@ fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { class: UpstreamFailureClass::RetryableStatus, })); - let LlmCallFailureV2::Http { failure } = error else { + let LlmContinuationFailureV2::Http { failure } = error else { panic!("expected an upstream failure"); }; let LlmHttpFailureV2 { body, headers, .. } = failure; @@ -1877,14 +1880,14 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { let output_stream_ref = Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &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", - nemo_relay_plugin::LlmDispatchRouteV2::AnthropicMessages, + nemo_relay_plugin::LlmContinuationRouteV2::AnthropicMessages, ), }) .unwrap(), @@ -1892,7 +1895,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { .unwrap(); let (open_sender, open_receiver) = - tokio::sync::oneshot::channel::>(); + tokio::sync::oneshot::channel::>(); assert_eq!( unsafe { native_async_llm_next_open_stream_v2( @@ -1933,7 +1936,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { }); let terminal = matches!( event, - LlmStreamEventV2::Done | LlmStreamEventV2::Failure { .. } + LlmContinuationStreamEventV2::Done | LlmContinuationStreamEventV2::Failure { .. } ); events.push(event); if terminal { @@ -1943,14 +1946,14 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { assert_eq!( events, vec![ - LlmStreamEventV2::Chunk { + LlmContinuationStreamEventV2::Chunk { chunk: json!({ "type": "content_block_delta", "delta": {"text": "hi"}, }), }, - LlmStreamEventV2::Failure { - error: LlmCallFailureV2::Http { + LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::Http { failure: LlmHttpFailureV2 { status: 503, body: "unavailable".into(), @@ -2003,7 +2006,8 @@ fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight NemoRelayStatus::Ok ); - let (overlap_sender, _overlap_receiver) = tokio::sync::oneshot::channel::(); + let (overlap_sender, _overlap_receiver) = + tokio::sync::oneshot::channel::(); let overlap_state = Box::into_raw(Box::new(overlap_sender)); assert_eq!( unsafe { @@ -2037,8 +2041,8 @@ fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight .expect("provider next callback should be delivered"); assert!(matches!( event, - LlmStreamEventV2::Failure { - error: LlmCallFailureV2::NonHttp { + LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { failure: LlmNonHttpFailureV2 { kind: LlmNonHttpFailureKindV2::Cancelled, .. @@ -2073,14 +2077,14 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &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", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ), }) .unwrap(), @@ -2089,7 +2093,7 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { let mut receivers = Vec::with_capacity(DISPATCH_COUNT); for _ in 0..DISPATCH_COUNT { - let (sender, receiver) = tokio::sync::oneshot::channel::(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); assert_eq!( unsafe { native_async_llm_next_invoke_result_v2( @@ -2115,7 +2119,7 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { assert_eq!(outcomes.len(), DISPATCH_COUNT); assert!(outcomes.into_iter().all(|outcome| { outcome - == Ok(LlmCallOutcomeV2::Success { + == Ok(LlmContinuationOutcomeV2::Success { response: json!({"ok": true}), }) })); @@ -2155,26 +2159,26 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { for index in 0..DISPATCH_COUNT { let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &serde_json::to_value(LlmContinuationInvocationV2 { request: LlmRequest { headers: Map::new(), content: json!({"index": index}), }, - target: nemo_relay_plugin::LlmDispatchTargetV2 { + 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"), - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ) }, }) .unwrap(), ) .unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); assert_eq!( unsafe { native_async_llm_next_invoke_result_v2( @@ -2196,7 +2200,7 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { receiver .await .expect("dispatch callback should be delivered"), - LlmCallOutcomeV2::Success { + LlmContinuationOutcomeV2::Success { response: json!({ "request": {"index": index}, "url": format!("https://provider-{index}.example/v1/chat/completions"), @@ -2247,14 +2251,14 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { let output_stream_ref = Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; let dispatch = native_string_from_json( - &serde_json::to_value(LlmDispatchRequestV2 { + &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", - nemo_relay_plugin::LlmDispatchRouteV2::OpenaiChat, + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, ), }) .unwrap(), @@ -2264,7 +2268,7 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { let mut open_receivers = Vec::with_capacity(STREAM_COUNT); for _ in 0..STREAM_COUNT { let (sender, receiver) = - tokio::sync::oneshot::channel::>(); + tokio::sync::oneshot::channel::>(); assert_eq!( unsafe { native_async_llm_next_open_stream_v2( @@ -2310,9 +2314,9 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { NemoRelayStatus::Ok ); match receiver.await.expect("next callback should be delivered") { - LlmStreamEventV2::Chunk { .. } => chunks += 1, - LlmStreamEventV2::Done => break, - LlmStreamEventV2::Failure { error } => { + LlmContinuationStreamEventV2::Chunk { .. } => chunks += 1, + LlmContinuationStreamEventV2::Done => break, + LlmContinuationStreamEventV2::Failure { error } => { panic!("provider stream failed during stress test: {error:?}") } } diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 887e05d88..688e8a194 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -49,9 +49,9 @@ the dynamic-library boundary on the stable C-compatible ABI. - **Raw async middleware**: Completion-based raw registrations for plugins that need asynchronous guardrails, intercepts, or event sanitizers. Typed Rust callbacks remain synchronous convenience APIs. -- **Native API v2 LLM dispatch**: Register v2-only execution callbacks that - send an explicit provider target through Relay and receive buffered JSON, - structured failures, or a bounded host-owned provider stream. +- **Native API v2 targeted LLM continuations**: Register v2-only execution + callbacks that send an explicit provider target through Relay and receive + buffered JSON, structured failures, or a bounded host-owned provider stream. ## Installation @@ -111,8 +111,8 @@ nemo_relay_plugin::nemo_relay_plugin_v2!( ``` Set `compat.native_api = "2"` in `relay-plugin.toml`. During registration, -`PluginContext::host_api_v4` exposes the C-safe typed LLM dispatch table and -the raw v2 registration helpers. +`PluginContext::host_api_v4` exposes the C-safe targeted LLM continuation table +and the raw v2 registration helpers. The plugin provides JSON plus an HTTP method, absolute target URL, protocol route, and explicit target headers. Relay binds that transport target to the diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index c6972ef1a..a183d65a6 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -43,8 +43,8 @@ use serde_json::Map; 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 typed LLM target dispatch and outcomes. -pub const NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH: u32 = 4; +/// 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; @@ -874,7 +874,7 @@ pub struct NemoRelayNativeAsyncStream { _marker: PhantomData<(*mut u8, PhantomPinned)>, } -/// Opaque host-owned provider stream returned by native API v2 dispatch. +/// 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 cancels /// or releases the handle exactly once. Relay pumps provider output into a @@ -912,10 +912,10 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( error: *const NemoRelayNativeString, ); -/// Provider protocol selected for one typed LLM dispatch. +/// Provider protocol selected for one targeted LLM continuation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] #[serde(rename_all = "snake_case")] -pub enum LlmDispatchRouteV2 { +pub enum LlmContinuationRouteV2 { /// OpenAI Chat Completions. OpenaiChat, /// OpenAI Responses. @@ -924,7 +924,7 @@ pub enum LlmDispatchRouteV2 { AnthropicMessages, } -impl LlmDispatchRouteV2 { +impl LlmContinuationRouteV2 { /// Returns the stable Relay gateway route identifier. pub const fn as_str(self) -> &'static str { match self { @@ -935,15 +935,15 @@ impl LlmDispatchRouteV2 { } } -/// Explicit provider target supplied to Relay through native API v2. +/// Explicit provider target for one native API v2 LLM continuation. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] -pub struct LlmDispatchTargetV2 { +pub struct LlmContinuationTargetV2 { /// HTTP method used for the provider call. pub method: String, /// Absolute HTTP(S) provider URL including the selected endpoint. pub url: String, /// Provider protocol used by the selected endpoint. - pub route: LlmDispatchRouteV2, + pub route: LlmContinuationRouteV2, /// Explicit outbound provider headers, including target credentials. /// /// Relay validates and transports these headers but never records their @@ -953,11 +953,11 @@ pub struct LlmDispatchTargetV2 { /// Typed LLM continuation invocation supplied through native API v2. #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] -pub struct LlmDispatchRequestV2 { +pub struct LlmContinuationInvocationV2 { /// Replacement request passed to the Relay execution continuation. pub request: LlmRequest, /// Explicit provider target selected by the plugin. - pub target: LlmDispatchTargetV2, + pub target: LlmContinuationTargetV2, } /// Bounded HTTP failure returned by a provider. @@ -981,7 +981,7 @@ pub enum LlmNonHttpFailureKindV2 { Timeout, /// The caller cancelled the operation. Cancelled, - /// Relay rejected an invalid dispatch request. + /// Relay rejected an invalid continuation invocation. InvalidRequest, /// A guardrail rejected the provider call. Guardrail, @@ -1001,7 +1001,7 @@ pub struct LlmNonHttpFailureV2 { /// Structured LLM continuation failure exposed through native API v2. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmCallFailureV2 { +pub enum LlmContinuationFailureV2 { /// A provider returned a non-success HTTP response. Http { /// Bounded HTTP failure details. @@ -1014,7 +1014,7 @@ pub enum LlmCallFailureV2 { }, } -impl LlmCallFailureV2 { +impl LlmContinuationFailureV2 { /// Return Relay's provider-neutral retry disposition. /// /// The disposition is derived rather than serialized so the wire contract @@ -1039,7 +1039,7 @@ pub const fn is_retryable_http_status_v2(status: u16) -> bool { /// Unary LLM continuation outcome delivered through native API v2. #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmCallOutcomeV2 { +pub enum LlmContinuationOutcomeV2 { /// Provider call completed successfully. Success { /// Provider response JSON. @@ -1048,14 +1048,14 @@ pub enum LlmCallOutcomeV2 { /// Provider call failed before producing a response. Failure { /// Structured Relay/provider failure. - error: LlmCallFailureV2, + error: LlmContinuationFailureV2, }, } /// Streaming LLM continuation event delivered through native API v2. #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmStreamEventV2 { +pub enum LlmContinuationStreamEventV2 { /// One provider stream event. Chunk { /// Provider event JSON. @@ -1066,13 +1066,13 @@ pub enum LlmStreamEventV2 { /// Provider stream failed before clean completion. Failure { /// Structured Relay/provider failure. - error: LlmCallFailureV2, + error: LlmContinuationFailureV2, }, } /// Receives one typed unary LLM continuation outcome. /// -/// `outcome_json` contains one serialized [`LlmCallOutcomeV2`] and is borrowed +/// `outcome_json` contains one serialized [`LlmContinuationOutcomeV2`] and is borrowed /// for the callback. pub type NemoRelayNativeAsyncLlmResultCbV2 = unsafe extern "C" fn(user_data: *mut c_void, outcome_json: *const NemoRelayNativeString); @@ -1089,7 +1089,7 @@ pub type NemoRelayNativeAsyncLlmStreamOpenCbV2 = unsafe extern "C" fn( /// Receives one item from a native API v2 provider stream. /// -/// `event_json` contains one serialized [`LlmStreamEventV2`] and is borrowed +/// `event_json` contains one serialized [`LlmContinuationStreamEventV2`] and is 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, event_json: *const NemoRelayNativeString); @@ -1259,7 +1259,7 @@ pub struct NemoRelayNativeHostApiV3 { ) -> NemoRelayStatus, } -/// ABI-v4 host extension implementing native API v2 typed LLM dispatch. +/// ABI-v4 host extension implementing native API v2 targeted LLM continuations. /// /// 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. @@ -1272,7 +1272,7 @@ pub struct NemoRelayNativeHostApiV4 { /// outcome. pub async_llm_next_invoke_result_v2: unsafe extern "C" fn( next: *const NemoRelayNativeAsyncNext, - dispatch_json: *const NemoRelayNativeString, + invocation_json: *const NemoRelayNativeString, cb: NemoRelayNativeAsyncLlmResultCbV2, user_data: *mut c_void, ) -> NemoRelayStatus, @@ -1282,7 +1282,7 @@ pub struct NemoRelayNativeHostApiV4 { /// items are read with `async_llm_stream_next_v2`. pub async_llm_next_open_stream_v2: unsafe extern "C" fn( next: *const NemoRelayNativeAsyncNext, - dispatch_json: *const NemoRelayNativeString, + invocation_json: *const NemoRelayNativeString, output_stream: *const NemoRelayNativeAsyncStream, cb: NemoRelayNativeAsyncLlmStreamOpenCbV2, user_data: *mut c_void, @@ -2061,7 +2061,7 @@ impl<'a> PluginContext<'a> { /// 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_TYPED_LLM_DISPATCH + (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) }) } @@ -3618,7 +3618,7 @@ enum OwnedHostApi { impl OwnedHostApi { unsafe fn copy_from(host: &NemoRelayNativeHostApiV1) -> Self { - if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH + 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) }) @@ -3941,7 +3941,7 @@ pub unsafe fn export_plugin_v2( export_plugin_checked( host_ref, out, - NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, std::mem::size_of::(), || plugin, ) @@ -3999,7 +3999,7 @@ where export_plugin_checked( host_ref, out, - NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, + NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, std::mem::size_of::(), constructor, ) @@ -4079,7 +4079,7 @@ 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 typed LLM dispatch from [`NemoRelayNativeHostApiV4`]. +/// the plugin requires targeted LLM continuations from [`NemoRelayNativeHostApiV4`]. #[macro_export] macro_rules! nemo_relay_plugin_v2 { ($symbol:ident, $constructor:expr) => { diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 8e4c9e361..575512bf4 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -14,11 +14,11 @@ use std::sync::{ use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, - Event, EventCategory, EventSanitizeFields, Json, LlmCallFailureV2, LlmCodecIdentity, + Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmContinuationFailureV2, LlmHttpFailureV2, LlmJsonStream, LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, - NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TYPED_LLM_DISPATCH, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareKind, + NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, + NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, @@ -39,7 +39,7 @@ use serde_json::{Map, json}; fn native_api_v2_retry_policy_is_derived_from_http_semantics() { for status in [408, 425, 429, 500, 502, 503, 504] { assert!( - LlmCallFailureV2::Http { + LlmContinuationFailureV2::Http { failure: LlmHttpFailureV2 { status, body: String::new(), @@ -52,7 +52,7 @@ fn native_api_v2_retry_policy_is_derived_from_http_semantics() { } for status in [400, 401, 404, 409, 422, 501] { assert!( - !LlmCallFailureV2::Http { + !LlmContinuationFailureV2::Http { failure: LlmHttpFailureV2 { status, body: String::new(), @@ -68,7 +68,7 @@ fn native_api_v2_retry_policy_is_derived_from_http_semantics() { LlmNonHttpFailureKindV2::Timeout, ] { assert!( - LlmCallFailureV2::NonHttp { + LlmContinuationFailureV2::NonHttp { failure: LlmNonHttpFailureV2 { kind, message: String::new(), @@ -85,7 +85,7 @@ fn native_api_v2_retry_policy_is_derived_from_http_semantics() { LlmNonHttpFailureKindV2::Internal, ] { assert!( - !LlmCallFailureV2::NonHttp { + !LlmContinuationFailureV2::NonHttp { failure: LlmNonHttpFailureV2 { kind, message: String::new(), @@ -399,7 +399,7 @@ static LLM_REQUEST_INTERCEPT_REGISTRATION: Mutex(), test_host().struct_size From a4558413abc7133710457df980dfaf6289be72e9 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 31 Jul 2026 14:24:10 -0600 Subject: [PATCH 06/32] refactor(plugin): move targeted LLM dispatch into core Signed-off-by: Bryan Bednarski --- Cargo.lock | 1 + crates/cli/src/gateway/mod.rs | 58 +-- crates/cli/src/server/mod.rs | 9 - .../tests/coverage/shared/gateway_tests.rs | 68 ---- crates/core/Cargo.toml | 1 + crates/core/src/api/llm.rs | 15 +- crates/core/src/api/runtime.rs | 7 +- .../src/api/runtime/llm_dispatch_context.rs | 348 ++++++++++++++++-- crates/core/src/error.rs | 8 +- crates/core/src/plugin/dynamic/native.rs | 72 +--- .../tests/fixtures/native_plugin/src/lib.rs | 173 ++++++++- .../tests/integration/native_plugin_tests.rs | 144 +++++++- .../tests/unit/llm_dispatch_context_tests.rs | 347 +++++++++++++++++ crates/core/tests/unit/native_plugin_tests.rs | 17 +- crates/plugin/README.md | 5 + .../dynamic-plugins/native-dynamic/about.mdx | 8 +- 16 files changed, 1044 insertions(+), 237 deletions(-) create mode 100644 crates/core/tests/unit/llm_dispatch_context_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 1ad534707..e56ad37fd 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", diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index c9a39942f..ad9893d45 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -26,8 +26,7 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::runtime::{ - LlmCodecIdentity, LlmDispatchTargetContext, LlmExecutionNextFn, LlmJsonStream, - LlmStreamExecutionNextFn, TASK_SCOPE_STACK, current_llm_dispatch_target, + LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, }; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ @@ -177,7 +176,6 @@ async fn run_unmanaged_gateway( &prepared.body_bytes, &prepared.headers, None, - None, ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config), ) .await?; @@ -320,7 +318,6 @@ fn build_buffered_func( upstream_failures: CapturedUpstreamFailuresRef, ) -> LlmExecutionNextFn { let http = state.http.clone(); - let targeted_http = state.targeted_http.clone(); let method = prepared.method.clone(); let url = prepared.upstream_url.clone(); let body_bytes = prepared.body_bytes.clone(); @@ -329,7 +326,6 @@ fn build_buffered_func( ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config); Arc::new(move |request| { let http = http.clone(); - let targeted_http = targeted_http.clone(); let forwarding = forwarding.clone(); let method = method.clone(); let url = url.clone(); @@ -337,17 +333,14 @@ fn build_buffered_func( let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); Box::pin(async move { - let typed_target = current_llm_dispatch_target(); - let retry_aware = typed_target.is_some() || retry_aware_dispatch(&request); - let http = typed_target.as_ref().map_or(&http, |_| &targeted_http); + let retry_aware = retry_aware_dispatch(&request); let response = match forward_upstream_request( - http, + &http, &method, &url, &body_bytes, &headers, Some(&request), - typed_target.as_ref(), forwarding, ) .await @@ -519,7 +512,6 @@ fn build_streaming_func( upstream_failures: CapturedUpstreamFailuresRef, ) -> LlmStreamExecutionNextFn { let http = state.http.clone(); - let targeted_http = state.targeted_http.clone(); let method = prepared.method.clone(); let url = prepared.upstream_url.clone(); let body_bytes = prepared.body_bytes.clone(); @@ -528,7 +520,6 @@ fn build_streaming_func( ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config); Arc::new(move |request| { let http = http.clone(); - let targeted_http = targeted_http.clone(); let forwarding = forwarding.clone(); let method = method.clone(); let url = url.clone(); @@ -536,17 +527,14 @@ fn build_streaming_func( let headers = headers.clone(); let upstream_failures = upstream_failures.clone(); Box::pin(async move { - let typed_target = current_llm_dispatch_target(); - let retry_aware = typed_target.is_some() || retry_aware_dispatch(&request); - let http = typed_target.as_ref().map_or(&http, |_| &targeted_http); + let retry_aware = retry_aware_dispatch(&request); let response = match forward_upstream_request( - http, + &http, &method, &url, &body_bytes, &headers, Some(&request), - typed_target.as_ref(), forwarding, ) .await @@ -800,7 +788,6 @@ async fn forward_upstream_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, - typed_target: Option<&LlmDispatchTargetContext>, forwarding: ProviderForwarding, ) -> Result { debug_assert_eq!( @@ -810,11 +797,10 @@ async fn forward_upstream_request( .provider_credential_present(), crate::provider_auth::has_provider_credential(headers) ); - let effective = effective_dispatch_request_with_target( + let effective = build_effective_dispatch_request( body_bytes, headers, effective_request, - typed_target, url, method, forwarding.source_route, @@ -881,22 +867,20 @@ fn effective_dispatch_request( url: &str, route: ProviderRoute, ) -> EffectiveUpstreamRequest { - effective_dispatch_request_with_target( + build_effective_dispatch_request( body_bytes, headers, effective_request, - None, url, &Method::POST, route, ) } -fn effective_dispatch_request_with_target( +fn build_effective_dispatch_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, - typed_target: Option<&LlmDispatchTargetContext>, url: &str, method: &Method, route: ProviderRoute, @@ -909,31 +893,6 @@ fn effective_dispatch_request_with_target( let Some((body_bytes, body_reencoded)) = reencode_request_body(request, body_bytes) else { return source_request(body_bytes, headers, method, url, route); }; - if let Some(target) = typed_target { - let mut target_headers = HeaderMap::new(); - for (name, value) in target.headers() { - let name = HeaderName::from_bytes(name.as_bytes()) - .expect("native API v2 target header names were validated"); - let value = HeaderValue::from_str(value) - .expect("native API v2 target header values were validated"); - target_headers.insert(name, value); - } - target_headers - .entry(header::CONTENT_TYPE) - .or_insert(HeaderValue::from_static("application/json")); - let target_route = ProviderRoute::from_dispatch_override(target.route()) - .expect("native API v2 target routes originate from a typed route enum"); - let method = Method::from_bytes(target.method().as_bytes()) - .expect("native API v2 target methods were validated"); - return EffectiveUpstreamRequest { - body_bytes, - headers: target_headers, - method, - url: target.url().to_owned(), - target_route, - credential_policy: TargetCredentialPolicy::ExplicitTarget, - }; - } let overrides = dispatch_overrides(&request.headers); let credential_policy = if overrides.is_explicit_target() { crate::provider_auth::remove_provider_credentials(&mut headers); @@ -1171,7 +1130,6 @@ async fn passthrough_streaming( &prepared.body_bytes, &prepared.headers, None, - None, ProviderForwarding::new(prepared.provider, prepared.authorization, &state.config), ) .await?; diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index b52a07e36..c4ca3ea62 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -63,7 +63,6 @@ pub(crate) struct AppState { pub(crate) transparent_proxy_credential: Option, pub(crate) http: Client, - pub(crate) targeted_http: Client, pub(crate) sessions: SessionManager, pub(crate) last_activity: Arc>, pub(crate) bootstrap_shutdown: Option, @@ -513,13 +512,6 @@ impl AppState { .read_timeout(HTTP_READ_TIMEOUT) .build() .expect("gateway HTTP client configuration is valid"); - let targeted_http = Client::builder() - .connect_timeout(HTTP_CONNECT_TIMEOUT) - .timeout(HTTP_REQUEST_TIMEOUT) - .read_timeout(HTTP_READ_TIMEOUT) - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("targeted gateway HTTP client configuration is valid"); Self { config, bootstrap_fingerprint, @@ -527,7 +519,6 @@ impl AppState { require_provider_client_token, transparent_proxy_credential, http, - targeted_http, sessions, last_activity: Arc::new(Mutex::new(Instant::now())), bootstrap_shutdown, diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index d71998f13..a775ddf93 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -13,7 +13,6 @@ use axum::response::IntoResponse; use http_body_util::BodyExt; use reqwest::Client; use serde_json::{Map, json}; -use std::collections::BTreeMap; use tokio::io::{AsyncReadExt, AsyncWriteExt}; fn test_http_client() -> Client { @@ -703,71 +702,6 @@ fn internal_dispatch_controls_are_consumed_and_never_forwarded() { assert!(retry_aware_dispatch(&request)); } -#[test] -fn typed_dispatch_target_overrides_legacy_headers_without_exposing_transport_data() { - let original_body = Bytes::from_static(br#"{"model":"original"}"#); - let mut original_headers = HeaderMap::new(); - original_headers.insert( - header::AUTHORIZATION, - HeaderValue::from_static("Bearer source-secret"), - ); - original_headers.insert("x-source", HeaderValue::from_static("source")); - let request = LlmRequest { - headers: Map::from_iter([ - ( - INTERNAL_DISPATCH_URL_HEADER.to_string(), - json!("http://attacker.invalid/v1/chat/completions"), - ), - ("authorization".into(), json!("Bearer request-secret")), - ("x-request".into(), json!("request")), - ]), - content: json!({"model": "selected"}), - }; - let target = LlmDispatchTargetContext::new( - "POST".into(), - "http://selected.invalid/v1/responses".into(), - "openai_responses".into(), - BTreeMap::from([ - ("authorization".into(), "Bearer target-secret".into()), - ("x-target".into(), "selected".into()), - ]), - ); - - let effective = effective_dispatch_request_with_target( - &original_body, - &original_headers, - Some(&request), - Some(&target), - "http://default.invalid/v1/chat/completions", - &Method::GET, - ProviderRoute::OpenAiChatCompletions, - ); - - assert_eq!(effective.method, Method::POST); - assert_eq!( - effective.url, - "http://selected.invalid/v1/responses".to_string() - ); - assert_eq!(effective.target_route, ProviderRoute::OpenAiResponses); - assert_eq!( - effective.headers.get(header::AUTHORIZATION).unwrap(), - "Bearer target-secret" - ); - assert_eq!(effective.headers.get("x-target").unwrap(), "selected"); - assert_eq!( - effective.headers.get(header::CONTENT_TYPE).unwrap(), - "application/json" - ); - assert!(effective.headers.get("x-source").is_none()); - assert!(effective.headers.get("x-request").is_none()); - assert!( - effective - .headers - .get(INTERNAL_DISPATCH_URL_HEADER) - .is_none() - ); -} - #[test] fn explicit_keyless_target_drops_source_credentials() { let mut source_headers = HeaderMap::new(); @@ -2010,7 +1944,6 @@ async fn passthrough_rejects_unsupported_provider_path_directly() { require_provider_client_token: false, transparent_proxy_credential: None, http: test_http_client(), - targeted_http: test_http_client(), sessions: SessionManager::new(config), last_activity: std::sync::Arc::new(std::sync::Mutex::new(std::time::Instant::now())), bootstrap_shutdown: None, @@ -2049,7 +1982,6 @@ async fn models_rejects_non_get_requests_directly() { require_provider_client_token: false, transparent_proxy_credential: None, http: test_http_client(), - targeted_http: test_http_client(), sessions: SessionManager::new(config), last_activity: std::sync::Arc::new(std::sync::Mutex::new(std::time::Instant::now())), bootstrap_shutdown: None, 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 b4e9165c6..52e065df1 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -24,8 +24,11 @@ pub use continuation_context::MiddlewareContinuationContext; #[cfg(test)] pub(crate) use continuation_context::MiddlewareContinuationLease; pub use global::global_context; -#[doc(hidden)] -pub use llm_dispatch_context::{LlmDispatchTargetContext, current_llm_dispatch_target}; +#[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/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 3b24cc619..9940be3e5 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -1,10 +1,30 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Invocation-scoped transport target for managed LLM continuations. +//! 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, Method, StatusCode, Url}; + +use crate::api::llm::LlmRequest; +use crate::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn}; +use crate::codec::streaming::SseEventDecoder; +use crate::error::{FlowError, Result, UpstreamFailure, UpstreamFailureClass}; +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_UPSTREAM_ERROR_BODY_BYTES: usize = 16 * 1024; +const MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES: usize = 1024; tokio::task_local! { static TASK_LLM_DISPATCH_TARGET: LlmDispatchTargetContext; @@ -12,69 +32,144 @@ tokio::task_local! { /// Validated provider transport target bound to one LLM continuation invocation. /// -/// This internal bridge keeps transport data out of [`crate::api::llm::LlmRequest`] -/// while allowing the terminal gateway callback to dispatch a plugin-selected -/// provider request. +/// 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, Debug, PartialEq, Eq)] +#[derive(Clone)] pub struct LlmDispatchTargetContext { - method: String, - url: String, + method: Method, + url: Url, route: String, - headers: BTreeMap, + headers: HeaderMap, } impl LlmDispatchTargetContext { - /// Construct a validated target from the native-plugin adapter. - #[doc(hidden)] - #[must_use] - pub fn new( + /// Validate and construct a target for one continuation invocation. + pub(crate) fn try_new( method: String, url: String, route: String, headers: BTreeMap, - ) -> Self { - Self { + ) -> Result { + let method = Method::from_bytes(method.as_bytes()).map_err(|_| { + FlowError::InvalidArgument("LLM continuation method was invalid or prohibited".into()) + })?; + if matches!(method, Method::CONNECT | Method::TRACE) { + return Err(FlowError::InvalidArgument( + "LLM continuation method was invalid or prohibited".into(), + )); + } + 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" + ))); + } + 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 { method, url, route, - headers, - } + headers: validated_headers, + }) } /// HTTP method selected for this invocation. #[doc(hidden)] #[must_use] - pub fn method(&self) -> &str { + pub(crate) fn method(&self) -> &Method { &self.method } /// Absolute provider URL selected for this invocation. #[doc(hidden)] #[must_use] - pub fn url(&self) -> &str { + pub(crate) fn url(&self) -> &Url { &self.url } - /// Relay provider-route identifier selected for this invocation. - #[doc(hidden)] - #[must_use] - pub fn route(&self) -> &str { + #[cfg(test)] + pub(crate) fn route(&self) -> &str { &self.route } /// Explicit provider headers selected for this invocation. #[doc(hidden)] #[must_use] - pub fn headers(&self) -> &BTreeMap { + pub(crate) fn headers(&self) -> &HeaderMap { &self.headers } } -/// Return the typed target bound to the current continuation invocation. -#[doc(hidden)] -#[must_use] -pub fn current_llm_dispatch_target() -> Option { +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("method", &self.method) + .field("url", &redacted_url) + .field("route", &self.route) + .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(Clone::clone).ok() } @@ -85,3 +180,202 @@ pub(crate) async fn scope_llm_dispatch_target( ) -> F::Output { TASK_LLM_DISPATCH_TARGET.scope(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).await?; + let status = response.status(); + let headers = safe_failure_headers(response.headers()); + if !status.is_success() { + let bytes = bounded_response_body(response).await?; + return Err(http_error(status, headers, &bytes)); + } + let bytes = response.bytes().await.map_err(transport_error)?; + serde_json::from_slice(&bytes).map_err(|_| http_error(status, headers, &bytes)) +} + +async fn dispatch_stream( + target: &LlmDispatchTargetContext, + request: LlmRequest, +) -> Result { + let response = send(target, request).await?; + let status = response.status(); + if !status.is_success() { + let headers = safe_failure_headers(response.headers()); + let body = bounded_response_body(response).await?; + return Err(http_error(status, headers, &body)); + } + + 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(error)); + return; + } + } + } + match decoder.finish() { + Ok(Some(event)) => yield Ok(event.data), + Ok(None) => {} + Err(error) => yield Err(error), + } + })) +} + +async fn send(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result { + let body = serde_json::to_vec(&request.content) + .map_err(|error| FlowError::InvalidArgument(error.to_string()))?; + let mut outbound = targeted_http_client() + .request(target.method().clone(), target.url().clone()) + .body(body); + for (name, value) in target.headers() { + outbound = outbound.header(name, value); + } + outbound.send().await.map_err(transport_error) +} + +fn targeted_http_client() -> &'static Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + Client::builder() + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .timeout(HTTP_REQUEST_TIMEOUT) + .read_timeout(HTTP_READ_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("core targeted LLM HTTP client configuration is valid") + }) +} + +async fn bounded_response_body(response: reqwest::Response) -> Result> { + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while body.len() < MAX_UPSTREAM_ERROR_BODY_BYTES { + let Some(chunk) = stream.next().await else { + break; + }; + let chunk = chunk.map_err(transport_error)?; + let remaining = MAX_UPSTREAM_ERROR_BODY_BYTES - body.len(); + body.extend_from_slice(&chunk[..chunk.len().min(remaining)]); + } + Ok(body) +} + +fn transport_error(error: reqwest::Error) -> FlowError { + let timeout = error.is_timeout(); + 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_ERROR_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 { + headers + .iter() + .filter_map(|(name, value)| { + let name = name.as_str(); + matches!( + name, + "retry-after" + | "request-id" + | "traceparent" + | "x-request-id" + | "x-ratelimit-limit" + | "x-ratelimit-remaining" + | "x-ratelimit-reset" + | "ratelimit-limit" + | "ratelimit-remaining" + | "ratelimit-reset" + ) + .then(|| { + ( + name.to_owned(), + bounded_utf8( + String::from_utf8_lossy(value.as_bytes()).into_owned(), + MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES, + ), + ) + }) + }) + .collect() +} + +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() +} + +#[cfg(test)] +#[path = "../../../tests/unit/llm_dispatch_context_tests.rs"] +mod tests; diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 270f9737a..0768388f7 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -12,7 +12,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use thiserror::Error; -/// Stable classification for a failure from an upstream provider attempt. +/// 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 +34,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 +48,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 +122,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 08e56f7d1..5f9068b47 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -2302,75 +2302,19 @@ 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 url = match reqwest::Url::parse(&invocation.target.url) { - Ok(url) - if matches!(url.scheme(), "http" | "https") - && url.has_host() - && url.username().is_empty() - && url.password().is_none() => - { - url - } - _ => { - set_native_last_error( - "LLM continuation target must be an absolute HTTP(S) URL without user info", - ); - return Err(NemoRelayStatus::InvalidArg); - } - }; - let method = match reqwest::Method::from_bytes(invocation.target.method.as_bytes()) { - Ok(method) if method != reqwest::Method::CONNECT && method != reqwest::Method::TRACE => { - method - } - _ => { - set_native_last_error("LLM continuation method was invalid or prohibited"); - return Err(NemoRelayStatus::InvalidArg); - } - }; - for (name, value) in &invocation.target.headers { - let Ok(parsed_name) = reqwest::header::HeaderName::from_bytes(name.as_bytes()) else { - set_native_last_error("LLM continuation contained an invalid target header name"); - return Err(NemoRelayStatus::InvalidArg); - }; - if prohibited_target_header(&parsed_name) { - set_native_last_error(format!( - "LLM continuation target header {parsed_name} is host-owned or prohibited" - )); - return Err(NemoRelayStatus::InvalidArg); - } - if reqwest::header::HeaderValue::from_str(value).is_err() { - set_native_last_error(format!( - "LLM continuation target header {parsed_name} had an invalid value" - )); - return Err(NemoRelayStatus::InvalidArg); - } - } - let target = LlmDispatchTargetContext::new( - method.as_str().to_owned(), - url.to_string(), + let target = LlmDispatchTargetContext::try_new( + invocation.target.method, + invocation.target.url, invocation.target.route.as_str().into(), invocation.target.headers, - ); + ) + .map_err(|error| { + set_native_last_error(error.to_string()); + NemoRelayStatus::InvalidArg + })?; Ok((invocation.request, target)) } -fn prohibited_target_header(name: &reqwest::header::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" - ) -} - fn non_http_llm_failure( kind: LlmNonHttpFailureKindV2, message: String, diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index f8c077123..384320a1c 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -7,10 +7,12 @@ use std::sync::atomic::{AtomicBool, Ordering}; use nemo_relay_plugin::{ CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, - Json, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareCb, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, - NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, + Json, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationRouteV2, + LlmContinuationTargetV2, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, NativePlugin, + NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, + NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, + NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, @@ -302,6 +304,169 @@ nemo_relay_plugin::nemo_relay_plugin_v2!( || FixtureNativePlugin ); +struct TargetedFixturePlugin; + +struct TargetedFixtureState { + host: NemoRelayNativeHostApiV4, + url: String, +} + +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 host = *ctx + .host_api_v4() + .ok_or_else(|| "targeted fixture requires native API v2".to_string())?; + 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 state = Box::into_raw(Box::new(TargetedFixtureState { host, url })).cast(); + let status = unsafe { + ctx.register_async_llm_execution_v2_raw( + "fixture_targeted_llm", + 0, + targeted_fixture_callback, + state, + Some(drop_targeted_fixture_state), + ) + }; + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + unsafe { drop(Box::from_raw(state.cast::())) }; + Err(format!("targeted fixture registration failed: {status:?}")) + } + } +} + +unsafe extern "C" fn drop_targeted_fixture_state(user_data: *mut c_void) { + if !user_data.is_null() { + unsafe { drop(Box::from_raw(user_data.cast::())) }; + } +} + +struct TargetedFixtureResult { + host: NemoRelayNativeHostApiV1, + sender: std::sync::mpsc::SyncSender>, +} + +unsafe extern "C" fn targeted_fixture_result( + user_data: *mut c_void, + outcome_json: *const NemoRelayNativeString, +) { + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let outcome = unsafe { raw_host_string_value(&state.host, outcome_json) } + .ok_or_else(|| "targeted fixture received an invalid result string".to_string()) + .and_then(|json| serde_json::from_str(&json).map_err(|error| error.to_string())); + let _ = state.sender.send(outcome); +} + +unsafe extern "C" fn targeted_fixture_callback( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let Some(state) = (unsafe { user_data.cast::().as_ref() }) else { + return NemoRelayNativeAsyncCallbackState::Complete as u32; + }; + let result = (|| { + let invocation = unsafe { raw_host_string_value(&state.host.v3.v1, invocation_json) } + .ok_or_else(|| "targeted fixture received an invalid invocation".to_string())?; + let mut invocation: Json = + serde_json::from_str(&invocation).map_err(|error| error.to_string())?; + let mut request: LlmRequest = serde_json::from_value( + invocation + .get_mut("request") + .ok_or_else(|| "targeted fixture invocation omitted request".to_string())? + .take(), + ) + .map_err(|error| error.to_string())?; + request.content["fixture_targeted"] = json!(true); + let dispatch = LlmContinuationInvocationV2 { + request, + target: LlmContinuationTargetV2 { + method: "POST".into(), + url: state.url.clone(), + route: LlmContinuationRouteV2::OpenaiChat, + headers: std::collections::BTreeMap::from([( + "authorization".into(), + "Bearer fixture-target".into(), + )]), + }, + }; + let dispatch = serde_json::to_string(&dispatch).map_err(|error| error.to_string())?; + let dispatch = unsafe { raw_host_string(&state.host.v3.v1, &dispatch) }; + if dispatch.is_null() { + return Err("targeted fixture could not allocate dispatch".into()); + } + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + let callback_state = Box::into_raw(Box::new(TargetedFixtureResult { + host: state.host.v3.v1, + sender, + })) + .cast(); + let status = unsafe { + (state.host.async_llm_next_invoke_result_v2)( + next, + dispatch, + targeted_fixture_result, + callback_state, + ) + }; + unsafe { (state.host.v3.v1.string_free)(dispatch) }; + if status != NemoRelayStatus::Ok { + unsafe { drop(Box::from_raw(callback_state.cast::())) }; + return Err(format!("targeted fixture dispatch failed: {status:?}")); + } + receiver + .recv_timeout(std::time::Duration::from_secs(5)) + .map_err(|error| error.to_string())? + })(); + unsafe { (state.host.v3.async_next_release)(next) }; + + match result { + Ok(LlmContinuationOutcomeV2::Success { response }) => { + if let Ok(response) = serde_json::to_string(&response) { + let response = unsafe { raw_host_string(&state.host.v3.v1, &response) }; + if !response.is_null() { + unsafe { + (state.host.v3.async_completion_resolve_json)(completion, response); + (state.host.v3.v1.string_free)(response); + } + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + } + unsafe { + reject_async_completion( + &state.host.v3, + completion, + "targeted fixture could not serialize response", + ) + }; + } + Ok(LlmContinuationOutcomeV2::Failure { error }) => unsafe { + reject_async_completion(&state.host.v3, completion, &format!("{error:?}")) + }, + Err(error) => unsafe { reject_async_completion(&state.host.v3, completion, &error) }, + } + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +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, diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 7b5c491ab..c00dade0a 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}; @@ -1301,6 +1301,62 @@ fn native_api_v2_plugin_requires_the_v2_manifest_contract() { assert!(error.contains("entry symbol"), "{error}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn native_api_v2_fixture_dispatches_through_core_without_a_cli_gateway() { + 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 original_provider_called = Arc::new(AtomicBool::new(false)); + let original_provider_called_for_fn = original_provider_called.clone(); + let response = 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 + .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")); + + clear_plugin_configuration().expect("targeted fixture configuration should clear"); + activation.clear(); +} + #[test] fn native_manifest_writer_escapes_toml_strings() { let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); @@ -2095,6 +2151,92 @@ struct BuiltFixture { library_path: PathBuf, } +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"); + let address = listener.local_addr().expect("embedded provider address"); + let (request_tx, request) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let (mut socket, _) = listener.accept().expect("embedded provider should accept"); + socket + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("embedded provider timeout should configure"); + let mut request_bytes = Vec::new(); + let mut buffer = [0_u8; 4096]; + loop { + let read = socket.read(&mut buffer).expect("provider request read"); + 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) + .expect("embedded test should receive provider request"); + 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()) + .expect("embedded provider should write headers"); + socket + .write_all(body) + .expect("embedded provider should write body"); + }); + 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() { + thread + .join() + .expect("embedded provider thread should finish"); + } + } +} + 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/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs new file mode 100644 index 000000000..a7ee1d730 --- /dev/null +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -0,0 +1,347 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +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 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"); + let address = listener.local_addr().expect("fake provider address"); + let (request_tx, request) = mpsc::channel(); + let thread = std::thread::spawn(move || { + let (mut socket, _) = listener.accept().expect("fake provider should accept"); + 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"); + socket + .write_all(&response) + .expect("fake provider should write response"); + }); + 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("POST".into(), url, "openai_chat".into(), 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"}), + } +} + +#[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_llm_dispatch_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 buffered_http_failure_is_bounded_and_filters_headers() { + let body = vec![b'x'; MAX_UPSTREAM_ERROR_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_ERROR_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 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")], + 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_llm_dispatch_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(); +} + +#[test] +fn target_validation_rejects_unsafe_transport_inputs() { + for (method, url, headers) in [ + ("TRACE", "https://provider.example/v1", BTreeMap::new()), + ("POST", "ftp://provider.example/v1", BTreeMap::new()), + ( + "POST", + "https://user:secret@provider.example/v1", + BTreeMap::new(), + ), + ( + "POST", + "https://provider.example/v1", + BTreeMap::from([("host".into(), "attacker.invalid".into())]), + ), + ( + "POST", + "https://provider.example/v1", + BTreeMap::from([( + "x-nemo-relay-internal-dispatch-url".into(), + "http://attacker.invalid".into(), + )]), + ), + ] { + assert!( + LlmDispatchTargetContext::try_new( + method.into(), + url.into(), + "openai_chat".into(), + headers, + ) + .is_err() + ); + } +} + +#[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_llm_dispatch_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)); +} diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index cf12e6593..206a7d5eb 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -1554,10 +1554,16 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { Box::pin(async move { assert!(request.headers.is_empty()); assert_eq!(target.method(), "POST"); - assert_eq!(target.url(), "https://provider.example/v1/chat/completions"); + assert_eq!( + target.url().as_str(), + "https://provider.example/v1/chat/completions" + ); assert_eq!(target.route(), "openai_chat"); assert_eq!( - target.headers().get("authorization").map(String::as_str), + target + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), Some("Bearer target-secret") ); Err(FlowError::Upstream(crate::error::UpstreamFailure { @@ -2146,8 +2152,11 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { let target = current_llm_dispatch_target().expect("typed target is bound"); Ok(json!({ "request": request.content, - "url": target.url(), - "authorization": target.headers().get("authorization"), + "url": target.url().as_str(), + "authorization": target + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), })) }) })), diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 688e8a194..b52da04e2 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -121,6 +121,11 @@ calls return provider JSON. Provider rejections return an HTTP status, bounded body, and safe response headers; failures without an HTTP response use a small transport-oriented kind. +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 an opaque host-owned stream; request one JSON event at a time, then cancel and release it exactly once. No Rust future, trait object, `serde_json::Value`, or allocator-owned Rust string crosses the ABI diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 652687ab1..2fd946b80 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -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. @@ -150,6 +150,12 @@ request middleware. Relay returns either: - a host-owned provider stream pulled one JSON event at a time, including typed setup and late failures. +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. + HTTP retryability is derived by the SDK from status alone: `408`, `425`, `429`, `500`, `502`, `503`, and `504` are retryable. Transport and timeout failures are retryable; other non-HTTP failures are not. Relay does not inspect provider From 79d6d1c3d943c55254e9d29ee4bd910e2fbe7e51 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 31 Jul 2026 14:50:27 -0600 Subject: [PATCH 07/32] fix(plugin): harden targeted LLM continuations Signed-off-by: Bryan Bednarski --- .../src/api/runtime/continuation_context.rs | 2 +- .../src/api/runtime/llm_dispatch_context.rs | 53 ++++++++++++++----- crates/core/src/plugin/dynamic/native.rs | 22 ++++++-- .../tests/fixtures/native_plugin/src/lib.rs | 1 - .../tests/integration/native_plugin_tests.rs | 24 ++++++++- .../tests/unit/llm_dispatch_context_tests.rs | 42 +++++++++++++-- crates/core/tests/unit/native_plugin_tests.rs | 4 ++ 7 files changed, 123 insertions(+), 25 deletions(-) diff --git a/crates/core/src/api/runtime/continuation_context.rs b/crates/core/src/api/runtime/continuation_context.rs index 71666a145..2d0aba777 100644 --- a/crates/core/src/api/runtime/continuation_context.rs +++ b/crates/core/src/api/runtime/continuation_context.rs @@ -117,7 +117,7 @@ impl MiddlewareContinuationContext { C: FnOnce() -> F, F: Future, { - scope_llm_dispatch_target(target, self.run(async move { callback().await })).await + scope_llm_dispatch_target(target, self.invoke(callback)).await } /// Invoke a callback and poll its future with the captured Relay context. diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 9940be3e5..cb389f6b2 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -79,6 +79,11 @@ impl LlmDispatchTargetContext { "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" @@ -210,14 +215,17 @@ pub(crate) fn targeted_llm_stream_execution( } async fn dispatch_buffered(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result { - let response = send(target, request).await?; + let response = send(target, request, Some(HTTP_REQUEST_TIMEOUT)).await?; let status = response.status(); let headers = safe_failure_headers(response.headers()); if !status.is_success() { - let bytes = bounded_response_body(response).await?; + let bytes = bounded_response_body(target, response).await?; return Err(http_error(status, headers, &bytes)); } - let bytes = response.bytes().await.map_err(transport_error)?; + let bytes = response + .bytes() + .await + .map_err(|error| transport_error(target, error))?; serde_json::from_slice(&bytes).map_err(|_| http_error(status, headers, &bytes)) } @@ -225,14 +233,15 @@ async fn dispatch_stream( target: &LlmDispatchTargetContext, request: LlmRequest, ) -> Result { - let response = send(target, request).await?; + 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(response).await?; + let body = bounded_response_body(target, response).await?; return Err(http_error(status, headers, &body)); } + let target = target.clone(); let mut decoder = SseEventDecoder::new(); let mut bytes = response.bytes_stream(); Ok(LlmJsonStream::new(stream! { @@ -250,7 +259,7 @@ async fn dispatch_stream( } } Err(error) => { - yield Err(transport_error(error)); + yield Err(transport_error(&target, error)); return; } } @@ -263,7 +272,11 @@ async fn dispatch_stream( })) } -async fn send(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result { +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() @@ -272,7 +285,13 @@ async fn send(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result< for (name, value) in target.headers() { outbound = outbound.header(name, value); } - outbound.send().await.map_err(transport_error) + 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 { @@ -280,7 +299,6 @@ fn targeted_http_client() -> &'static Client { CLIENT.get_or_init(|| { Client::builder() .connect_timeout(HTTP_CONNECT_TIMEOUT) - .timeout(HTTP_REQUEST_TIMEOUT) .read_timeout(HTTP_READ_TIMEOUT) .redirect(reqwest::redirect::Policy::none()) .build() @@ -288,22 +306,33 @@ fn targeted_http_client() -> &'static Client { }) } -async fn bounded_response_body(response: reqwest::Response) -> Result> { +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_ERROR_BODY_BYTES { let Some(chunk) = stream.next().await else { break; }; - let chunk = chunk.map_err(transport_error)?; + let chunk = chunk.map_err(|error| transport_error(target, error))?; let remaining = MAX_UPSTREAM_ERROR_BODY_BYTES - body.len(); body.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } Ok(body) } -fn transport_error(error: reqwest::Error) -> FlowError { +fn transport_error(target: &LlmDispatchTargetContext, error: reqwest::Error) -> FlowError { let timeout = error.is_timeout(); + let diagnostic = error.without_url(); + log::warn!( + target: "nemo_relay.llm", + 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 { diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 5f9068b47..9cd4bdc4e 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -1774,11 +1774,23 @@ async fn invoke_native_async_callback_with_lane( }) }; let callback_result = if blocking { - tokio::task::spawn_blocking(invoke).await.map_err(|error| { - FlowError::Internal(format!( - "native API v2 blocking callback task failed: {error}" - )) - })? + match tokio::task::spawn_blocking(invoke).await { + Ok(result) => result, + Err(error) => { + unsafe { + drop(Arc::from_raw( + completion_ref as *const NativeAsyncCompletion, + )); + if let Some(next) = next_ref { + drop(Arc::from_raw(next as *const NativeAsyncNext)); + } + native_string_free(invocation as *mut NemoRelayNativeString); + } + return Err(FlowError::Internal(format!( + "native API v2 blocking callback task failed: {error}" + ))); + } + } } else { invoke() }; diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 384320a1c..4c401a139 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -342,7 +342,6 @@ impl NativePlugin for TargetedFixturePlugin { if status == NemoRelayStatus::Ok { Ok(()) } else { - unsafe { drop(Box::from_raw(state.cast::())) }; Err(format!("targeted fixture registration failed: {status:?}")) } } diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index c00dade0a..6864e8965 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -1328,6 +1328,8 @@ async fn native_api_v2_fixture_dispatches_through_core_without_a_cli_gateway() { 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 original_provider_called = Arc::new(AtomicBool::new(false)); let original_provider_called_for_fn = original_provider_called.clone(); @@ -1353,7 +1355,7 @@ async fn native_api_v2_fixture_dispatches_through_core_without_a_cli_gateway() { assert!(captured.contains("authorization: Bearer fixture-target\r\n")); assert!(captured.contains("\"fixture_targeted\":true")); - clear_plugin_configuration().expect("targeted fixture configuration should clear"); + drop(cleanup); activation.clear(); } @@ -2164,10 +2166,28 @@ impl EmbeddedFakeProvider { 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 || { - let (mut socket, _) = listener.accept().expect("embedded provider should accept"); + 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!("embedded provider should accept: {error}"), + } + }; + socket + .set_nonblocking(false) + .expect("embedded provider socket should be blocking"); socket .set_read_timeout(Some(Duration::from_secs(5))) .expect("embedded provider timeout should configure"); diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index a7ee1d730..4995f8708 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -22,10 +22,28 @@ struct FakeProvider { 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 (mut socket, _) = listener.accept().expect("fake provider should accept"); + 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"); @@ -33,9 +51,17 @@ impl FakeProvider { request_tx .send(request_bytes) .expect("test should receive provider request"); - socket - .write_all(&response) - .expect("fake provider should write response"); + 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"), @@ -305,6 +331,14 @@ fn target_validation_rejects_unsafe_transport_inputs() { "http://attacker.invalid".into(), )]), ), + ( + "POST", + "https://provider.example/v1", + BTreeMap::from([ + ("Authorization".into(), "Bearer first".into()), + ("authorization".into(), "Bearer second".into()), + ]), + ), ] { assert!( LlmDispatchTargetContext::try_new( diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 206a7d5eb..338e2e5b0 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -1699,6 +1699,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { .unwrap_err(), NemoRelayStatus::InvalidArg ); + assert_last_error_contains("method was invalid or prohibited"); let mut target = test_dispatch_target( "https://provider.example/v1/chat/completions", @@ -1715,6 +1716,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { .unwrap_err(), NemoRelayStatus::InvalidArg ); + assert_last_error_contains("host-owned or prohibited"); let mut target = test_dispatch_target( "https://provider.example/v1/chat/completions", @@ -1728,6 +1730,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { .unwrap_err(), NemoRelayStatus::InvalidArg ); + assert_last_error_contains("host-owned or prohibited"); } #[test] @@ -1747,6 +1750,7 @@ fn native_api_v2_rejects_target_url_credentials() { .unwrap_err(), NemoRelayStatus::InvalidArg ); + assert_last_error_contains("absolute HTTP(S) URL"); } #[test] From c960ec9f54ea8f2e25361c353d3f561b282ba686 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 00:40:14 -0600 Subject: [PATCH 08/32] feat(plugin): add safe native API v2 continuations Signed-off-by: Bryan Bednarski --- Cargo.lock | 1 + crates/core/src/plugin/dynamic/native.rs | 682 ++++++++-- .../tests/fixtures/native_plugin/Cargo.toml | 1 + .../tests/fixtures/native_plugin/src/lib.rs | 202 +-- .../tests/integration/native_plugin_tests.rs | 265 ++++ crates/core/tests/unit/native_plugin_tests.rs | 457 +++++++ crates/plugin/Cargo.toml | 1 + crates/plugin/src/lib.rs | 50 +- crates/plugin/src/native_v2.rs | 1010 +++++++++++++++ crates/plugin/tests/typed_callbacks.rs | 1115 ++++++++++++++++- 10 files changed, 3510 insertions(+), 274 deletions(-) create mode 100644 crates/plugin/src/native_v2.rs diff --git a/Cargo.lock b/Cargo.lock index e56ad37fd..2a1123651 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1887,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/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 9cd4bdc4e..dfd7de46d 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -56,9 +56,10 @@ use nemo_relay_plugin::{ NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, - NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeAsyncLlmStreamForwardCbV2, NemoRelayNativeAsyncLlmStreamNextCbV2, + NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, @@ -928,6 +929,7 @@ fn build_native_host_api_v4() -> NemoRelayNativeHostApiV4 { native_plugin_context_register_async_llm_execution_v2, plugin_context_register_async_llm_stream_execution_v2: native_plugin_context_register_async_llm_stream_execution_v2, + async_llm_next_forward_stream_v2: native_async_llm_next_forward_stream_v2, } } @@ -1408,6 +1410,75 @@ 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) }; + } + } +} + +struct NativeStreamHandoff { + raw: usize, + armed: bool, +} + +impl NativeStreamHandoff { + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for NativeStreamHandoff { + fn drop(&mut self) { + if self.armed { + unsafe { native_async_stream_release(self.raw as *const NemoRelayNativeAsyncStream) }; + } + } +} + unsafe impl Send for NativeCallbackUserData {} unsafe impl Sync for NativeCallbackUserData {} @@ -1490,6 +1561,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>, @@ -1505,11 +1577,24 @@ 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, @@ -1564,6 +1649,92 @@ struct NativeLlmStreamOpenCallbackGuardV2 { _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, outcome: &LlmContinuationOutcomeV2) { + if !self.active { + return; + } + unsafe { + invoke_typed_llm_result_callback(self.cb, self.user_data as *mut c_void, outcome); + } + self.active = false; + } +} + +impl Drop for NativeLlmResultCallbackGuardV2 { + fn drop(&mut self) { + if self.active { + self.complete(&LlmContinuationOutcomeV2::Failure { + error: non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM continuation was cancelled".into(), + ), + }); + } + } +} + impl NativeLlmStreamOpenCallbackGuardV2 { fn success(&mut self, stream: Arc) { if !self.active { @@ -1602,6 +1773,78 @@ impl Drop for NativeLlmStreamOpenCallbackGuardV2 { } } +struct NativeLlmStreamForwardCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: usize, + stream: Arc, + active: bool, + _library_guard: Option>, +} + +impl NativeLlmStreamForwardCallbackGuardV2 { + fn complete(&mut self) { + if !self.active { + return; + } + self.active = false; + unsafe { (self.cb)(self.user_data as *mut c_void, ptr::null()) }; + } + + fn fail(&mut self, message: &str) { + if !self.active { + return; + } + self.active = false; + if let Some(message) = native_string_from_str(message) + .or_else(|| native_string_from_str("native LLM stream forwarding failed")) + { + unsafe { + (self.cb)(self.user_data as *mut c_void, message); + native_string_free(message); + } + } else { + // Allocation failure must not strand the SDK trampoline. The + // output is already failed or cancelled, so null only acts as the + // terminal wake-up in this exceptional case. + unsafe { (self.cb)(self.user_data as *mut c_void, ptr::null()) }; + } + } + + 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.fail("downstream LLM stream forwarding was cancelled"); + } +} + impl NativeAsyncStreamCallbackGuard { fn finish(&mut self) { self.active = false; @@ -1759,9 +2002,24 @@ async fn invoke_native_async_callback_with_lane( }; 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 next_handoff = NativeNextHandoff { + raw: next_ref, + armed: true, + }; let invoke = move || { - with_scope_stack(callback_scope_stack, || { - catch_unwind(AssertUnwindSafe(|| unsafe { + 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, @@ -1770,72 +2028,61 @@ async fn invoke_native_async_callback_with_lane( .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 }; - let callback_result = if blocking { - match tokio::task::spawn_blocking(invoke).await { - Ok(result) => result, - Err(error) => { - unsafe { - drop(Arc::from_raw( - completion_ref as *const NativeAsyncCompletion, - )); - if let Some(next) = next_ref { - drop(Arc::from_raw(next as *const NativeAsyncNext)); - } - native_string_free(invocation as *mut NemoRelayNativeString); - } - return Err(FlowError::Internal(format!( + let state = if blocking { + // Cleanup lives in guards captured by the blocking closure. The + // monitor reports the callback state only; runtime shutdown cannot + // strand host-owned strings or references in that task. + let blocking_task = tokio::task::spawn_blocking(invoke); + let (state_sender, state_receiver) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let state = match blocking_task.await { + Ok(Ok(state)) => NemoRelayNativeAsyncCallbackState::try_from(state).map_err(|()| { + FlowError::Internal("native async callback returned an invalid state".into()) + }), + Ok(Err(_)) => Err(FlowError::Internal("native async callback panicked".into())), + Err(error) => Err(FlowError::Internal(format!( "native API v2 blocking callback task failed: {error}" - ))); - } - } + ))), + }; + let _ = state_sender.send(state); + }); + state_receiver.await.map_err(|_| { + FlowError::Internal("native API v2 blocking callback monitor stopped".into()) + })?? } else { - invoke() - }; - let state = match callback_result { - 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())); - } - }; - 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 callback_result = invoke(); + 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 { - unsafe { - drop(Arc::from_raw( - completion_ref as *const NativeAsyncCompletion, - )); - } - if completion + 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 } @@ -2265,9 +2512,22 @@ unsafe extern "C" fn native_async_next_invoke_result( 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 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 @@ -2277,33 +2537,15 @@ unsafe extern "C" fn native_async_next_invoke_result( 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); - } - } - } - } - drop(_library_guard); + 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 } @@ -2458,9 +2700,22 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( Err(error) => return status_from_flow_error(error), }; let next_fn = next_fn.clone(); - let user_data = user_data as usize; - let library_guard = next._callback_user_data.clone(); - next.runtime.spawn(async move { + 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)), ) @@ -2478,11 +2733,15 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( error: typed_llm_failure(error), }, }; - unsafe { - invoke_typed_llm_result_callback(cb, user_data as *mut c_void, &outcome); - } - drop(library_guard); + callback_guard.complete(&outcome); + 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 } @@ -2644,6 +2903,203 @@ async fn forward_native_async_next_stream_with( 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) -> bool { + 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 false; + } + let sender = stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if sender.is_none() { + return false; + } + stream.settled.store(true, Ordering::Release); + sender + }; + abort_native_stream_downstream_tasks(stream); + drop(sender); + true +} + +async fn reject_forwarded_native_stream(stream: &NativeAsyncStream, error: FlowError) -> bool { + 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 false; + } + let sender = stream + .sender + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + let Some(sender) = sender else { + return false; + }; + stream.settled.store(true, Ordering::Release); + sender + }; + abort_native_stream_downstream_tasks(stream); + let _ = sender.send(Err(error)).await; + true +} + +/// 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 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 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, + _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 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(()) if finish_forwarded_native_stream(&callback_guard.stream) => { + callback_guard.complete(); + } + Ok(()) => { + callback_guard.fail("downstream LLM stream forwarding was cancelled"); + } + Err(error) => { + if reject_forwarded_native_stream(&callback_guard.stream, error).await { + callback_guard.fail("downstream LLM stream failed"); + } else { + callback_guard.fail("downstream LLM stream forwarding was cancelled"); + } + } + } + }); + let abort = task.abort_handle(); + downstream_aborts.insert(task.id(), abort); + drop(downstream_aborts); + drop(settlement); + let _ = start_tx.send(()); + NemoRelayStatus::Ok +} + /// Opens a streaming LLM continuation through native API v2. unsafe extern "C" fn native_async_llm_next_open_stream_v2( next: *const NemoRelayNativeAsyncNext, @@ -3313,26 +3769,40 @@ fn wrap_native_incremental_llm_stream_execution_v2( let stream_ref = Arc::into_raw(stream.clone()) as usize; let callback_user_data = user_data.ptr as usize; let callback_scope_stack = current_scope_stack(); + let invocation_guard = NativeInvocationStringGuard(invocation); + let next_handoff = NativeNextHandoff { + raw: Some(next_ref), + armed: true, + }; + let output_handoff = NativeStreamHandoff { + raw: stream_ref, + armed: true, + }; let blocking_task = tokio::task::spawn_blocking(move || { - with_scope_stack(callback_scope_stack, || { - catch_unwind(AssertUnwindSafe(|| unsafe { + let _invocation = invocation_guard; + let mut next = next_handoff; + let mut output = output_handoff; + catch_unwind(AssertUnwindSafe(|| { + // Both handles transfer to the plugin at callback entry. + next.disarm(); + output.disarm(); + with_scope_stack(callback_scope_stack, || unsafe { cb( callback_user_data as *mut c_void, invocation as *const NemoRelayNativeString, next_ref as *const NemoRelayNativeAsyncNext, stream_ref as *const NemoRelayNativeAsyncStream, ) - })) - }) + }) + })) }); let stream_for_monitor = Arc::clone(&stream); tokio::spawn(async move { - let state = blocking_task - .await + let callback_result = blocking_task.await; + let state = callback_result .ok() .and_then(std::result::Result::ok) .and_then(|state| NemoRelayNativeAsyncCallbackState::try_from(state).ok()); - unsafe { native_string_free(invocation as *mut NemoRelayNativeString) }; let error = match state { Some(NemoRelayNativeAsyncCallbackState::Pending) => None, Some(NemoRelayNativeAsyncCallbackState::Complete) 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 4c401a139..e1b9711b9 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -5,14 +5,15 @@ use std::ffi::c_void; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; +use futures::StreamExt; use nemo_relay_plugin::{ CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, - Json, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationRouteV2, - LlmContinuationTargetV2, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, + Json, LlmContinuationInvocationV2, LlmContinuationRouteV2, LlmContinuationTargetV2, + LlmJsonAsyncStreamV2, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, + LlmStreamExecutionOutcomeV2, NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, - NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativePluginContext, + NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, + NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, @@ -306,11 +307,6 @@ nemo_relay_plugin::nemo_relay_plugin_v2!( struct TargetedFixturePlugin; -struct TargetedFixtureState { - host: NemoRelayNativeHostApiV4, - url: String, -} - impl NativePlugin for TargetedFixturePlugin { fn plugin_kind(&self) -> &str { "fixture_native" @@ -321,144 +317,76 @@ impl NativePlugin for TargetedFixturePlugin { plugin_config: &Map, ctx: &mut PluginContext<'_>, ) -> nemo_relay_plugin::Result<()> { - let host = *ctx - .host_api_v4() - .ok_or_else(|| "targeted fixture requires native API v2".to_string())?; 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 state = Box::into_raw(Box::new(TargetedFixtureState { host, url })).cast(); - let status = unsafe { - ctx.register_async_llm_execution_v2_raw( - "fixture_targeted_llm", - 0, - targeted_fixture_callback, - state, - Some(drop_targeted_fixture_state), - ) - }; - if status == NemoRelayStatus::Ok { - Ok(()) - } else { - Err(format!("targeted fixture registration failed: {status:?}")) - } - } -} + let buffered_url = url.clone(); + ctx.register_async_llm_execution_v2( + "fixture_targeted_llm", + 0, + move |name, mut request, continuation| { + let url = buffered_url.clone(); + async move { + if name == "fixture_passthrough_llm" { + return continuation.call_passthrough(request).await; + } + request.content["fixture_targeted"] = json!(true); + continuation + .call(targeted_fixture_invocation(url, request)) + .await + .map_err(|error| format!("targeted fixture dispatch failed: {error:?}")) + } + }, + )?; -unsafe extern "C" fn drop_targeted_fixture_state(user_data: *mut c_void) { - if !user_data.is_null() { - unsafe { drop(Box::from_raw(user_data.cast::())) }; + 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)) + } + }, + ) } } -struct TargetedFixtureResult { - host: NemoRelayNativeHostApiV1, - sender: std::sync::mpsc::SyncSender>, -} - -unsafe extern "C" fn targeted_fixture_result( - user_data: *mut c_void, - outcome_json: *const NemoRelayNativeString, -) { - let state = unsafe { Box::from_raw(user_data.cast::()) }; - let outcome = unsafe { raw_host_string_value(&state.host, outcome_json) } - .ok_or_else(|| "targeted fixture received an invalid result string".to_string()) - .and_then(|json| serde_json::from_str(&json).map_err(|error| error.to_string())); - let _ = state.sender.send(outcome); -} - -unsafe extern "C" fn targeted_fixture_callback( - user_data: *mut c_void, - invocation_json: *const NemoRelayNativeString, - next: *const NemoRelayNativeAsyncNext, - completion: *const NemoRelayNativeAsyncCompletion, -) -> u32 { - let Some(state) = (unsafe { user_data.cast::().as_ref() }) else { - return NemoRelayNativeAsyncCallbackState::Complete as u32; - }; - let result = (|| { - let invocation = unsafe { raw_host_string_value(&state.host.v3.v1, invocation_json) } - .ok_or_else(|| "targeted fixture received an invalid invocation".to_string())?; - let mut invocation: Json = - serde_json::from_str(&invocation).map_err(|error| error.to_string())?; - let mut request: LlmRequest = serde_json::from_value( - invocation - .get_mut("request") - .ok_or_else(|| "targeted fixture invocation omitted request".to_string())? - .take(), - ) - .map_err(|error| error.to_string())?; - request.content["fixture_targeted"] = json!(true); - let dispatch = LlmContinuationInvocationV2 { - request, - target: LlmContinuationTargetV2 { - method: "POST".into(), - url: state.url.clone(), - route: LlmContinuationRouteV2::OpenaiChat, - headers: std::collections::BTreeMap::from([( - "authorization".into(), - "Bearer fixture-target".into(), - )]), - }, - }; - let dispatch = serde_json::to_string(&dispatch).map_err(|error| error.to_string())?; - let dispatch = unsafe { raw_host_string(&state.host.v3.v1, &dispatch) }; - if dispatch.is_null() { - return Err("targeted fixture could not allocate dispatch".into()); - } - let (sender, receiver) = std::sync::mpsc::sync_channel(1); - let callback_state = Box::into_raw(Box::new(TargetedFixtureResult { - host: state.host.v3.v1, - sender, - })) - .cast(); - let status = unsafe { - (state.host.async_llm_next_invoke_result_v2)( - next, - dispatch, - targeted_fixture_result, - callback_state, - ) - }; - unsafe { (state.host.v3.v1.string_free)(dispatch) }; - if status != NemoRelayStatus::Ok { - unsafe { drop(Box::from_raw(callback_state.cast::())) }; - return Err(format!("targeted fixture dispatch failed: {status:?}")); - } - receiver - .recv_timeout(std::time::Duration::from_secs(5)) - .map_err(|error| error.to_string())? - })(); - unsafe { (state.host.v3.async_next_release)(next) }; - - match result { - Ok(LlmContinuationOutcomeV2::Success { response }) => { - if let Ok(response) = serde_json::to_string(&response) { - let response = unsafe { raw_host_string(&state.host.v3.v1, &response) }; - if !response.is_null() { - unsafe { - (state.host.v3.async_completion_resolve_json)(completion, response); - (state.host.v3.v1.string_free)(response); - } - return NemoRelayNativeAsyncCallbackState::Complete as u32; - } - } - unsafe { - reject_async_completion( - &state.host.v3, - completion, - "targeted fixture could not serialize response", - ) - }; - } - Ok(LlmContinuationOutcomeV2::Failure { error }) => unsafe { - reject_async_completion(&state.host.v3, completion, &format!("{error:?}")) +fn targeted_fixture_invocation( + url: String, + request: LlmRequest, +) -> LlmContinuationInvocationV2 { + LlmContinuationInvocationV2 { + request, + target: LlmContinuationTargetV2 { + method: "POST".into(), + url, + route: LlmContinuationRouteV2::OpenaiChat, + headers: std::collections::BTreeMap::from([( + "authorization".into(), + "Bearer fixture-target".into(), + )]), }, - Err(error) => unsafe { reject_async_completion(&state.host.v3, completion, &error) }, } - NemoRelayNativeAsyncCallbackState::Complete as u32 } nemo_relay_plugin::nemo_relay_plugin_v2!( diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 6864e8965..f156f3a9f 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -1359,6 +1359,250 @@ async fn native_api_v2_fixture_dispatches_through_core_without_a_cli_gateway() { activation.clear(); } +#[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 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(); + Box::pin(async move { Ok(LlmJsonStream::new(PendingDropStream { dropped })) }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .expect("pending pass-through stream should open"); + 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(); @@ -2153,6 +2397,27 @@ 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>, diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 338e2e5b0..f8a4d35f6 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -183,6 +183,28 @@ unsafe extern "C" fn record_typed_llm_stream_next( let _ = sender.send(event); } +#[derive(Default)] +struct NativeForwardTerminalState { + callbacks: AtomicUsize, + error: Mutex>, + notified: tokio::sync::Notify, +} + +unsafe extern "C" fn record_native_forward_terminal( + user_data: *mut c_void, + error: *const NemoRelayNativeString, +) { + let state = unsafe { &*user_data.cast::() }; + state.callbacks.fetch_add(1, Ordering::SeqCst); + if !error.is_null() { + *state + .error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = read_native_string(error).ok(); + } + state.notified.notify_one(); +} + #[derive(Default)] struct NativeStreamCallbackState { error: Mutex>, @@ -1001,6 +1023,36 @@ unsafe extern "C" fn invoke_native_next_then_return_state( state.callback_state } +struct BlockingSafeCallbackState { + started: std::sync::mpsc::Sender<()>, + release: Mutex>, + freed: std::sync::mpsc::Sender<()>, +} + +unsafe extern "C" fn blocking_safe_callback( + user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + _completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let state = unsafe { &*user_data.cast::() }; + let _ = state.started.send(()); + let _ = state + .release + .lock() + .unwrap_or_else(|error| error.into_inner()) + .recv_timeout(Duration::from_secs(1)); + unsafe { native_async_next_release(next) }; + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +unsafe extern "C" fn free_blocking_safe_callback(user_data: *mut c_void) { + let state = unsafe { Box::from_raw(user_data.cast::()) }; + let freed = state.freed.clone(); + drop(state); + let _ = freed.send(()); +} + unsafe extern "C" fn invoke_native_stream_next_then_return_state( user_data: *mut c_void, invocation_json: *const NemoRelayNativeString, @@ -1627,6 +1679,86 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { } } +#[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", + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, + ), + }) + .unwrap(), + ) + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + 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 { + failure: LlmNonHttpFailureV2 { + 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() @@ -1986,6 +2118,290 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { } } +#[test] +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 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 (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_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); + assert!( + terminal + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + ); + + 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 (sender, receiver) = tokio::sync::mpsc::channel(2); + 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_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); + assert_eq!( + terminal + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_deref(), + Some("downstream LLM stream failed") + ); + + 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 (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_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); + assert!( + terminal + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_deref() + .is_some_and(|error| error.contains("cancelled")) + ); + + 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() { let runtime = tokio::runtime::Builder::new_current_thread() @@ -3752,6 +4168,47 @@ fn native_async_callback_contract_errors_abort_an_invoked_next() { } } +#[test] +fn cancelling_blocking_v2_callback_reclaims_host_resources_after_it_returns() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let (freed_tx, freed_rx) = std::sync::mpsc::channel(); + let callback_state = Box::into_raw(Box::new(BlockingSafeCallbackState { + started: started_tx, + release: Mutex::new(release_rx), + freed: freed_tx, + })); + let user_data = Arc::new(NativeCallbackUserData { + ptr: callback_state.cast(), + free_fn: Some(free_blocking_safe_callback), + _instance: None, + }); + let task = runtime.spawn(invoke_native_async_callback_blocking( + blocking_safe_callback, + user_data, + json!({"name": "cancelled"}), + Some(NativeAsyncNextInner::Llm(Arc::new(|_| { + Box::pin(async { std::future::pending::>().await }) + }))), + )); + + started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("blocking callback should start"); + task.abort(); + assert!(runtime.block_on(task).unwrap_err().is_cancelled()); + runtime.shutdown_timeout(Duration::from_millis(1)); + release_tx.send(()).unwrap(); + freed_rx + .recv_timeout(Duration::from_secs(1)) + .expect("detached cleanup must release completion, next, and callback state"); +} + #[test] fn native_async_stream_contract_errors_abort_an_invoked_next() { struct DropSignal(std::sync::mpsc::Sender<()>); 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/src/lib.rs b/crates/plugin/src/lib.rs index a183d65a6..50ce6ba5c 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -36,6 +36,13 @@ pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; +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. /// /// Relay preserves this value so plugins rebuilt with the current SDK and the @@ -1094,6 +1101,14 @@ pub type NemoRelayNativeAsyncLlmStreamOpenCbV2 = unsafe extern "C" fn( pub type NemoRelayNativeAsyncLlmStreamNextCbV2 = unsafe extern "C" fn(user_data: *mut c_void, event_json: *const NemoRelayNativeString); +/// Receives terminal settlement of a directly forwarded downstream stream. +/// +/// `error` is null after clean completion and contains a borrowed UTF-8 error +/// message after downstream failure or cancellation. Relay settles the output +/// stream before invoking this callback. +pub type NemoRelayNativeAsyncLlmStreamForwardCbV2 = + unsafe extern "C" fn(user_data: *mut c_void, error: *const NemoRelayNativeString); + /// Incremental native LLM stream intercept callback. /// /// The callback owns `next` and `stream` and must release each exactly once. @@ -1303,7 +1318,9 @@ pub struct NemoRelayNativeHostApiV4 { /// /// Relay invokes the callback on its reusable blocking executor with the /// active scope stack bound to that worker. Provider continuations invoked - /// through this table continue to execute on Relay's Tokio runtime. + /// through this table continue to execute on Relay's Tokio runtime. Once + /// invoked, the host consumes `user_data` and calls `free_fn` exactly once, + /// including when registration fails. pub plugin_context_register_async_llm_execution_v2: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -1316,7 +1333,9 @@ pub struct NemoRelayNativeHostApiV4 { /// Registers a native API v2 incremental LLM stream execution callback. /// /// Relay starts the callback on its reusable blocking executor and exposes - /// the bounded output stream to the caller concurrently. + /// the bounded output stream to the caller concurrently. Once invoked, the + /// host consumes `user_data` and calls `free_fn` exactly once, including + /// when registration fails. pub plugin_context_register_async_llm_stream_execution_v2: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, @@ -1326,6 +1345,19 @@ pub struct NemoRelayNativeHostApiV4 { user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus, + /// 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`]. + 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 NemoRelayNativeHostApiV4 {} @@ -2897,9 +2929,10 @@ impl<'a> PluginContext<'a> { /// Registers a native API v2 unary LLM execution callback. /// /// # Safety - /// The callback and user data must remain valid until deregistration or - /// `free_fn`. The callback owns its completion and `next` handles and must - /// settle/release them exactly once. + /// The host consumes `user_data` as soon as the registration function is + /// invoked and calls `free_fn` exactly once on registration failure or + /// eventual deregistration. The callback owns its completion and `next` + /// handles and must settle/release them exactly once. pub unsafe fn register_async_llm_execution_v2_raw( &mut self, name: &str, @@ -2921,9 +2954,10 @@ impl<'a> PluginContext<'a> { /// Registers a native API v2 incremental LLM stream execution callback. /// /// # 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. + /// The host consumes `user_data` as soon as the registration function is + /// invoked and calls `free_fn` exactly once on registration failure or + /// eventual deregistration. Callback-owned `next` and stream handles must + /// each be released exactly once. pub unsafe fn register_async_llm_stream_execution_v2_raw( &mut self, name: &str, diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs new file mode 100644 index 000000000..a4a6c2569 --- /dev/null +++ b/crates/plugin/src/native_v2.rs @@ -0,0 +1,1010 @@ +// 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::ptr; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::thread; +use std::time::Duration; + +use futures::channel::oneshot; +use futures::task::{ArcWake, waker_ref}; +use futures::{Stream, StreamExt}; +use serde::Deserialize; + +use super::{ + HostString, Json, LlmContinuationFailureV2, LlmContinuationInvocationV2, + LlmContinuationOutcomeV2, LlmContinuationStreamEventV2, LlmNonHttpFailureKindV2, + LlmNonHttpFailureV2, LlmRequest, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, + NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV4, NemoRelayNativeLlmStreamV2, + NemoRelayNativeString, NemoRelayStatus, PluginContext, Result, read_json_value, + read_required_host_string, set_last_error, +}; + +const BACKPRESSURE_RETRY_DELAY: Duration = Duration::from_millis(1); +const CANCELLATION_POLL_DELAY: Duration = Duration::from_millis(10); +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>, + finished: bool, + terminal: bool, +} + +struct ProviderItem { + value: std::result::Result, LlmContinuationFailureV2>, + terminal: bool, +} + +struct ContinuationInner { + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, +} + +struct BlockingThreadWaker(thread::Thread); + +impl ArcWake for BlockingThreadWaker { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.unpark(); + } +} + +fn block_on_complete_callback(future: F, is_cancelled: C) -> Option +where + F: Future, + C: Fn() -> bool, +{ + let mut future = Box::pin(future); + let thread_waker = Arc::new(BlockingThreadWaker(thread::current())); + let waker = waker_ref(&thread_waker); + let mut context = Context::from_waker(&waker); + loop { + if is_cancelled() { + return None; + } + match future.as_mut().poll(&mut context) { + Poll::Ready(output) => return Some(output), + Poll::Pending => thread::park_timeout(CANCELLATION_POLL_DELAY), + } + } +} + +// 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 { + continuation: Arc, + 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.continuation.host.v3.async_stream_release)(self.output) }; + } +} + +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::(); + 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); + } + let continuation = unsafe { LlmContinuationV2::from_raw(host, next) }?; + Ok(Self { + inner: Arc::new(StreamContinuationInner { + continuation: continuation.inner, + output, + }), + }) + } + + /// Opens one explicitly targeted provider stream. + pub async fn open_stream( + &self, + invocation: LlmContinuationInvocationV2, + ) -> std::result::Result { + let host = self.inner.continuation.host; + let (sender, receiver) = oneshot::channel(); + let state = Box::new(StreamOpenCallback { host, sender }); + let state = Box::into_raw(state).cast::(); + let status = match HostString::from_json(&host.v3.v1, &invocation) { + Some(invocation) => unsafe { + (host.async_llm_next_open_stream_v2)( + self.inner.continuation.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)); + } + let raw = receiver.await.unwrap_or_else(|_| { + Err(internal_failure( + "targeted LLM stream setup callback closed without an outcome", + )) + })? as *const NemoRelayNativeLlmStreamV2; + if raw.is_null() { + return Err(internal_failure( + "targeted LLM stream setup returned a null stream", + )); + } + Ok(LlmProviderStreamV2 { + host, + raw, + pending: None, + finished: false, + terminal: false, + }) + } + + async fn forward_passthrough(&self, request: LlmRequest) -> Result<()> { + let host = self.inner.continuation.host; + let (sender, receiver) = oneshot::channel(); + let state = Box::new(ForwardStreamCallback { + host: host.v3.v1, + sender, + }); + let state = Box::into_raw(state).cast::(); + let status = match HostString::from_json(&host.v3.v1, &request) { + Some(request) => unsafe { + (host.async_llm_next_forward_stream_v2)( + self.inner.continuation.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.unwrap_or_else(|_| { + 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::(); + 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(|_| ProviderItem { + value: Err(internal_failure( + "provider stream callback closed without an event", + )), + terminal: false, + }); + self.terminal = item.terminal; + match item.value { + 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) { + if !self.terminal { + let _ = unsafe { (self.host.async_llm_stream_cancel_v2)(self.raw) }; + } + unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; + self.raw = ptr::null(); + } +} + +#[derive(Deserialize)] +struct LlmCallbackInvocation { + name: String, + request: LlmRequest, +} + +struct SafeV2Callback { + host: NemoRelayNativeHostApiV4, + callback: F, +} + +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. + /// Its future is driven by a cancellation-aware local executor on Relay's + /// reusable blocking callback lane. It is not polled inside an entered + /// Tokio runtime; runtime-specific APIs should not be assumed. + 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 }); + 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.plugin_context_register_async_llm_execution_v2)( + self.raw, + name.as_ptr(), + priority, + 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. Its future and returned stream are driven on Relay's + /// reusable blocking callback lane, outside an entered Tokio runtime. + 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 }); + let user_data = Box::into_raw(state).cast::(); + let status = unsafe { + // The V4 host has the same consume-on-invocation ownership + // contract as buffered registration. + (host.plugin_context_register_async_llm_stream_execution_v2)( + 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 result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let invocation: LlmCallbackInvocation = read_json_value( + &state.host.v3.v1, + invocation_json, + "native API v2 LLM invocation", + ) + .map_err(|status| format!("invalid native API v2 LLM invocation: {status:?}"))?; + let execution = (state.callback)(invocation.name, invocation.request, continuation.clone()); + block_on_complete_callback(execution, || unsafe { + (state.host.v3.async_completion_is_cancelled)(completion) + }) + .unwrap_or_else(|| Err("native API v2 callback was cancelled".into())) + })); + match result { + Ok(Ok(value)) => resolve_completion(&state.host, completion, &value), + Ok(Err(error)) => reject_completion(&state.host, completion, &error), + Err(_) => reject_completion( + &state.host, + completion, + "native API v2 buffered LLM callback panicked", + ), + } + drop(continuation); + 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( + &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 result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let invocation: LlmCallbackInvocation = read_json_value( + &state.host.v3.v1, + invocation_json, + "native API v2 streaming LLM invocation", + ) + .map_err(|status| format!("invalid native API v2 stream invocation: {status:?}"))?; + let execution = async { + let outcome = + (state.callback)(invocation.name, invocation.request, continuation.clone()).await?; + match outcome { + LlmStreamExecutionOutcomeV2::Stream(stream) => { + pump_output_stream(&continuation, stream).await + } + LlmStreamExecutionOutcomeV2::Passthrough(request) => { + continuation.forward_passthrough(request).await + } + } + }; + block_on_complete_callback(execution, || unsafe { + (state.host.v3.async_stream_is_cancelled)(continuation.inner.output) + }) + .unwrap_or_else(|| Err("native API v2 output stream was cancelled".into())) + })); + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => reject_output(&state.host, output, &error), + Err(_) => reject_output( + &state.host, + output, + "native API v2 streaming LLM callback panicked", + ), + } + drop(continuation); + 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)?; + } + finish_output(continuation) +} + +fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { + let host = &continuation.inner.continuation.host; + let chunk = HostString::from_json(&host.v3.v1, chunk) + .ok_or_else(|| "failed to serialize native API v2 output chunk".to_string())?; + loop { + 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_push_json)(continuation.inner.output, chunk.as_ptr()) }; + match status { + NemoRelayStatus::Ok => return Ok(()), + // For this operation the V3 ABI contract reserves `Internal` for + // a full bounded queue. Serialization and lifecycle faults use + // distinct statuses, so retrying cannot mask another ABI error. + NemoRelayStatus::Internal => thread::sleep(BACKPRESSURE_RETRY_DELAY), + status => return Err(format!("native API v2 output push failed: {status:?}")), + } + } +} + +fn finish_output(continuation: &LlmStreamContinuationV2) -> Result<()> { + let host = &continuation.inner.continuation.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:?}")) + } +} + +fn reject_output( + 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; + }; + loop { + if unsafe { (host.v3.async_stream_is_cancelled)(output) } { + return; + } + let status = unsafe { (host.v3.async_stream_reject)(output, message.as_ptr()) }; + match status { + NemoRelayStatus::Ok => return, + // `Internal` has the same queue-full-only contract for rejection. + NemoRelayStatus::Internal => thread::sleep(BACKPRESSURE_RETRY_DELAY), + status => { + set_last_error( + &host.v3.v1, + &format!("native API v2 output rejection failed: {status:?}"), + ); + return; + } + } + } +} + +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, + outcome_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(|| { + let outcome: LlmContinuationOutcomeV2 = read_json_value( + &state.host, + outcome_json, + "targeted LLM continuation outcome", + ) + .map_err(|status| status_failure("targeted LLM continuation outcome", status))?; + match outcome { + LlmContinuationOutcomeV2::Success { response } => Ok(response), + LlmContinuationOutcomeV2::Failure { error } => Err(error), + } + })) + .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>, +} + +struct OwnedProviderStream { + host: NemoRelayNativeHostApiV4, + raw: *const NemoRelayNativeLlmStreamV2, + armed: bool, +} + +impl OwnedProviderStream { + fn new(host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeLlmStreamV2) -> Self { + Self { + host, + raw, + armed: !raw.is_null(), + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for OwnedProviderStream { + fn drop(&mut self) { + if !self.armed { + return; + } + let _ = unsafe { (self.host.async_llm_stream_cancel_v2)(self.raw) }; + unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; + } +} + +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 owned_stream = OwnedProviderStream::new(state.host, stream); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match (stream.is_null(), error_json.is_null()) { + (false, true) => Ok(stream as usize), + (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 transfers_stream = result.is_ok(); + if state.sender.send(result).is_ok() && transfers_stream { + owned_stream.disarm(); + } +} + +struct ProviderNextCallback { + host: NemoRelayNativeHostApiV1, + sender: oneshot::Sender, +} + +unsafe extern "C" fn provider_next_callback( + user_data: *mut c_void, + event_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(|| { + let event: LlmContinuationStreamEventV2 = + read_json_value(&state.host, event_json, "targeted provider stream event") + .map_err(|status| status_failure("targeted provider stream event", status))?; + Ok::<_, LlmContinuationFailureV2>(match event { + LlmContinuationStreamEventV2::Chunk { chunk } => ProviderItem { + value: Ok(Some(chunk)), + terminal: false, + }, + LlmContinuationStreamEventV2::Done => ProviderItem { + value: Ok(None), + terminal: true, + }, + LlmContinuationStreamEventV2::Failure { error } => ProviderItem { + value: Err(error), + terminal: true, + }, + }) + })) + .unwrap_or_else(|_| { + Err(internal_failure( + "targeted provider stream callback panicked", + )) + }) + .unwrap_or_else(|error| ProviderItem { + value: Err(error), + terminal: false, + }); + let _ = state.sender.send(result); +} + +struct ForwardStreamCallback { + host: NemoRelayNativeHostApiV1, + sender: oneshot::Sender>, +} + +unsafe extern "C" fn forward_stream_callback( + user_data: *mut c_void, + 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() { + Ok(()) + } else { + // Relay has already settled the output before this wake-up. Keep + // the terminal context available for diagnostics, but report + // completion to the trampoline so it cannot reject a second time. + let message = + read_required_host_string(&state.host, error, "streaming pass-through error") + .unwrap_or_else(|status| { + format!("invalid streaming pass-through error: {status:?}") + }); + set_last_error(&state.host, &message); + Ok(()) + } + })) + .unwrap_or_else(|_| Err("streaming pass-through terminal callback panicked".into())); + let _ = state.sender.send(result); +} + +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 { + failure: LlmNonHttpFailureV2 { + 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 575512bf4..362774ab9 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -9,29 +9,36 @@ 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, LlmContinuationFailureV2, - LlmHttpFailureV2, LlmJsonStream, LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, - LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, - NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, - NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareKind, - 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, + LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationRouteV2, + LlmContinuationStreamEventV2, LlmContinuationTargetV2, LlmHttpFailureV2, LlmJsonStream, + LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, LlmRequestInterceptOutcome, + LlmStream, 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, NemoRelayNativeEventSanitizeCb, + NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, + NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, + NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, + NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, + NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, + NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativeLlmStreamV2, + NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, + NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, + NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, + NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, PendingMarkSpec, + PluginContext, PluginRuntime, ScopeType, ToolExecutionInterceptOutcome, ToolNext, }; use serde_json::{Map, json}; @@ -300,6 +307,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 { @@ -336,6 +375,8 @@ impl_captured_registration!( RegisteredLlmExecution, RegisteredLlmStreamExecution, RegisteredLlmRequestIntercept, + RegisteredAsyncV2, + RegisteredAsyncStreamV2, ); fn replace_registration(slot: &Mutex>, registration: T) { @@ -395,6 +436,26 @@ 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()); +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_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_OUTPUT_RELEASES: AtomicUsize = AtomicUsize::new(0); +static SAFE_V2_PROVIDER_CANCELS: 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); #[test] fn native_abi_v3_struct_sizes_are_self_describing() { @@ -435,10 +496,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 496); + assert_eq!(size_of::(), 504); assert_eq!( host_api_v4_offsets(), - [0, 440, 448, 456, 464, 472, 480, 488] + [0, 440, 448, 456, 464, 472, 480, 488, 496] ); assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 56); @@ -469,10 +530,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 244); + assert_eq!(size_of::(), 248); assert_eq!( host_api_v4_offsets(), - [0, 216, 220, 224, 228, 232, 236, 240] + [0, 216, 220, 224, 228, 232, 236, 240, 244] ); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); @@ -483,7 +544,7 @@ fn native_abi_v3_struct_sizes_are_self_describing() { } } -fn host_api_v4_offsets() -> [usize; 8] { +fn host_api_v4_offsets() -> [usize; 9] { [ offset_of!(NemoRelayNativeHostApiV4, v3), offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_result_v2), @@ -499,6 +560,7 @@ fn host_api_v4_offsets() -> [usize; 8] { NemoRelayNativeHostApiV4, plugin_context_register_async_llm_stream_execution_v2 ), + offset_of!(NemoRelayNativeHostApiV4, async_llm_next_forward_stream_v2), ] } @@ -1421,6 +1483,8 @@ 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_eq!( STRING_LIVE_COUNT.load(Ordering::SeqCst), 0, @@ -1448,6 +1512,22 @@ 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_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_OUTPUT_RELEASES.store(0, Ordering::SeqCst); + SAFE_V2_PROVIDER_CANCELS.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; } fn test_context(host: &NemoRelayNativeHostApiV1) -> PluginContext<'_> { @@ -6057,3 +6137,992 @@ 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_string(&host, value_json) + .ok() + .and_then(|value| serde_json::from_str(&value).ok()) + { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + *SAFE_V2_COMPLETION.lock().unwrap() = Some(Ok(value)); + NemoRelayStatus::Ok +} + +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)); + NemoRelayStatus::Ok +} + +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, +) { +} + +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 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 { + NemoRelayStatus::InvalidArg +} + +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_string(&host, chunk_json) + .ok() + .and_then(|value| serde_json::from_str(&value).ok()) + { + Some(chunk) => chunk, + None => return NemoRelayStatus::InvalidJson, + }; + SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_stream_finish( + _stream: *const NemoRelayNativeAsyncStream, +) -> NemoRelayStatus { + SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); + NemoRelayStatus::Ok +} + +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:?}")); + SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); + NemoRelayStatus::Ok +} + +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 { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn safe_v2_passthrough_result( + _next: *const NemoRelayNativeAsyncNext, + invocation_json: *const NemoRelayNativeString, + cb: NemoRelayNativeAsyncNextResultCb, + user_data: *mut c_void, +) -> NemoRelayStatus { + let host = test_host(); + if required_host_string(&host, invocation_json) + .ok() + .and_then(|value| serde_json::from_str::(&value).ok()) + .is_none() + { + return NemoRelayStatus::InvalidJson; + } + 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 host = test_host(); + if required_host_string(&host, invocation_json) + .ok() + .and_then(|value| serde_json::from_str::(&value).ok()) + .is_none() + { + return NemoRelayStatus::InvalidJson; + } + 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; + } + let outcome = serde_json::to_value(LlmContinuationOutcomeV2::Success { + response: json!({ "targeted": true }), + }) + .unwrap(); + let outcome = json_host_string(&host, outcome); + unsafe { cb(user_data, outcome) }; + unsafe { (host.string_free)(outcome) }; + 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 host = test_host(); + if required_host_string(&host, invocation_json) + .ok() + .and_then(|value| serde_json::from_str::(&value).ok()) + .is_none() + { + return NemoRelayStatus::InvalidJson; + } + 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 host = test_host(); + let event = SAFE_V2_PROVIDER_EVENTS + .lock() + .unwrap() + .pop_front() + .unwrap_or(LlmContinuationStreamEventV2::Done); + let event = json_host_string(&host, serde_json::to_value(event).unwrap()); + unsafe { cb(user_data, event) }; + unsafe { (host.string_free)(event) }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn safe_v2_provider_cancel( + _stream: *const NemoRelayNativeLlmStreamV2, +) -> NemoRelayStatus { + SAFE_V2_PROVIDER_CANCELS.fetch_add(1, Ordering::SeqCst); + 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_register_buffered( + _ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => return status, + }; + 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_register_streaming( + _ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => return status, + }; + 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_forward_stream( + _next: *const NemoRelayNativeAsyncNext, + request_json: *const NemoRelayNativeString, + _output_stream: *const NemoRelayNativeAsyncStream, + cb: NemoRelayNativeAsyncLlmStreamForwardCbV2, + user_data: *mut c_void, +) -> NemoRelayStatus { + let host = test_host(); + let request = match required_host_string(&host, request_json) + .ok() + .and_then(|value| serde_json::from_str::(&value).ok()) + { + Some(request) => request, + None => return NemoRelayStatus::InvalidJson, + }; + SAFE_V2_FORWARDED_REQUESTS.lock().unwrap().push(request); + SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); + unsafe { cb(user_data, ptr::null()) }; + 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_cancel_v2: safe_v2_provider_cancel, + async_llm_stream_release_v2: safe_v2_provider_release, + plugin_context_register_async_llm_execution_v2: safe_v2_register_buffered, + plugin_context_register_async_llm_stream_execution_v2: safe_v2_register_streaming, + async_llm_next_forward_stream_v2: safe_v2_forward_stream, + } +} + +fn safe_v2_target() -> LlmContinuationTargetV2 { + LlmContinuationTargetV2 { + method: "POST".into(), + url: "https://provider.example/v1/chat/completions".into(), + route: LlmContinuationRouteV2::OpenaiChat, + headers: Default::default(), + } +} + +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() +} + +#[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::Complete) + ); + 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); + 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() }), + ) as usize; + let user_data = registration.user_data; + let callback = registration.cb; + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let callback_thread = std::thread::spawn(move || { + let state = unsafe { + callback( + user_data as *mut c_void, + invocation as *const NemoRelayNativeString, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + done_tx.send(state).unwrap(); + }); + + std::thread::sleep(std::time::Duration::from_millis(20)); + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + let state = done_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("caller cancellation must wake a pending safe buffered callback"); + callback_thread.join().unwrap(); + unsafe { (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString) }; + + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Complete) + ); + assert_eq!(SAFE_V2_NEXT_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_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() }), + ) as usize; + let user_data = registration.user_data; + let callback = registration.cb; + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let callback_thread = std::thread::spawn(move || { + let state = unsafe { + callback( + user_data as *mut c_void, + invocation as *const NemoRelayNativeString, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + done_tx.send(state).unwrap(); + }); + while SAFE_V2_HELD_TARGETED_CALLBACK.lock().unwrap().is_none() { + std::thread::yield_now(); + } + SAFE_V2_COMPLETION_CANCELLED.store(true, Ordering::SeqCst); + let state = done_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("cancelling the managed call must stop the targeted callback"); + callback_thread.join().unwrap(); + + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Complete) + ); + 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 outcome = json_host_string( + &host.v3.v1, + serde_json::to_value(LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Cancelled, + message: "cancelled by host".into(), + }, + }, + }) + .unwrap(), + ); + unsafe { targeted_callback(targeted_user_data as *mut c_void, outcome) }; + unsafe { + (host.v3.v1.string_free)(outcome); + (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString); + 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([ + LlmContinuationStreamEventV2::Chunk { + chunk: json!({ "delta": "hello" }), + }, + LlmContinuationStreamEventV2::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 provider = next + .open_stream(LlmContinuationInvocationV2 { + request, + target: safe_v2_target(), + }) + .await + .map_err(|error| format!("{error:?}"))?; + let output = provider.map(|item| item.map_err(|error| format!("{error:?}"))); + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(output))) + }, + ) + .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::Complete) + ); + 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_CANCELS.load(Ordering::SeqCst), 0); + 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::Complete) + ); + 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_cancels_unfinished_production() { + let _guard = begin_test(); + SAFE_V2_PROVIDER_EVENTS + .lock() + .unwrap() + .push_back(LlmContinuationStreamEventV2::Chunk { + 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) }; + assert_eq!(SAFE_V2_PROVIDER_CANCELS.load(Ordering::SeqCst), 1); + assert_eq!(SAFE_V2_PROVIDER_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() }), + ) as usize; + let user_data = registration.user_data; + let callback = registration.cb; + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let callback_thread = std::thread::spawn(move || { + let state = unsafe { + callback( + user_data as *mut c_void, + invocation as *const NemoRelayNativeString, + NonNull::::dangling().as_ptr(), + NonNull::::dangling().as_ptr(), + ) + }; + done_tx.send(state).unwrap(); + }); + + std::thread::sleep(std::time::Duration::from_millis(20)); + SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); + let state = done_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("caller cancellation must wake a pending safe stream callback"); + callback_thread.join().unwrap(); + unsafe { (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString) }; + + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Complete) + ); + 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 { + failure: LlmHttpFailureV2 { + 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); + assert!(error.is_retryable()); + 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) }; + 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_cancels_and_releases_the_owned_stream() { + let _guard = begin_test(); + *SAFE_V2_OPEN_FAILURE.lock().unwrap() = Some(LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + 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) }; + + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("invalid outcome") + )); + assert_eq!(SAFE_V2_PROVIDER_CANCELS.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_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_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); +} From 2025301ad0bf744f49030ce4ebe521ab6f32e817 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 00:40:23 -0600 Subject: [PATCH 09/32] docs(plugin): document safe native API v2 facade Signed-off-by: Bryan Bednarski --- crates/plugin/README.md | 73 +++++++++++++++---- docs/build-plugins/dynamic-plugins/about.mdx | 4 +- .../dynamic-plugins/native-dynamic/about.mdx | 62 +++++++++++++++- docs/reference/migration-guides.mdx | 12 +-- 4 files changed, 128 insertions(+), 23 deletions(-) diff --git a/crates/plugin/README.md b/crates/plugin/README.md index b52da04e2..e5207d895 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -42,23 +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. -- **Native API v2 targeted LLM continuations**: Register v2-only execution - callbacks that send an explicit provider target through Relay and receive - buffered JSON, structured failures, or a bounded host-owned provider stream. +- **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: @@ -110,9 +111,47 @@ nemo_relay_plugin::nemo_relay_plugin_v2!( ); ``` -Set `compat.native_api = "2"` in `relay-plugin.toml`. During registration, -`PluginContext::host_api_v4` exposes the C-safe targeted LLM continuation table -and the raw v2 registration helpers. +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))) + } +})?; +``` + +For an unmanaged buffered request, call +`LlmContinuationV2::call_passthrough`. For an unmanaged streaming request, +return `LlmStreamExecutionOutcomeV2::Passthrough(request)`. Relay then pumps +the original downstream stream directly through its bounded queue; provider +events do not cross into the plugin merely to be forwarded. The plugin provides JSON plus an HTTP method, absolute target URL, protocol route, and explicit target headers. Relay binds that transport target to the @@ -126,10 +165,16 @@ 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 an opaque host-owned stream; request one JSON event -at a time, then cancel and release it exactly once. No Rust future, trait -object, `serde_json::Value`, or allocator-owned Rust string crosses the ABI -boundary. +Streaming dispatch returns `LlmProviderStreamV2`, which implements Rust +`Stream` and cancels unfinished provider work on drop. Safe callback futures +and returned streams run to completion on Relay's reusable blocking callback +lane, outside an entered Tokio runtime. No Rust future, trait object, +`serde_json::Value`, or allocator-owned Rust string crosses the C ABI boundary. + +The raw `PluginContext::host_api_v4` table and `_raw` v2 registration methods +remain available for advanced ABI consumers and non-Rust bindings. Code using +them is responsible for every callback lifetime, host string, completion, +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. diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 772a0b827..0df4516a1 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, or `"2"` for targeted LLM continuations | `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 2fd946b80..103220916 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -174,6 +174,57 @@ restored. Continuations run on Relay's existing Tokio runtime; Relay does not create an OS thread per continuation. 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: + +```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` for an unmanaged buffered call. For +an unmanaged stream, return +`LlmStreamExecutionOutcomeV2::Passthrough(request)`. Relay connects the +original downstream stream directly to the caller through its bounded queue, +so pass-through events do not cross the plugin boundary. + +The SDK drives safe callback futures and returned streams to completion on the +blocking callback lane. It does not enter Relay's Tokio runtime on that thread. +`LlmProviderStreamV2` implements `Stream`, enforces one pending pull, and +cancels unfinished provider production on drop. + Export a v2-only plugin with: ```rust @@ -184,8 +235,7 @@ nemo_relay_plugin::nemo_relay_plugin_v2!( ``` Set `compat.native_api = "2"` in its manifest. Use -`PluginContext::host_api_v4` and the raw v2 LLM registration methods when the -plugin needs the typed dispatch contract. +the safe registration methods above for normal Rust plugins. ## Native C ABI Table Versioning @@ -203,6 +253,14 @@ extern "C" fn nemo_relay_register_plugin( ) -> NemoRelayStatus ``` +`PluginContext::host_api_v4` and the `_raw` v2 registration methods are the +advanced escape hatch for raw ABI consumers and non-Rust bindings. They expose +opaque callbacks and handles intentionally: the caller must own host strings, +completion settlement, stream backpressure, cancellation, panic fencing, and +every release operation. The safe Rust facade is implemented on top of these +same C functions; 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 diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index ed76d4bb1..1dead7a9e 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -335,7 +335,8 @@ 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. @@ -347,9 +348,10 @@ 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. Request and response callbacks now receive distinct context structures. Each structure contains structured codec identity and a borrowed directional codec @@ -362,7 +364,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 From 43589eab584c79bc3ba4aa52d4a53ffa8a724a2e Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 00:40:32 -0600 Subject: [PATCH 10/32] fix(runtime): use registered targeted dispatch log target Signed-off-by: Bryan Bednarski --- crates/core/src/api/runtime/llm_dispatch_context.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index cb389f6b2..1f16f216c 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -327,7 +327,7 @@ fn transport_error(target: &LlmDispatchTargetContext, error: reqwest::Error) -> let timeout = error.is_timeout(); let diagnostic = error.without_url(); log::warn!( - target: "nemo_relay.llm", + target: "nemo_relay.runtime", event = "targeted_llm_transport_failed", provider_host = target.url().host_str().unwrap_or(""), failure_kind = if timeout { "timeout" } else { "transport" }; From ec2639060776766cd1fd6641810160af4232c8dc Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 01:12:33 -0600 Subject: [PATCH 11/32] fix(plugin): avoid safe callback thread-local teardown Signed-off-by: Bryan Bednarski --- .../tests/integration/native_plugin_tests.rs | 16 ++++++- crates/plugin/src/native_v2.rs | 42 +++++++++++++++---- 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index f156f3a9f..9fcc270cc 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -1545,6 +1545,8 @@ async fn native_api_v2_safe_fixture_forwards_passthrough_inside_embedded_core() 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") @@ -1554,7 +1556,12 @@ async fn native_api_v2_safe_fixture_forwards_passthrough_inside_embedded_core() }) .func(Arc::new(move |_| { let dropped = provider_dropped_for_fn.clone(); - Box::pin(async move { Ok(LlmJsonStream::new(PendingDropStream { dropped })) }) + 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)) @@ -1562,6 +1569,13 @@ async fn native_api_v2_safe_fixture_forwards_passthrough_inside_embedded_core() ) .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) { diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index a4a6c2569..1b0eaf16f 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -7,7 +7,7 @@ use std::ffi::c_void; use std::future::Future; use std::pin::Pin; use std::ptr; -use std::sync::Arc; +use std::sync::{Arc, Condvar, Mutex}; use std::task::{Context, Poll}; use std::thread; use std::time::Duration; @@ -86,11 +86,36 @@ struct ContinuationInner { next: *const NemoRelayNativeAsyncNext, } -struct BlockingThreadWaker(thread::Thread); +// Do not use `thread::current()` here. Safe callbacks run on reusable host +// threads, and Rust's thread handle installs a TLS destructor in this plugin +// library. Relay may unload the library before that host thread exits. +struct BlockingCallbackWaker { + notified: Mutex, + ready: Condvar, +} -impl ArcWake for BlockingThreadWaker { +impl BlockingCallbackWaker { + fn wait_timeout(&self, timeout: Duration) { + let notified = self + .notified + .lock() + .unwrap_or_else(|error| error.into_inner()); + let (mut notified, _) = self + .ready + .wait_timeout_while(notified, timeout, |notified| !*notified) + .unwrap_or_else(|error| error.into_inner()); + *notified = false; + } +} + +impl ArcWake for BlockingCallbackWaker { fn wake_by_ref(arc_self: &Arc) { - arc_self.0.unpark(); + let mut notified = arc_self + .notified + .lock() + .unwrap_or_else(|error| error.into_inner()); + *notified = true; + arc_self.ready.notify_one(); } } @@ -100,8 +125,11 @@ where C: Fn() -> bool, { let mut future = Box::pin(future); - let thread_waker = Arc::new(BlockingThreadWaker(thread::current())); - let waker = waker_ref(&thread_waker); + let callback_waker = Arc::new(BlockingCallbackWaker { + notified: Mutex::new(false), + ready: Condvar::new(), + }); + let waker = waker_ref(&callback_waker); let mut context = Context::from_waker(&waker); loop { if is_cancelled() { @@ -109,7 +137,7 @@ where } match future.as_mut().poll(&mut context) { Poll::Ready(output) => return Some(output), - Poll::Pending => thread::park_timeout(CANCELLATION_POLL_DELAY), + Poll::Pending => callback_waker.wait_timeout(CANCELLATION_POLL_DELAY), } } } From 1821bc41cce7863f4cd7ba68c6d1dc4d177a5dd3 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 01:57:26 -0600 Subject: [PATCH 12/32] test(plugin): cover safe v2 failure paths Signed-off-by: Bryan Bednarski --- crates/plugin/tests/typed_callbacks.rs | 792 ++++++++++++++++++++++++- 1 file changed, 775 insertions(+), 17 deletions(-) diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 362774ab9..7b1b40d3d 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -5,6 +5,7 @@ use std::collections::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::{ @@ -17,9 +18,10 @@ use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationRouteV2, - LlmContinuationStreamEventV2, LlmContinuationTargetV2, LlmHttpFailureV2, LlmJsonStream, - LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, LlmRequestInterceptOutcome, - LlmStream, LlmStreamExecutionOutcomeV2, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, + LlmContinuationStreamEventV2, LlmContinuationTargetV2, LlmContinuationV2, LlmHttpFailureV2, + LlmJsonStream, LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, + LlmRequestInterceptOutcome, LlmStream, LlmStreamContinuationV2, LlmStreamExecutionOutcomeV2, + LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, NemoRelayNativeAsyncLlmStreamForwardCbV2, @@ -456,6 +458,20 @@ 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_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_FORWARD_STATUS: Mutex = Mutex::new(NemoRelayStatus::Ok); +static SAFE_V2_FORWARD_ERROR: Mutex> = Mutex::new(None); +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()); #[test] fn native_abi_v3_struct_sizes_are_self_describing() { @@ -1528,6 +1544,19 @@ fn reset_state() { 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_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_FORWARD_STATUS.lock().unwrap() = NemoRelayStatus::Ok; + *SAFE_V2_FORWARD_ERROR.lock().unwrap() = None; + *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(); } fn test_context(host: &NemoRelayNativeHostApiV1) -> PluginContext<'_> { @@ -6151,7 +6180,7 @@ unsafe extern "C" fn safe_v2_completion_resolve( None => return NemoRelayStatus::InvalidJson, }; *SAFE_V2_COMPLETION.lock().unwrap() = Some(Ok(value)); - NemoRelayStatus::Ok + *SAFE_V2_COMPLETION_RESOLVE_STATUS.lock().unwrap() } unsafe extern "C" fn safe_v2_completion_reject( @@ -6162,7 +6191,7 @@ unsafe extern "C" fn safe_v2_completion_reject( let message = required_host_string(&host, message) .unwrap_or_else(|status| format!("invalid rejection: {status:?}")); *SAFE_V2_COMPLETION.lock().unwrap() = Some(Err(message)); - NemoRelayStatus::Ok + *SAFE_V2_COMPLETION_REJECT_STATUS.lock().unwrap() } unsafe extern "C" fn safe_v2_completion_is_cancelled( @@ -6213,15 +6242,22 @@ unsafe extern "C" fn safe_v2_stream_push( Some(chunk) => chunk, None => return NemoRelayStatus::InvalidJson, }; - SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); - NemoRelayStatus::Ok + let status = SAFE_V2_OUTPUT_PUSH_STATUSES + .lock() + .unwrap() + .pop_front() + .unwrap_or(NemoRelayStatus::Ok); + if status == NemoRelayStatus::Ok { + SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); + } + status } unsafe extern "C" fn safe_v2_stream_finish( _stream: *const NemoRelayNativeAsyncStream, ) -> NemoRelayStatus { SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); - NemoRelayStatus::Ok + *SAFE_V2_OUTPUT_FINISH_STATUS.lock().unwrap() } unsafe extern "C" fn safe_v2_stream_reject( @@ -6231,8 +6267,15 @@ unsafe extern "C" fn safe_v2_stream_reject( let host = test_host(); let message = required_host_string(&host, message) .unwrap_or_else(|status| format!("invalid rejection: {status:?}")); - SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); - NemoRelayStatus::Ok + let status = SAFE_V2_OUTPUT_REJECT_STATUSES + .lock() + .unwrap() + .pop_front() + .unwrap_or(NemoRelayStatus::Ok); + if status == NemoRelayStatus::Ok { + SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); + } + status } unsafe extern "C" fn safe_v2_stream_is_cancelled( @@ -6272,6 +6315,10 @@ unsafe extern "C" fn safe_v2_passthrough_result( 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 required_host_string(&host, invocation_json) .ok() @@ -6280,9 +6327,15 @@ unsafe extern "C" fn safe_v2_passthrough_result( { return NemoRelayStatus::InvalidJson; } - let value = json_host_string(&host, json!({ "passthrough": true })); - unsafe { cb(user_data, value, ptr::null()) }; - unsafe { (host.string_free)(value) }; + 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 } @@ -6292,6 +6345,10 @@ unsafe extern "C" fn safe_v2_targeted_result( 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 required_host_string(&host, invocation_json) .ok() @@ -6321,6 +6378,10 @@ unsafe extern "C" fn safe_v2_stream_open( 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 required_host_string(&host, invocation_json) .ok() @@ -6355,7 +6416,17 @@ unsafe extern "C" fn safe_v2_provider_next( 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 event = host_string(&host, &event); + unsafe { cb(user_data, event) }; + unsafe { (host.string_free)(event) }; + return NemoRelayStatus::Ok; + } let event = SAFE_V2_PROVIDER_EVENTS .lock() .unwrap() @@ -6451,6 +6522,10 @@ unsafe extern "C" fn safe_v2_forward_stream( 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_string(&host, request_json) .ok() @@ -6461,7 +6536,13 @@ unsafe extern "C" fn safe_v2_forward_stream( }; SAFE_V2_FORWARDED_REQUESTS.lock().unwrap().push(request); SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); - unsafe { cb(user_data, ptr::null()) }; + if let Some(error) = SAFE_V2_FORWARD_ERROR.lock().unwrap().clone() { + let error = host_string(&host, &error); + unsafe { cb(user_data, error) }; + unsafe { (host.string_free)(error) }; + } else { + unsafe { cb(user_data, ptr::null()) }; + } NemoRelayStatus::Ok } @@ -6516,6 +6597,205 @@ 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() }; + 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() }; + 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:?}"))), + ))) +} + +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_raw_registration_remains_an_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_llm_execution_v2_raw( + "raw-buffered", + 11, + 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_llm_stream_execution_v2_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_llm_execution_v2_raw( + "unsupported-buffered", + 0, + raw_v2_buffered_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_eq!( + unsafe { + v1_ctx.register_async_llm_stream_execution_v2_raw( + "unsupported-streaming", + 0, + raw_v2_streaming_probe, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::InvalidArg + ); + + assert_eq!(LlmContinuationRouteV2::OpenaiChat.as_str(), "openai_chat"); + assert_eq!( + LlmContinuationRouteV2::OpenaiResponses.as_str(), + "openai_responses" + ); + assert_eq!( + LlmContinuationRouteV2::AnthropicMessages.as_str(), + "anthropic_messages" + ); + assert_eq!(live_host_strings(), 0); +} + #[test] fn safe_v2_buffered_registration_wraps_targeted_and_passthrough_calls() { let _guard = begin_test(); @@ -6707,15 +6987,24 @@ fn safe_v2_stream_registration_pumps_provider_stream_and_releases_handles() { "safe-stream", 3, |_name, request, next| async move { - let provider = next + let mut provider = next .open_stream(LlmContinuationInvocationV2 { request, target: safe_v2_target(), }) .await .map_err(|error| format!("{error:?}"))?; - let output = provider.map(|item| item.map_err(|error| format!("{error:?}"))); - Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(output))) + 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(); @@ -7126,3 +7415,472 @@ fn safe_v2_malformed_invocations_settle_and_release_callback_handles() { unsafe { streaming.free() }; assert_eq!(live_host_strings(), 0); } + +#[test] +fn safe_v2_executor_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_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_streaming(&host, "open-status", |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); + 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_streaming(&host, "poll-status", |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); + 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(LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Transport, + message: "late provider failure".into(), + }, + }, + }); + run_safe_v2_streaming(&host, "late-failure", |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); + 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_streaming(&host, "malformed-event", |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("stream 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::Internal, NemoRelayStatus::Ok]); + run_safe_v2_streaming(&host, "push-backpressure", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( + async { Ok(json!({ "chunk": true })) }, + )))) + }); + assert_eq!( + *SAFE_V2_OUTPUT.lock().unwrap(), + vec![Ok(json!({ "chunk": true }))] + ); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + + SAFE_V2_OUTPUT_PUSH_STATUSES + .lock() + .unwrap() + .push_back(NemoRelayStatus::InvalidArg); + run_safe_v2_streaming(&host, "push-failure", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( + async { Ok(json!({ "chunk": true })) }, + )))) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("output push failed") + )); + 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::Internal, 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(); + + *SAFE_V2_FORWARD_ERROR.lock().unwrap() = Some("downstream stream failed".into()); + *LAST_ERROR.lock().unwrap() = None; + run_safe_v2_streaming(&host, "forward-error", |_, request, _| async move { + Ok(LlmStreamExecutionOutcomeV2::Passthrough(request)) + }); + assert_eq!( + LAST_ERROR.lock().unwrap().as_deref(), + Some("downstream stream failed") + ); + assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); + + 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); +} From 007b19b72862900b3f934b64476a43bbaf9f7ce9 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 03:28:08 -0600 Subject: [PATCH 13/32] test(plugin): harden native v2 host failure paths Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 16 +- .../tests/unit/llm_dispatch_context_tests.rs | 84 + crates/core/tests/unit/native_plugin_tests.rs | 1822 ++++++++++++++++- 3 files changed, 1867 insertions(+), 55 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index dfd7de46d..f533886fc 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -3619,6 +3619,13 @@ fn wrap_native_async_llm_execution_v2( free_fn: NemoRelayNativeFreeFn, ) -> LlmExecutionFn { let user_data = make_user_data(instance, user_data, free_fn); + wrap_native_async_llm_execution_v2_with_user_data(cb, user_data) +} + +fn wrap_native_async_llm_execution_v2_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(); @@ -3729,6 +3736,13 @@ fn wrap_native_incremental_llm_stream_execution_v2( free_fn: NemoRelayNativeFreeFn, ) -> LlmStreamExecutionFn { let user_data = make_user_data(instance, user_data, free_fn); + wrap_native_incremental_llm_stream_execution_v2_with_user_data(cb, user_data) +} + +fn wrap_native_incremental_llm_stream_execution_v2_with_user_data( + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: Arc, +) -> LlmStreamExecutionFn { Arc::new(move |name, request, next| { let user_data = user_data.clone(); let name = name.to_owned(); @@ -3756,6 +3770,7 @@ fn wrap_native_incremental_llm_stream_execution_v2( "failed to allocate native API v2 stream invocation".into(), ) })? as usize; + let invocation_guard = NativeInvocationStringGuard(invocation); let runtime = tokio::runtime::Handle::try_current().map_err(|error| { FlowError::Internal(format!( "native API v2 stream intercept requires a Tokio runtime: {error}" @@ -3769,7 +3784,6 @@ fn wrap_native_incremental_llm_stream_execution_v2( let stream_ref = Arc::into_raw(stream.clone()) as usize; let callback_user_data = user_data.ptr as usize; let callback_scope_stack = current_scope_stack(); - let invocation_guard = NativeInvocationStringGuard(invocation); let next_handoff = NativeNextHandoff { raw: Some(next_ref), armed: true, diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index 4995f8708..8dbf61ea8 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -227,6 +227,33 @@ async fn buffered_http_failure_is_bounded_and_filters_headers() { 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( @@ -352,6 +379,63 @@ fn target_validation_rejects_unsafe_transport_inputs() { } } +#[test] +fn target_validation_reports_malformed_method_and_headers() { + let error = LlmDispatchTargetContext::try_new( + "P OST".into(), + "https://provider.example/v1".into(), + "openai_chat".into(), + BTreeMap::new(), + ) + .expect_err("method token containing a space should be rejected"); + let FlowError::InvalidArgument(message) = error else { + panic!("expected invalid method argument"); + }; + assert_eq!(message, "LLM continuation method was invalid or prohibited"); + + let error = LlmDispatchTargetContext::try_new( + "POST".into(), + "https://provider.example/v1".into(), + "openai_chat".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( + "POST".into(), + "https://provider.example/v1".into(), + "openai_chat".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_ERROR_HEADER_VALUE_BYTES - 1); + let value = format!("{expected}\u{e9}"); + assert_eq!(value.len(), MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES + 1); + + let bounded = bounded_utf8(value, MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES); + + assert_eq!(bounded, expected); + assert_eq!(bounded.len(), MAX_UPSTREAM_ERROR_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"); diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index f8a4d35f6..de8ec69db 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -144,16 +144,15 @@ unsafe extern "C" fn record_typed_llm_stream_open( let result = if error_json.is_null() { Ok(stream as usize) } else { - parse_json_arg(error_json, "typed LLM stream open error") - .and_then(|value| { - serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson) - }) - .map_err(|status| { - non_http_llm_failure( - LlmNonHttpFailureKindV2::Internal, - format!("invalid typed stream open error: {status:?}"), - ) - }) + 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:?}"), + )), + } }; let _ = sender.send(result); } @@ -205,6 +204,60 @@ unsafe extern "C" fn record_native_forward_terminal( 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 +} + #[derive(Default)] struct NativeStreamCallbackState { error: Mutex>, @@ -1292,6 +1345,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 @@ -1329,6 +1390,13 @@ 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!( @@ -1465,7 +1533,7 @@ 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 { sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), @@ -1896,7 +1964,10 @@ fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { ("Retry-After".into(), "1".into()), ( "X-Request-ID".into(), - "x".repeat(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES + 50), + format!( + "{}é", + "x".repeat(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES - 1) + ), ), ]), class: UpstreamFailureClass::RetryableStatus, @@ -1911,7 +1982,7 @@ fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { assert_eq!(headers.get("retry-after").map(String::as_str), Some("1")); assert_eq!( headers.get("x-request-id").map(String::len), - Some(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES) + Some(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES - 1) ); assert!(!headers.contains_key("authorization")); assert!(!headers.contains_key("set-cookie")); @@ -2760,16 +2831,1414 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { .await .expect("provider stream stress test should not deadlock") }); - assert_eq!(event_counts, vec![EVENT_COUNT; STREAM_COUNT]); + 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 { + 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 native_api_v2_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_v2_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_v2_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_v2_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_v2_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_api_v2_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_v2_with_user_data( + return_v2_stream_state_and_release_handles, + user_data, + ); + let mut stream = runtime + .block_on(wrapped( + "stream-contract", + test_llm_request(), + empty_llm_stream_next(), + )) + .unwrap(); + let item = runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .expect("stream callback monitor should settle") + }); + match expected_error { + Some(expected) => { + let error = item.expect("contract failure item").unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + assert!(runtime.block_on(stream.next()).is_none()); + } + None => assert!(item.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_v2_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_v2_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(&LlmContinuationOutcomeV2::Success { + response: json!({"ok": true}), + }); + typed_guard.complete(&LlmContinuationOutcomeV2::Success { + response: 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 { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + } + } + )); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (provider_sender, provider_receiver) = tokio::sync::mpsc::channel(1); + let provider = Arc::new(NativeLlmProviderStreamV2 { + receiver: tokio::sync::Mutex::new(provider_receiver), + producer_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + next_in_flight: AtomicBool::new(false), + cancelled: AtomicBool::new(false), + _library_guard: None, + }); + 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 { + failure: LlmNonHttpFailureV2 { + 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 { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + }) + )); + drop(provider_sender); + 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, + _library_guard: None, + }; + forward_guard.complete(); + forward_guard.complete(); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + + let (stream, _receiver) = test_native_output_stream(1); + let terminal = NativeForwardTerminalState::default(); + fail_native_string_allocation_after(0); + let mut forward_guard = NativeLlmStreamForwardCallbackGuardV2 { + cb: record_native_forward_terminal, + user_data: (&terminal as *const NativeForwardTerminalState) as usize, + stream: Arc::clone(&stream), + active: true, + _library_guard: None, + }; + forward_guard.fail("allocation failure"); + forward_guard.fail("ignored"); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + assert_eq!( + terminal + .error + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_deref(), + Some("native LLM stream forwarding failed") + ); + + 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, + _library_guard: None, + }); + assert!(stream.cancelled.load(Ordering::Acquire)); + assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); +} + +#[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 { failure } = typed_llm_failure(error) else { + panic!("expected non-HTTP failure") + }; + assert_eq!(failure.kind, expected_kind); + } + + let message = format!("{}é", "x".repeat(NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES)); + let LlmContinuationFailureV2::NonHttp { failure } = + non_http_llm_failure(LlmNonHttpFailureKindV2::Internal, message) + else { + panic!("expected non-HTTP failure") + }; + assert_eq!( + failure.message.len(), + NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES + ); + assert!(failure.message.is_char_boundary(failure.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", + nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, + ), + }) + .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 (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + next_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + let outcome = runtime.block_on(receiver).unwrap(); + match outcome { + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::Http { failure }, + } => assert_eq!(Some(failure.status), http_status), + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { failure }, + } => assert_eq!(Some(failure.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 (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_next_invoke_result_v2( + panicking_ref, + dispatch, + complete_typed_llm_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + assert!(matches!( + runtime.block_on(receiver).unwrap(), + LlmContinuationOutcomeV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + 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 (sender, open_receiver) = + tokio::sync::oneshot::channel::>(); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + next_ref, + dispatch, + output_ref, + record_typed_llm_stream_open, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + 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 (sender, open_receiver) = + tokio::sync::oneshot::channel::>(); + assert_eq!( + unsafe { + native_async_llm_next_open_stream_v2( + panicking_ref, + dispatch, + panic_output_ref, + record_typed_llm_stream_open, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + assert!(matches!( + runtime.block_on(open_receiver).unwrap(), + Err(LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + 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); + assert!(finish_forwarded_native_stream(&settled)); + assert!(!finish_forwarded_native_stream(&settled)); + assert!( + runtime + .block_on(push_forwarded_native_stream_chunk(&settled, Json::Null)) + .is_err() + ); + assert!(!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_closed_and_cancelled_states() { + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + ptr::null(), + record_typed_llm_stream_next, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { native_async_llm_stream_cancel_v2(ptr::null()) }, + NemoRelayStatus::NullPointer + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (sender, receiver) = tokio::sync::mpsc::channel(1); + drop(sender); + let stream = Arc::new(NativeLlmProviderStreamV2 { + receiver: tokio::sync::Mutex::new(receiver), + producer_abort: Mutex::new(None), + runtime: runtime.handle().clone(), + next_in_flight: AtomicBool::new(false), + cancelled: AtomicBool::new(false), + _library_guard: None, + }); + let stream_ref = Arc::into_raw(stream) as *const NemoRelayNativeLlmStreamV2; + let (sender, receiver) = tokio::sync::oneshot::channel(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + stream_ref, + record_typed_llm_stream_next, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + assert!(matches!( + runtime.block_on(receiver).unwrap(), + LlmContinuationStreamEventV2::Failure { + error: LlmContinuationFailureV2::NonHttp { + failure: LlmNonHttpFailureV2 { + kind: LlmNonHttpFailureKindV2::Cancelled, + .. + } + } + } + )); + assert_eq!( + unsafe { native_async_llm_stream_cancel_v2(stream_ref) }, + NemoRelayStatus::Ok + ); + 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("cancelled"); + unsafe { + drop(Box::from_raw(callback_state)); + native_async_llm_stream_release_v2(stream_ref); + native_async_llm_stream_release_v2(ptr::null()); + } +} + +#[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_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 { + 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()); - 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); + native_string_free(invalid_utf8); + native_async_completion_release(completion_ref); } } @@ -5198,6 +6667,28 @@ fn native_registration_entrypoints_reject_null_contexts() { ), NemoRelayStatus::NullPointer ); + assert_eq!( + native_plugin_context_register_async_llm_execution_v2( + ptr::null_mut(), + ptr::null(), + 0, + return_v2_callback_state_and_release_next, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_async_llm_stream_execution_v2( + ptr::null_mut(), + ptr::null(), + 0, + return_v2_stream_state_and_release_handles, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); } assert_last_error_contains("plugin context is null"); } @@ -6059,45 +7550,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( @@ -6108,9 +7604,17 @@ 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"); let request_decode_sentinel = native_string("request-decode-sentinel"); - let mut output = request_decode_sentinel; + output = request_decode_sentinel; set_native_last_error("stale request decode error"); assert_eq!( unsafe { @@ -6309,6 +7813,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, @@ -6330,6 +7844,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, @@ -6379,6 +7912,48 @@ unsafe extern "C" fn llm_response_alias( NemoRelayStatus::Ok } +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 +} + +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 +} + +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 +} + +unsafe extern "C" fn llm_stream_execution_error( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _request_json: *const NemoRelayNativeString, + _next_fn: NemoRelayNativeLlmStreamNextFn, + _next_ctx: *mut c_void, + _out_stream: *mut NemoRelayNativeLlmStreamV1, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + #[test] fn native_callback_helpers_cover_success_error_and_invalid_output() { assert_eq!( @@ -6392,6 +7967,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"}), @@ -6421,6 +8027,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!( @@ -6595,6 +8220,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] @@ -6742,6 +8408,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, @@ -6818,9 +8488,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 @@ -6834,6 +8516,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!( @@ -6856,9 +8542,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 @@ -6984,6 +8682,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()); @@ -6994,19 +8697,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] @@ -7123,9 +8824,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 From 59ca42200602ef37a154f45589ec0f0846cdd1f8 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sat, 1 Aug 2026 11:32:16 -0600 Subject: [PATCH 14/32] refactor(plugin): narrow continuation transport contract Signed-off-by: Bryan Bednarski --- .../src/api/runtime/llm_dispatch_context.rs | 9 -- crates/core/src/plugin/dynamic/native.rs | 1 - .../tests/fixtures/native_plugin/src/lib.rs | 5 +- .../tests/unit/llm_dispatch_context_tests.rs | 15 +--- crates/core/tests/unit/native_plugin_tests.rs | 74 ++++------------ crates/plugin/README.md | 12 +-- crates/plugin/src/lib.rs | 47 ---------- crates/plugin/tests/typed_callbacks.rs | 86 ++----------------- .../dynamic-plugins/native-dynamic/about.mdx | 17 ++-- 9 files changed, 42 insertions(+), 224 deletions(-) diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 1f16f216c..02bdf480c 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -39,7 +39,6 @@ tokio::task_local! { pub struct LlmDispatchTargetContext { method: Method, url: Url, - route: String, headers: HeaderMap, } @@ -48,7 +47,6 @@ impl LlmDispatchTargetContext { pub(crate) fn try_new( method: String, url: String, - route: String, headers: BTreeMap, ) -> Result { let method = Method::from_bytes(method.as_bytes()).map_err(|_| { @@ -97,7 +95,6 @@ impl LlmDispatchTargetContext { Ok(Self { method, url, - route, headers: validated_headers, }) } @@ -116,11 +113,6 @@ impl LlmDispatchTargetContext { &self.url } - #[cfg(test)] - pub(crate) fn route(&self) -> &str { - &self.route - } - /// Explicit provider headers selected for this invocation. #[doc(hidden)] #[must_use] @@ -138,7 +130,6 @@ impl fmt::Debug for LlmDispatchTargetContext { .debug_struct("LlmDispatchTargetContext") .field("method", &self.method) .field("url", &redacted_url) - .field("route", &self.route) .field( "header_names", &self diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index f533886fc..28411a2dc 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -2559,7 +2559,6 @@ fn prepare_llm_continuation_invocation( let target = LlmDispatchTargetContext::try_new( invocation.target.method, invocation.target.url, - invocation.target.route.as_str().into(), invocation.target.headers, ) .map_err(|error| { diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index e1b9711b9..cfba96e46 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -8,8 +8,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use futures::StreamExt; use nemo_relay_plugin::{ CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, - Json, LlmContinuationInvocationV2, LlmContinuationRouteV2, LlmContinuationTargetV2, - LlmJsonAsyncStreamV2, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, + Json, LlmContinuationInvocationV2, LlmContinuationTargetV2, LlmJsonAsyncStreamV2, + LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, LlmStreamExecutionOutcomeV2, NativePlugin, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, @@ -380,7 +380,6 @@ fn targeted_fixture_invocation( target: LlmContinuationTargetV2 { method: "POST".into(), url, - route: LlmContinuationRouteV2::OpenaiChat, headers: std::collections::BTreeMap::from([( "authorization".into(), "Bearer fixture-target".into(), diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index 8dbf61ea8..29763aba4 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -131,7 +131,7 @@ fn response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { } fn target(url: String, headers: BTreeMap) -> LlmDispatchTargetContext { - LlmDispatchTargetContext::try_new("POST".into(), url, "openai_chat".into(), headers) + LlmDispatchTargetContext::try_new("POST".into(), url, headers) .expect("test target should be valid") } @@ -367,15 +367,7 @@ fn target_validation_rejects_unsafe_transport_inputs() { ]), ), ] { - assert!( - LlmDispatchTargetContext::try_new( - method.into(), - url.into(), - "openai_chat".into(), - headers, - ) - .is_err() - ); + assert!(LlmDispatchTargetContext::try_new(method.into(), url.into(), headers,).is_err()); } } @@ -384,7 +376,6 @@ fn target_validation_reports_malformed_method_and_headers() { let error = LlmDispatchTargetContext::try_new( "P OST".into(), "https://provider.example/v1".into(), - "openai_chat".into(), BTreeMap::new(), ) .expect_err("method token containing a space should be rejected"); @@ -396,7 +387,6 @@ fn target_validation_reports_malformed_method_and_headers() { let error = LlmDispatchTargetContext::try_new( "POST".into(), "https://provider.example/v1".into(), - "openai_chat".into(), BTreeMap::from([("bad header".into(), "value".into())]), ) .expect_err("header name containing a space should be rejected"); @@ -411,7 +401,6 @@ fn target_validation_reports_malformed_method_and_headers() { let error = LlmDispatchTargetContext::try_new( "POST".into(), "https://provider.example/v1".into(), - "openai_chat".into(), BTreeMap::from([("x-target".into(), "line one\nline two".into())]), ) .expect_err("header value containing a newline should be rejected"); diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index de8ec69db..9d79e5efc 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -157,14 +157,10 @@ unsafe extern "C" fn record_typed_llm_stream_open( let _ = sender.send(result); } -fn test_dispatch_target( - url: &str, - route: nemo_relay_plugin::LlmContinuationRouteV2, -) -> nemo_relay_plugin::LlmContinuationTargetV2 { +fn test_dispatch_target(url: &str) -> nemo_relay_plugin::LlmContinuationTargetV2 { nemo_relay_plugin::LlmContinuationTargetV2 { method: "POST".into(), url: url.into(), - route, headers: BTreeMap::new(), } } @@ -1678,7 +1674,6 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { target.url().as_str(), "https://provider.example/v1/chat/completions" ); - assert_eq!(target.route(), "openai_chat"); assert_eq!( target .headers() @@ -1706,10 +1701,7 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { }, target: nemo_relay_plugin::LlmContinuationTargetV2 { headers: BTreeMap::from([("authorization".into(), "Bearer target-secret".into())]), - ..test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ) + ..test_dispatch_target("https://provider.example/v1/chat/completions") }, }) .unwrap(), @@ -1783,10 +1775,7 @@ fn native_api_v2_releasing_next_cancels_a_pending_targeted_call() { headers: Map::new(), content: json!({}), }, - target: test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ), + target: test_dispatch_target("https://provider.example/v1/chat/completions"), }) .unwrap(), ) @@ -1852,10 +1841,7 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { headers: Map::new(), content: json!({}), }, - target: test_dispatch_target( - "/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ), + target: test_dispatch_target("/v1/chat/completions"), }) .unwrap(), ) @@ -1886,10 +1872,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { headers: Map::new(), content: json!({}), }; - let mut target = test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ); + let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); target.method = "CONNECT".into(); assert_eq!( prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { @@ -1901,10 +1884,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { ); assert_last_error_contains("method was invalid or prohibited"); - let mut target = test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ); + let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); target .headers .insert("x-nemo-relay-internal-dispatch-url".into(), "secret".into()); @@ -1918,10 +1898,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { ); assert_last_error_contains("host-owned or prohibited"); - let mut target = test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ); + let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); target .headers .insert("host".into(), "attacker.invalid".into()); @@ -1935,10 +1912,7 @@ fn native_api_v2_rejects_prohibited_target_methods_and_headers() { #[test] fn native_api_v2_rejects_target_url_credentials() { - let target = test_dispatch_target( - "https://user:secret@provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ); + let target = test_dispatch_target("https://user:secret@provider.example/v1/chat/completions"); assert_eq!( prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { request: LlmRequest { @@ -2062,8 +2036,9 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { assert_eq!( current_llm_dispatch_target() .expect("typed target is bound") - .route(), - "anthropic_messages" + .url() + .as_str(), + "https://provider.example/v1/messages" ); Ok(LlmJsonStream::new(tokio_stream::iter(vec![ Ok(json!({"type": "content_block_delta", "delta": {"text": "hi"}})), @@ -2098,10 +2073,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { headers: Map::new(), content: json!({"model": "provider/model", "stream": true}), }, - target: test_dispatch_target( - "https://provider.example/v1/messages", - nemo_relay_plugin::LlmContinuationRouteV2::AnthropicMessages, - ), + target: test_dispatch_target("https://provider.example/v1/messages"), }) .unwrap(), ) @@ -2579,10 +2551,7 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { headers: Map::new(), content: json!({"model": "provider/model"}), }, - target: test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ), + target: test_dispatch_target("https://provider.example/v1/chat/completions"), }) .unwrap(), ) @@ -2669,10 +2638,9 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { "authorization".into(), format!("Bearer target-{index}"), )]), - ..test_dispatch_target( - &format!("https://provider-{index}.example/v1/chat/completions"), - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ) + ..test_dispatch_target(&format!( + "https://provider-{index}.example/v1/chat/completions" + )) }, }) .unwrap(), @@ -2756,10 +2724,7 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { headers: Map::new(), content: json!({"model": "provider/model", "stream": true}), }, - target: test_dispatch_target( - "https://provider.example/v1/chat/completions", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ), + target: test_dispatch_target("https://provider.example/v1/chat/completions"), }) .unwrap(), ) @@ -3301,10 +3266,7 @@ 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", - nemo_relay_plugin::LlmContinuationRouteV2::OpenaiChat, - ), + target: test_dispatch_target("https://provider.example/v1/chat/completions"), }) .unwrap(), ) diff --git a/crates/plugin/README.md b/crates/plugin/README.md index e5207d895..6d6440f39 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -153,12 +153,12 @@ return `LlmStreamExecutionOutcomeV2::Passthrough(request)`. Relay then pumps the original downstream stream directly through its bounded queue; provider events do not cross into the plugin merely to be forwarded. -The plugin provides JSON plus an HTTP method, absolute target URL, protocol -route, and explicit target headers. Relay binds that transport target to the -current LLM continuation without storing it in `LlmRequest.headers`. Successful -calls return provider JSON. 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 provides JSON plus an HTTP method, absolute target URL, and explicit +target headers. Relay binds that transport target to the current LLM +continuation without storing it in `LlmRequest.headers`. Successful calls return +provider JSON. 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 diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 50ce6ba5c..ebdc1f690 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -919,29 +919,6 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( error: *const NemoRelayNativeString, ); -/// Provider protocol selected for one targeted LLM continuation. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LlmContinuationRouteV2 { - /// OpenAI Chat Completions. - OpenaiChat, - /// OpenAI Responses. - OpenaiResponses, - /// Anthropic Messages. - AnthropicMessages, -} - -impl LlmContinuationRouteV2 { - /// Returns the stable Relay gateway route identifier. - pub const fn as_str(self) -> &'static str { - match self { - Self::OpenaiChat => "openai_chat", - Self::OpenaiResponses => "openai_responses", - Self::AnthropicMessages => "anthropic_messages", - } - } -} - /// Explicit provider target for one native API v2 LLM continuation. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] pub struct LlmContinuationTargetV2 { @@ -949,8 +926,6 @@ pub struct LlmContinuationTargetV2 { pub method: String, /// Absolute HTTP(S) provider URL including the selected endpoint. pub url: String, - /// Provider protocol used by the selected endpoint. - pub route: LlmContinuationRouteV2, /// Explicit outbound provider headers, including target credentials. /// /// Relay validates and transports these headers but never records their @@ -1021,28 +996,6 @@ pub enum LlmContinuationFailureV2 { }, } -impl LlmContinuationFailureV2 { - /// Return Relay's provider-neutral retry disposition. - /// - /// The disposition is derived rather than serialized so the wire contract - /// contains only HTTP semantics and the minimal non-HTTP failure kind. - pub const fn is_retryable(&self) -> bool { - match self { - Self::Http { failure } => is_retryable_http_status_v2(failure.status), - Self::NonHttp { failure } => matches!( - failure.kind, - LlmNonHttpFailureKindV2::Transport | LlmNonHttpFailureKindV2::Timeout - ), - } - } -} - -/// Return Relay's provider-neutral retry disposition for an HTTP status. -#[must_use] -pub const fn is_retryable_http_status_v2(status: u16) -> bool { - matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504) -} - /// Unary LLM continuation outcome delivered through native API v2. #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 7b1b40d3d..7ad20e349 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -17,13 +17,12 @@ use futures::{StreamExt, stream}; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmContinuationFailureV2, - LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationRouteV2, - LlmContinuationStreamEventV2, LlmContinuationTargetV2, LlmContinuationV2, LlmHttpFailureV2, - LlmJsonStream, LlmNext, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, - LlmRequestInterceptOutcome, LlmStream, LlmStreamContinuationV2, LlmStreamExecutionOutcomeV2, - LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, - NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, + LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationStreamEventV2, + LlmContinuationTargetV2, LlmContinuationV2, LlmHttpFailureV2, LlmJsonStream, LlmNext, + LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, 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, @@ -44,68 +43,6 @@ use nemo_relay_plugin::{ }; use serde_json::{Map, json}; -#[test] -fn native_api_v2_retry_policy_is_derived_from_http_semantics() { - for status in [408, 425, 429, 500, 502, 503, 504] { - assert!( - LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { - status, - body: String::new(), - headers: Default::default(), - }, - } - .is_retryable(), - "status={status}" - ); - } - for status in [400, 401, 404, 409, 422, 501] { - assert!( - !LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { - status, - body: String::new(), - headers: Default::default(), - }, - } - .is_retryable(), - "status={status}" - ); - } - for kind in [ - LlmNonHttpFailureKindV2::Transport, - LlmNonHttpFailureKindV2::Timeout, - ] { - assert!( - LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind, - message: String::new(), - }, - } - .is_retryable(), - "kind={kind:?}" - ); - } - for kind in [ - LlmNonHttpFailureKindV2::Cancelled, - LlmNonHttpFailureKindV2::InvalidRequest, - LlmNonHttpFailureKindV2::Guardrail, - LlmNonHttpFailureKindV2::Internal, - ] { - assert!( - !LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind, - message: String::new(), - }, - } - .is_retryable(), - "kind={kind:?}" - ); - } -} - #[test] fn async_abi_discriminants_reject_unknown_values() { use NemoRelayNativeAsyncMiddlewareKind as Kind; @@ -6584,7 +6521,6 @@ fn safe_v2_target() -> LlmContinuationTargetV2 { LlmContinuationTargetV2 { method: "POST".into(), url: "https://provider.example/v1/chat/completions".into(), - route: LlmContinuationRouteV2::OpenaiChat, headers: Default::default(), } } @@ -6784,15 +6720,6 @@ fn native_api_v2_raw_registration_remains_an_advanced_escape_hatch() { NemoRelayStatus::InvalidArg ); - assert_eq!(LlmContinuationRouteV2::OpenaiChat.as_str(), "openai_chat"); - assert_eq!( - LlmContinuationRouteV2::OpenaiResponses.as_str(), - "openai_responses" - ); - assert_eq!( - LlmContinuationRouteV2::AnthropicMessages.as_str(), - "anthropic_messages" - ); assert_eq!(live_host_strings(), 0); } @@ -7206,7 +7133,6 @@ fn safe_v2_stream_open_preserves_structured_failure() { Err(error) => error, }; assert_eq!(error, expected); - assert!(error.is_retryable()); Err("observed typed stream-open failure".into()) } }, diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 103220916..6fc9e1def 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -133,10 +133,10 @@ 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 HTTP method, absolute -HTTP(S) URL, provider route (`openai_chat`, `openai_responses`, or -`anthropic_messages`), and explicit outbound headers. Target headers may -contain provider credentials. Relay validates and transports them but never -places their values in diagnostics or observability. +HTTP(S) URL, and explicit outbound headers. 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 @@ -156,11 +156,10 @@ 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. -HTTP retryability is derived by the SDK from status alone: `408`, `425`, `429`, -`500`, `502`, `503`, and `504` are retryable. Transport and timeout failures -are retryable; other non-HTTP failures are not. Relay does not inspect provider -bodies to infer context-window or model-availability errors, so ordinary `400` -and `404` responses do not trigger automatic reselection. +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 From 38a11fb2a953c6f4a45b960e5f6c882dd1b0559d Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 10:17:23 -0600 Subject: [PATCH 15/32] refactor(plugin): poll native v2 callbacks cooperatively Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 962 ++++++++++++------ .../tests/fixtures/native_plugin/src/lib.rs | 106 +- .../tests/integration/native_plugin_tests.rs | 283 +++++- crates/core/tests/unit/native_plugin_tests.rs | 752 ++++++++++++-- crates/plugin/README.md | 28 +- crates/plugin/src/lib.rs | 129 +-- crates/plugin/src/native_v2.rs | 535 +++++++--- crates/plugin/tests/typed_callbacks.rs | 678 +++++++++--- .../dynamic-plugins/native-dynamic/about.mdx | 42 +- 9 files changed, 2629 insertions(+), 886 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 28411a2dc..7b44b234e 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; @@ -60,10 +60,11 @@ use nemo_relay_plugin::{ NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, - NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, - NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, - NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeAsyncTaskPollCbV2, + NemoRelayNativeAsyncTaskV2, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, + NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, + NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, @@ -98,8 +99,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)>, @@ -115,7 +117,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)] @@ -132,6 +151,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(); } } @@ -286,8 +306,10 @@ struct NativePluginInstance { plugin_kind: String, relay_compat: String, allows_multiple_components: bool, + uses_native_api_v2: bool, plugin: Mutex, - _library: Library, + library: Option, + retain_library_on_drop: AtomicBool, } unsafe impl Send for NativePluginInstance {} @@ -298,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); + } } } @@ -436,8 +465,10 @@ fn load_one_native_plugin( plugin_kind, relay_compat, allows_multiple_components: plugin.allows_multiple_components, + uses_native_api_v2: native_api == Some("2"), plugin: Mutex::new(plugin), - _library: library, + library: Some(library), + retain_library_on_drop: AtomicBool::new(false), })) } @@ -544,6 +575,7 @@ struct NativeHostScopeStackBinding(ThreadScopeStackBinding); thread_local! { static NATIVE_LAST_ERROR: RefCell> = const { RefCell::new(None) }; + static CURRENT_NATIVE_ASYNC_TASK_V2: RefCell>> = const { RefCell::new(None) }; #[cfg(test)] static NATIVE_STRING_LIVE_ALLOCATIONS: RefCell> = RefCell::new(HashSet::new()); #[cfg(test)] @@ -925,10 +957,11 @@ fn build_native_host_api_v4() -> NemoRelayNativeHostApiV4 { async_llm_stream_next_v2: native_async_llm_stream_next_v2, async_llm_stream_cancel_v2: native_async_llm_stream_cancel_v2, async_llm_stream_release_v2: native_async_llm_stream_release_v2, - plugin_context_register_async_llm_execution_v2: - native_plugin_context_register_async_llm_execution_v2, - plugin_context_register_async_llm_stream_execution_v2: - native_plugin_context_register_async_llm_stream_execution_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, } } @@ -1460,25 +1493,6 @@ impl Drop for NativeNextHandoff { } } -struct NativeStreamHandoff { - raw: usize, - armed: bool, -} - -impl NativeStreamHandoff { - fn disarm(&mut self) { - self.armed = false; - } -} - -impl Drop for NativeStreamHandoff { - fn drop(&mut self) { - if self.armed { - unsafe { native_async_stream_release(self.raw as *const NemoRelayNativeAsyncStream) }; - } - } -} - unsafe impl Send for NativeCallbackUserData {} unsafe impl Sync for NativeCallbackUserData {} @@ -1502,14 +1516,28 @@ fn make_user_data( }) } -const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY: usize = 64; -const NATIVE_API_V2_STREAM_CHANNEL_CAPACITY: usize = 32; +// Native API v1's V3 incremental stream contract shipped with a 64-event +// queue. Keep that observable backpressure boundary stable for existing +// plugins while native API v2 uses the documented tighter bound. +const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1: usize = 64; +const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2: usize = 32; + +fn native_async_stream_channel_capacity(uses_native_api_v2: bool) -> usize { + if uses_native_api_v2 { + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2 + } else { + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1 + } +} struct NativeAsyncCompletion { sender: Mutex>>>, 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 @@ -1548,6 +1576,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); } } @@ -1601,11 +1631,294 @@ struct NativeAsyncStream { settled: AtomicBool, downstream_aborts: Mutex>, settlement: Mutex<()>, + runtime: tokio::runtime::Handle, + context: MiddlewareContinuationContext, + task: Mutex>>, + backpressure_waiter: 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 is_stream(&self, stream: *const NativeAsyncStream) -> bool { + match self { + Self::Stream(owner) => owner.as_ptr() == stream, + Self::Completion(_) => false, + } + } + + 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); + clear_native_async_task_waiter_slot(&owner.backpressure_waiter, 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(); + } +} + +fn clear_native_async_task_waiter_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, + _library_guard: Option>, +} + +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 { + runtime: tokio::runtime::Handle, + 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>, +} + +struct CurrentNativeAsyncTaskBindingV2(Option>); + +impl CurrentNativeAsyncTaskBindingV2 { + fn bind(task: &Arc) -> Self { + let previous = CURRENT_NATIVE_ASYNC_TASK_V2 + .with(|current| current.replace(Some(Arc::downgrade(task)))); + Self(previous) + } +} + +impl Drop for CurrentNativeAsyncTaskBindingV2 { + fn drop(&mut self) { + let previous = self.0.take(); + CURRENT_NATIVE_ASYNC_TASK_V2.with(|current| { + current.replace(previous); + }); + } +} + +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 { + let _binding = CurrentNativeAsyncTaskBindingV2::bind(&task); + 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, @@ -1909,7 +2222,11 @@ impl Stream for NativeAsyncStreamReceiver { type Item = FlowResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.receiver.poll_recv(cx) + let result = self.receiver.poll_recv(cx); + if result.is_ready() { + wake_native_async_stream_backpressure_waiter(&self.stream); + } + result } } @@ -1935,43 +2252,77 @@ impl Drop for NativeAsyncStreamReceiver { .lock() .unwrap_or_else(|error| error.into_inner()) .take(); + drop(_settlement); + wake_native_async_stream_backpressure_waiter(&self.stream); + cancel_native_async_task_slot(&self.stream.task); } } -async fn invoke_native_async_callback( - cb: NemoRelayNativeAsyncMiddlewareCb, - user_data: Arc, - invocation: Json, - next: Option, -) -> FlowResult { - invoke_native_async_callback_with_lane(cb, user_data, invocation, next, false).await +fn wake_native_async_stream_backpressure_waiter(stream: &NativeAsyncStream) { + let task = stream + .backpressure_waiter + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .and_then(|task| task.upgrade()); + if let Some(task) = task { + NativeAsyncTaskV2::wake(&task); + } } -async fn invoke_native_async_callback_blocking( - cb: NemoRelayNativeAsyncMiddlewareCb, - user_data: Arc, - invocation: Json, - next: Option, -) -> FlowResult { - invoke_native_async_callback_with_lane(cb, user_data, invocation, next, true).await +fn register_current_native_async_stream_backpressure_waiter( + stream: &NativeAsyncStream, + sender: &tokio::sync::mpsc::Sender>, +) { + let task = CURRENT_NATIVE_ASYNC_TASK_V2.with(|current| { + current + .borrow() + .as_ref() + .and_then(Weak::upgrade) + .filter(|task| task.owner.is_stream(stream as *const NativeAsyncStream)) + }); + let Some(task) = task else { + return; + }; + *stream + .backpressure_waiter + .lock() + .unwrap_or_else(|error| error.into_inner()) = Some(Arc::downgrade(&task)); + // The consumer can free capacity between `try_send` and waiter + // registration. Recheck after publishing the waiter to avoid a lost wake. + if sender.capacity() > 0 || sender.is_closed() { + wake_native_async_stream_backpressure_waiter(stream); + } +} + +fn clear_current_native_async_stream_backpressure_waiter(stream: &NativeAsyncStream) { + let current = + CURRENT_NATIVE_ASYNC_TASK_V2.with(|task| task.borrow().as_ref().map(Weak::as_ptr)); + let mut waiter = stream + .backpressure_waiter + .lock() + .unwrap_or_else(|error| error.into_inner()); + if current.is_some_and(|current| { + waiter + .as_ref() + .is_some_and(|waiter| waiter.as_ptr() == current) + }) { + waiter.take(); + } } -async fn invoke_native_async_callback_with_lane( +async fn invoke_native_async_callback( cb: NemoRelayNativeAsyncMiddlewareCb, user_data: Arc, invocation: Json, next: Option, - blocking: bool, ) -> 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; @@ -1981,6 +2332,9 @@ async fn invoke_native_async_callback_with_lane( 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()), @@ -1991,15 +2345,13 @@ async fn invoke_native_async_callback_with_lane( 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); @@ -2042,36 +2394,13 @@ async fn invoke_native_async_callback_with_lane( } result }; - let state = if blocking { - // Cleanup lives in guards captured by the blocking closure. The - // monitor reports the callback state only; runtime shutdown cannot - // strand host-owned strings or references in that task. - let blocking_task = tokio::task::spawn_blocking(invoke); - let (state_sender, state_receiver) = tokio::sync::oneshot::channel(); - tokio::spawn(async move { - let state = match blocking_task.await { - Ok(Ok(state)) => NemoRelayNativeAsyncCallbackState::try_from(state).map_err(|()| { - FlowError::Internal("native async callback returned an invalid state".into()) - }), - Ok(Err(_)) => Err(FlowError::Internal("native async callback panicked".into())), - Err(error) => Err(FlowError::Internal(format!( - "native API v2 blocking callback task failed: {error}" - ))), - }; - let _ = state_sender.send(state); - }); - state_receiver.await.map_err(|_| { - FlowError::Internal("native API v2 blocking callback monitor stopped".into()) - })?? - } else { - let callback_result = invoke(); - 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())), - }? - }; + 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 @@ -2122,6 +2451,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 } @@ -2167,6 +2498,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 } @@ -2185,6 +2518,176 @@ unsafe extern "C" fn native_async_completion_release( } } +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; + } + 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; + } + let mut slot = completion + .task + .lock() + .unwrap_or_else(|error| error.into_inner()); + if slot.as_ref().and_then(Weak::upgrade).is_some() { + set_native_last_error("an async completion task is already active"); + return NemoRelayStatus::InvalidArg; + } + let task = Arc::new(NativeAsyncTaskV2 { + runtime: completion.runtime.clone(), + context: completion.context.clone(), + owner: NativeAsyncTaskOwnerV2::Completion(unsafe { + weak_from_arc_raw(completion as *const NativeAsyncCompletion) + }), + wake: tokio::sync::Notify::new(), + state: Mutex::new(NativeAsyncTaskStateV2 { + polling: false, + cancel_requested: false, + complete: false, + plugin: Some(NativeAsyncTaskPluginStateV2 { + cb, + user_data: user_data as usize, + free_fn, + _library_guard: completion._callback_user_data.clone(), + }), + }), + _library_guard: completion._callback_user_data.clone(), + }); + *slot = Some(Arc::downgrade(&task)); + drop(slot); + task.runtime.spawn(Arc::clone(&task).run()); + NativeAsyncTaskV2::wake(&task); + NemoRelayStatus::Ok +} + +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; + } + let mut slot = stream + .task + .lock() + .unwrap_or_else(|error| error.into_inner()); + if slot.as_ref().and_then(Weak::upgrade).is_some() { + set_native_last_error("an async stream task is already active"); + return NemoRelayStatus::InvalidArg; + } + let task = Arc::new(NativeAsyncTaskV2 { + runtime: stream.runtime.clone(), + context: stream.context.clone(), + owner: NativeAsyncTaskOwnerV2::Stream(unsafe { + weak_from_arc_raw(stream as *const NativeAsyncStream) + }), + wake: tokio::sync::Notify::new(), + state: Mutex::new(NativeAsyncTaskStateV2 { + polling: false, + cancel_requested: false, + complete: false, + plugin: Some(NativeAsyncTaskPluginStateV2 { + cb, + user_data: user_data as usize, + free_fn, + _library_guard: stream._callback_user_data.clone(), + }), + }), + _library_guard: stream._callback_user_data.clone(), + }); + *slot = Some(Arc::downgrade(&task)); + drop(slot); + task.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, +) -> NemoRelayStatus { + if task.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { Arc::increment_strong_count(task as *const NativeAsyncTaskV2) }; + let task = unsafe { Arc::from_raw(task as *const NativeAsyncTaskV2) }; + NativeAsyncTaskV2::wake(&task); + NemoRelayStatus::Ok +} + +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)) }; @@ -2229,8 +2732,12 @@ unsafe extern "C" fn native_async_stream_push_json( return NemoRelayStatus::InvalidArg; }; match sender.try_send(Ok(chunk)) { - Ok(()) => NemoRelayStatus::Ok, + Ok(()) => { + clear_current_native_async_stream_backpressure_waiter(stream); + NemoRelayStatus::Ok + } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + register_current_native_async_stream_backpressure_waiter(stream, &sender); set_native_last_error( "native async stream is backpressured; retry the chunk after the consumer advances", ); @@ -2275,6 +2782,10 @@ unsafe extern "C" fn native_async_stream_finish( for (_, abort) in downstream_aborts.drain() { abort.abort(); } + drop(downstream_aborts); + drop(_settlement); + wake_native_async_stream_backpressure_waiter(stream); + cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } else { NemoRelayStatus::InvalidArg @@ -2323,9 +2834,15 @@ unsafe extern "C" fn native_async_stream_reject( for (_, abort) in downstream_aborts.drain() { abort.abort(); } + drop(downstream_aborts); + drop(sender_guard); + drop(_settlement); + wake_native_async_stream_backpressure_waiter(stream); + cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + register_current_native_async_stream_backpressure_waiter(stream, sender); set_native_last_error( "native async stream is backpressured; retry rejection after the consumer advances", ); @@ -3181,7 +3698,7 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( match next_fn(request).await { Ok(mut provider_stream) => { let (sender, receiver) = - tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); let provider = Arc::new(NativeLlmProviderStreamV2 { receiver: tokio::sync::Mutex::new(receiver), producer_abort: Mutex::new(None), @@ -3596,32 +4113,10 @@ fn wrap_native_async_llm_execution( free_fn: NemoRelayNativeFreeFn, ) -> LlmExecutionFn { let user_data = make_user_data(instance, user_data, free_fn); - Arc::new(move |name, request, next| { - let user_data = user_data.clone(); - let name = name.to_owned(); - Box::pin(async move { - invoke_native_async_callback( - cb, - user_data, - serde_json::json!({"name": name, "request": request}), - Some(NativeAsyncNextInner::Llm(next)), - ) - .await - }) - }) + wrap_native_async_llm_execution_with_user_data(cb, user_data) } -fn wrap_native_async_llm_execution_v2( - instance: Arc, - cb: NemoRelayNativeAsyncMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> LlmExecutionFn { - let user_data = make_user_data(instance, user_data, free_fn); - wrap_native_async_llm_execution_v2_with_user_data(cb, user_data) -} - -fn wrap_native_async_llm_execution_v2_with_user_data( +fn wrap_native_async_llm_execution_with_user_data( cb: NemoRelayNativeAsyncMiddlewareCb, user_data: Arc, ) -> LlmExecutionFn { @@ -3629,7 +4124,7 @@ fn wrap_native_async_llm_execution_v2_with_user_data( let user_data = user_data.clone(); let name = name.to_owned(); Box::pin(async move { - invoke_native_async_callback_blocking( + invoke_native_async_callback( cb, user_data, serde_json::json!({"name": name, "request": request}), @@ -3646,26 +4141,53 @@ fn wrap_native_incremental_llm_stream_execution( user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> LlmStreamExecutionFn { + let channel_capacity = native_async_stream_channel_capacity(instance.uses_native_api_v2); 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, + 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_V1, + ) +} + +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), + backpressure_waiter: Mutex::new(None), #[cfg(test)] before_settlement_lock: None, _callback_user_data: Some(user_data.clone()), @@ -3682,11 +4204,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, @@ -3728,121 +4245,7 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data( }) } -fn wrap_native_incremental_llm_stream_execution_v2( - instance: Arc, - cb: NemoRelayNativeAsyncStreamMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> LlmStreamExecutionFn { - let user_data = make_user_data(instance, user_data, free_fn); - wrap_native_incremental_llm_stream_execution_v2_with_user_data(cb, user_data) -} - -fn wrap_native_incremental_llm_stream_execution_v2_with_user_data( - cb: NemoRelayNativeAsyncStreamMiddlewareCb, - user_data: Arc, -) -> 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_API_V2_STREAM_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(()), - #[cfg(test)] - before_settlement_lock: None, - _callback_user_data: Some(user_data.clone()), - }); - let output = NativeAsyncStreamReceiver { - receiver, - stream: Arc::clone(&stream), - }; - let invocation = - native_string_from_json(&serde_json::json!({"name": name, "request": request})) - .ok_or_else(|| { - FlowError::Internal( - "failed to allocate native API v2 stream invocation".into(), - ) - })? as usize; - let invocation_guard = NativeInvocationStringGuard(invocation); - let runtime = tokio::runtime::Handle::try_current().map_err(|error| { - FlowError::Internal(format!( - "native API v2 stream intercept requires a Tokio runtime: {error}" - )) - })?; - let next_ref = Arc::into_raw(Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(next), - runtime, - Some(user_data.clone()), - ))) as usize; - let stream_ref = Arc::into_raw(stream.clone()) as usize; - let callback_user_data = user_data.ptr as usize; - let callback_scope_stack = current_scope_stack(); - let next_handoff = NativeNextHandoff { - raw: Some(next_ref), - armed: true, - }; - let output_handoff = NativeStreamHandoff { - raw: stream_ref, - armed: true, - }; - let blocking_task = tokio::task::spawn_blocking(move || { - let _invocation = invocation_guard; - let mut next = next_handoff; - let mut output = output_handoff; - catch_unwind(AssertUnwindSafe(|| { - // Both handles transfer to the plugin at callback entry. - next.disarm(); - output.disarm(); - with_scope_stack(callback_scope_stack, || unsafe { - cb( - callback_user_data as *mut c_void, - invocation as *const NemoRelayNativeString, - next_ref as *const NemoRelayNativeAsyncNext, - stream_ref as *const NemoRelayNativeAsyncStream, - ) - }) - })) - }); - let stream_for_monitor = Arc::clone(&stream); - tokio::spawn(async move { - let callback_result = blocking_task.await; - let state = callback_result - .ok() - .and_then(std::result::Result::ok) - .and_then(|state| NemoRelayNativeAsyncCallbackState::try_from(state).ok()); - let error = match state { - Some(NemoRelayNativeAsyncCallbackState::Pending) => None, - Some(NemoRelayNativeAsyncCallbackState::Complete) - if stream_for_monitor.settled.load(Ordering::Acquire) - || stream_for_monitor.cancelled.load(Ordering::Acquire) => - { - None - } - Some(NemoRelayNativeAsyncCallbackState::Complete) => Some( - "native API v2 stream callback returned Complete without finishing" - .to_string(), - ), - None => Some( - "native API v2 stream callback panicked or returned an invalid state" - .to_string(), - ), - }; - if let Some(error) = error { - settle_native_api_v2_stream_error(stream_for_monitor, error).await; - } - }); - Ok(LlmJsonStream::new(output)) - }) - }) -} - -async fn settle_native_api_v2_stream_error(stream: Arc, message: String) { +async fn settle_native_async_stream_error(stream: Arc, message: String) { let sender = { let _settlement = stream .settlement @@ -3867,6 +4270,9 @@ async fn settle_native_api_v2_stream_error(stream: Arc, messa for (_, abort) in downstream_aborts.drain() { abort.abort(); } + drop(downstream_aborts); + wake_native_async_stream_backpressure_waiter(&stream); + cancel_native_async_task_slot(&stream.task); } unsafe extern "C" fn native_plugin_context_register_async_stream_middleware( @@ -3900,68 +4306,6 @@ unsafe extern "C" fn native_plugin_context_register_async_stream_middleware( } } -unsafe extern "C" fn native_plugin_context_register_async_llm_stream_execution_v2( - ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncStreamMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> NemoRelayStatus { - clear_native_last_error(); - let user_data_guard = NativeCallbackUserDataGuard::new(user_data, free_fn); - let host_ctx = match host_ctx_mut(ctx) { - Ok(ctx) => ctx, - Err(status) => return status, - }; - let instance = host_ctx.instance.clone(); - let name = match read_name(name) { - Ok(name) => name, - Err(status) => return status, - }; - let (user_data, free_fn) = user_data_guard.transfer(); - let context = unsafe { &mut *host_ctx.ctx }; - match context.register_llm_stream_execution_intercept( - &name, - priority, - wrap_native_incremental_llm_stream_execution_v2(instance, cb, user_data, free_fn), - ) { - Ok(()) => NemoRelayStatus::Ok, - Err(error) => status_from_plugin_error(error), - } -} - -unsafe extern "C" fn native_plugin_context_register_async_llm_execution_v2( - ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> NemoRelayStatus { - clear_native_last_error(); - let user_data_guard = NativeCallbackUserDataGuard::new(user_data, free_fn); - let host_ctx = match host_ctx_mut(ctx) { - Ok(ctx) => ctx, - Err(status) => return status, - }; - let instance = host_ctx.instance.clone(); - let name = match read_name(name) { - Ok(name) => name, - Err(status) => return status, - }; - let (user_data, free_fn) = user_data_guard.transfer(); - let context = unsafe { &mut *host_ctx.ctx }; - match context.register_llm_execution_intercept( - &name, - priority, - wrap_native_async_llm_execution_v2(instance, cb, user_data, free_fn), - ) { - Ok(()) => NemoRelayStatus::Ok, - Err(error) => status_from_plugin_error(error), - } -} - unsafe extern "C" fn native_plugin_context_register_async_middleware( ctx: *mut NemoRelayNativePluginContext, kind: u32, diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index cfba96e46..050b2c7fa 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -3,20 +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, LlmContinuationInvocationV2, LlmContinuationTargetV2, LlmJsonAsyncStreamV2, - LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, + 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, + NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, + NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, NemoRelayNativePluginV1, + NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, + PluginContext, PluginRuntime, ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, }; use serde_json::{Map, json}; @@ -300,13 +300,40 @@ 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 -); +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" @@ -317,21 +344,40 @@ impl NativePlugin for TargetedFixturePlugin { 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)) @@ -357,13 +403,11 @@ impl NativePlugin for TargetedFixturePlugin { .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:?}") - }) - }), - ); + let stream: LlmJsonAsyncStreamV2 = Box::pin(stream.map(|item| { + item.map_err(|error| { + format!("targeted fixture provider stream failed: {error:?}") + }) + })); Ok(LlmStreamExecutionOutcomeV2::Stream(stream)) } }, @@ -371,10 +415,21 @@ impl NativePlugin for TargetedFixturePlugin { } } -fn targeted_fixture_invocation( - url: String, - request: LlmRequest, -) -> LlmContinuationInvocationV2 { +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 { @@ -388,10 +443,9 @@ fn targeted_fixture_invocation( } } -nemo_relay_plugin::nemo_relay_plugin_v2!( - nemo_relay_fixture_targeted_v2_plugin, - || TargetedFixturePlugin -); +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( diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 9fcc270cc..c830549cf 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -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; @@ -1301,62 +1304,238 @@ fn native_api_v2_plugin_requires_the_v2_manifest_contract() { assert!(error.contains("entry symbol"), "{error}"); } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn native_api_v2_fixture_dispatches_through_core_without_a_cli_gateway() { - 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))]), +#[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(); }); - 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 original_provider_called = Arc::new(AtomicBool::new(false)); - let original_provider_called_for_fn = original_provider_called.clone(); - let response = 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 - .expect("embedded targeted dispatch should succeed"); +#[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; + } - 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")); + 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"); - drop(cleanup); - activation.clear(); + 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)] diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 9d79e5efc..6ff8a5a25 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -254,6 +254,127 @@ unsafe extern "C" fn return_v2_stream_state_and_release_handles( 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 result = if poll == 0 { + assert_eq!( + unsafe { native_async_stream_push_json(stream, first) }, + NemoRelayStatus::Ok + ); + assert_eq!( + unsafe { native_async_stream_push_json(stream, second) }, + NemoRelayStatus::Internal + ); + if let Some(started) = &state.started { + let _ = started.send(0); + } + NemoRelayNativeAsyncCallbackState::Pending as u32 + } else { + assert_eq!( + unsafe { native_async_stream_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>, @@ -1072,36 +1193,6 @@ unsafe extern "C" fn invoke_native_next_then_return_state( state.callback_state } -struct BlockingSafeCallbackState { - started: std::sync::mpsc::Sender<()>, - release: Mutex>, - freed: std::sync::mpsc::Sender<()>, -} - -unsafe extern "C" fn blocking_safe_callback( - user_data: *mut c_void, - _invocation_json: *const NemoRelayNativeString, - next: *const NemoRelayNativeAsyncNext, - _completion: *const NemoRelayNativeAsyncCompletion, -) -> u32 { - let state = unsafe { &*user_data.cast::() }; - let _ = state.started.send(()); - let _ = state - .release - .lock() - .unwrap_or_else(|error| error.into_inner()) - .recv_timeout(Duration::from_secs(1)); - unsafe { native_async_next_release(next) }; - NemoRelayNativeAsyncCallbackState::Complete as u32 -} - -unsafe extern "C" fn free_blocking_safe_callback(user_data: *mut c_void) { - let state = unsafe { Box::from_raw(user_data.cast::()) }; - let freed = state.freed.clone(); - drop(state); - let _ = freed.send(()); -} - unsafe extern "C" fn invoke_native_stream_next_then_return_state( user_data: *mut c_void, invocation_json: *const NemoRelayNativeString, @@ -1435,6 +1526,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), @@ -1472,6 +1566,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), @@ -1531,6 +1628,9 @@ fn native_async_next_reports_a_revoked_continuation_without_calling_the_provider 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), @@ -2057,6 +2157,10 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::mpsc::channel(1); let output_stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2187,6 +2291,10 @@ fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2286,6 +2394,10 @@ fn native_api_v2_direct_stream_forwarding_preserves_downstream_failure() { let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (sender, receiver) = tokio::sync::mpsc::channel(2); let stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2385,6 +2497,10 @@ fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2452,7 +2568,7 @@ fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight .build() .unwrap(); let (provider_sender, provider_receiver) = - tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); let provider = Arc::new(NativeLlmProviderStreamV2 { receiver: tokio::sync::Mutex::new(provider_receiver), producer_abort: Mutex::new(None), @@ -2686,7 +2802,8 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { const STREAM_COUNT: usize = 64; const EVENT_COUNT: usize = 100; - assert_eq!(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY, 32); + assert_eq!(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1, 64); + assert_eq!(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2, 32); let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) @@ -2706,8 +2823,12 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (output_sender, output_receiver) = - tokio::sync::mpsc::channel(NATIVE_API_V2_STREAM_CHANNEL_CAPACITY); + tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); let output_stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(output_sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2809,6 +2930,12 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { } } +#[test] +fn native_async_stream_capacity_preserves_v1_and_bounds_v2() { + assert_eq!(native_async_stream_channel_capacity(false), 64); + assert_eq!(native_async_stream_channel_capacity(true), 32); +} + fn test_v2_callback_user_data(ptr: *mut c_void) -> Arc { Arc::new(NativeCallbackUserData { ptr, @@ -2841,6 +2968,10 @@ fn test_native_output_stream( let (sender, receiver) = tokio::sync::mpsc::channel(capacity); ( Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2854,7 +2985,7 @@ fn test_native_output_stream( } #[test] -fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { +fn native_async_buffered_wrapper_enforces_complete_callback_contract() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() @@ -2862,7 +2993,7 @@ fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { .unwrap(); let user_data = test_v2_callback_user_data(ptr::null_mut()); - let wrapped = wrap_native_async_llm_execution_v2_with_user_data( + let wrapped = wrap_native_async_llm_execution_with_user_data( complete_v2_callback_and_release_next, user_data, ); @@ -2883,7 +3014,7 @@ fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { 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_v2_with_user_data( + let wrapped = wrap_native_async_llm_execution_with_user_data( return_v2_callback_state_and_release_next, user_data, ); @@ -2896,7 +3027,7 @@ fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { 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_v2_with_user_data( + let wrapped = wrap_native_async_llm_execution_with_user_data( return_v2_callback_state_and_release_next, user_data, ); @@ -2909,7 +3040,7 @@ fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { 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_v2_with_user_data( + let wrapped = wrap_native_async_llm_execution_with_user_data( return_v2_callback_state_and_release_next, user_data, ); @@ -2923,7 +3054,7 @@ fn native_api_v2_buffered_wrapper_enforces_complete_callback_contract() { } #[test] -fn native_api_v2_stream_wrapper_enforces_callback_and_ownership_contracts() { +fn native_async_stream_wrapper_enforces_callback_and_ownership_contracts() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() @@ -2955,29 +3086,24 @@ fn native_api_v2_stream_wrapper_enforces_callback_and_ownership_contracts() { ] { let user_data = test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); - let wrapped = wrap_native_incremental_llm_stream_execution_v2_with_user_data( + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( return_v2_stream_state_and_release_handles, user_data, ); - let mut stream = runtime - .block_on(wrapped( - "stream-contract", - test_llm_request(), - empty_llm_stream_next(), - )) - .unwrap(); - let item = runtime.block_on(async { - tokio::time::timeout(Duration::from_secs(1), stream.next()) - .await - .expect("stream callback monitor should settle") - }); + let result = runtime.block_on(wrapped( + "stream-contract", + test_llm_request(), + empty_llm_stream_next(), + )); match expected_error { Some(expected) => { - let error = item.expect("contract failure item").unwrap_err(); + 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()); } - None => assert!(item.is_none()), } } @@ -2987,7 +3113,7 @@ fn native_api_v2_stream_wrapper_enforces_callback_and_ownership_contracts() { }; let user_data = test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); - let wrapped = wrap_native_incremental_llm_stream_execution_v2_with_user_data( + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( return_v2_stream_state_and_release_handles, user_data, ); @@ -3009,7 +3135,7 @@ fn native_api_v2_stream_wrapper_enforces_callback_and_ownership_contracts() { }; let user_data = test_v2_callback_user_data((&state as *const V2StreamContractState).cast_mut().cast()); - let wrapped = wrap_native_incremental_llm_stream_execution_v2_with_user_data( + let wrapped = wrap_native_incremental_llm_stream_execution_with_user_data( return_v2_stream_state_and_release_handles, user_data, ); @@ -4164,6 +4290,9 @@ fn native_async_completion_reject_covers_null_invalid_and_abort_paths() { 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), @@ -4272,6 +4401,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), @@ -4328,6 +4460,10 @@ fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuation 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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -4512,6 +4648,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), @@ -4554,6 +4693,10 @@ 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::mpsc::channel(1); let output_stream = Arc::new(NativeAsyncStream { + runtime: native_runtime().handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -4630,6 +4773,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), @@ -4681,6 +4827,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), @@ -4711,6 +4860,10 @@ 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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(true), settled: AtomicBool::new(false), @@ -4774,6 +4927,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), @@ -4821,6 +4977,10 @@ fn native_async_stream_next_supports_repeated_concurrent_calls() { let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -4937,6 +5097,10 @@ fn native_async_stream_settlement_rejects_late_next_and_aborts_in_flight_next() let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5057,6 +5221,10 @@ fn native_async_stream_next_stops_callbacks_after_false() { let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5130,6 +5298,10 @@ fn native_async_stream_in_flight_cancellation_releases_callback_state() { let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5211,6 +5383,10 @@ fn native_async_stream_cancellation_before_first_poll_releases_callback_state() let next_ref = Arc::into_raw(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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5279,6 +5455,385 @@ fn native_async_stream_cancellation_before_first_poll_releases_callback_state() } } +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; + assert_eq!( + unsafe { native_async_task_wake_v2(task) }, + NemoRelayStatus::Ok + ); + 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), + backpressure_waiter: 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_stream_spawn_task_v2( + stream_ref, + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + stream_started_rx + .recv_timeout(Duration::from_secs(1)) + .expect("stream task should poll once"); + drop(NativeAsyncStreamReceiver { + receiver, + stream: Arc::clone(&stream), + }); + 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_task_duplicate_spawn_does_not_consume_user_data() { + 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 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_completion_spawn_task_v2( + completion_ref, + poll_cooperative_task, + first.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + 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, + completed: false, + }); + wait_for_task_free(&runtime, &first_frees); +} + +#[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_task_stream_backpressure_wakes_after_consumer_drain() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .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(()), + runtime: runtime.handle().clone(), + context: MiddlewareContinuationContext::capture(), + task: Mutex::new(None), + backpressure_waiter: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }); + let stream_ref = Arc::into_raw(Arc::clone(&stream)) as *const NemoRelayNativeAsyncStream; + 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_stream_spawn_task_v2( + stream_ref, + poll_cooperative_task, + state.cast(), + Some(free_cooperative_task), + ) + }, + NemoRelayStatus::Ok + ); + 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] + }); + assert_eq!(chunks, [json!({"chunk": 1}), json!({"chunk": 2})]); + wait_for_task_free(&runtime, &frees); +} + #[test] fn native_async_completion_abi_rejects_invalid_duplicate_and_cancelled_settlement() { let runtime = tokio::runtime::Builder::new_current_thread() @@ -5287,6 +5842,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), @@ -5322,6 +5880,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), @@ -5348,6 +5909,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), @@ -5402,6 +5966,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), @@ -5483,6 +6050,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), @@ -5599,47 +6169,6 @@ fn native_async_callback_contract_errors_abort_an_invoked_next() { } } -#[test] -fn cancelling_blocking_v2_callback_reclaims_host_resources_after_it_returns() { - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(2) - .enable_all() - .build() - .unwrap(); - let (started_tx, started_rx) = std::sync::mpsc::channel(); - let (release_tx, release_rx) = std::sync::mpsc::channel(); - let (freed_tx, freed_rx) = std::sync::mpsc::channel(); - let callback_state = Box::into_raw(Box::new(BlockingSafeCallbackState { - started: started_tx, - release: Mutex::new(release_rx), - freed: freed_tx, - })); - let user_data = Arc::new(NativeCallbackUserData { - ptr: callback_state.cast(), - free_fn: Some(free_blocking_safe_callback), - _instance: None, - }); - let task = runtime.spawn(invoke_native_async_callback_blocking( - blocking_safe_callback, - user_data, - json!({"name": "cancelled"}), - Some(NativeAsyncNextInner::Llm(Arc::new(|_| { - Box::pin(async { std::future::pending::>().await }) - }))), - )); - - started_rx - .recv_timeout(Duration::from_secs(1)) - .expect("blocking callback should start"); - task.abort(); - assert!(runtime.block_on(task).unwrap_err().is_cancelled()); - runtime.shutdown_timeout(Duration::from_millis(1)); - release_tx.send(()).unwrap(); - freed_rx - .recv_timeout(Duration::from_secs(1)) - .expect("detached cleanup must release completion, next, and callback state"); -} - #[test] fn native_async_stream_contract_errors_abort_an_invoked_next() { struct DropSignal(std::sync::mpsc::Sender<()>); @@ -5798,6 +6327,10 @@ 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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5866,6 +6399,10 @@ fn native_async_stream_push_is_bounded_retryable_and_incremental() { .unwrap(); 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), + backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -6630,22 +7167,31 @@ fn native_registration_entrypoints_reject_null_contexts() { NemoRelayStatus::NullPointer ); assert_eq!( - native_plugin_context_register_async_llm_execution_v2( + 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(), - 0, - return_v2_callback_state_and_release_next, + 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_async_llm_stream_execution_v2( + native_plugin_context_register_subscriber( ptr::null_mut(), ptr::null(), - 0, - return_v2_stream_state_and_release_handles, + noop_subscriber, ptr::null_mut(), None, ), diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 6d6440f39..35b20fc2f 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -166,14 +166,20 @@ 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. Safe callback futures -and returned streams run to completion on Relay's reusable blocking callback -lane, outside an entered Tokio runtime. No Rust future, trait object, -`serde_json::Value`, or allocator-owned Rust string crosses the C ABI boundary. - -The raw `PluginContext::host_api_v4` table and `_raw` v2 registration methods -remain available for advanced ABI consumers and non-Rust bindings. Code using -them is responsible for every callback lifetime, host string, completion, +`Stream` and cancels unfinished provider work on drop. Safe callbacks register +through the generic V3 completion API 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. + +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: @@ -182,6 +188,12 @@ 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 ebdc1f690..e88d27a48 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -881,6 +881,28 @@ 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. +#[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 cancels @@ -1227,7 +1249,8 @@ pub struct NemoRelayNativeHostApiV3 { ) -> NemoRelayStatus, } -/// ABI-v4 host extension implementing native API v2 targeted LLM continuations. +/// 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. @@ -1267,37 +1290,37 @@ pub struct NemoRelayNativeHostApiV4 { /// Releases the plugin-owned provider-stream reference. pub async_llm_stream_release_v2: unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2), - /// Registers a native API v2 unary LLM execution callback. + /// Starts a cooperative task associated with an async completion. /// - /// Relay invokes the callback on its reusable blocking executor with the - /// active scope stack bound to that worker. Provider continuations invoked - /// through this table continue to execute on Relay's Tokio runtime. Once - /// invoked, the host consumes `user_data` and calls `free_fn` exactly once, - /// including when registration fails. - pub plugin_context_register_async_llm_execution_v2: unsafe extern "C" fn( - ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncMiddlewareCb, + /// 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, - /// Registers a native API v2 incremental LLM stream execution callback. + ) -> NemoRelayStatus, + /// Starts a cooperative task associated with an incremental output stream. /// - /// Relay starts the callback on its reusable blocking executor and exposes - /// the bounded output stream to the caller concurrently. Once invoked, the - /// host consumes `user_data` and calls `free_fn` exactly once, including - /// when registration fails. - pub plugin_context_register_async_llm_stream_execution_v2: - unsafe extern "C" fn( - ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncStreamMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, - ) -> NemoRelayStatus, + /// 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) -> NemoRelayStatus, + /// 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. /// @@ -2879,56 +2902,6 @@ impl<'a> PluginContext<'a> { }) } - /// Registers a native API v2 unary LLM execution callback. - /// - /// # Safety - /// The host consumes `user_data` as soon as the registration function is - /// invoked and calls `free_fn` exactly once on registration failure or - /// eventual deregistration. The callback owns its completion and `next` - /// handles and must settle/release them exactly once. - pub unsafe fn register_async_llm_execution_v2_raw( - &mut self, - name: &str, - priority: i32, - cb: NemoRelayNativeAsyncMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, - ) -> NemoRelayStatus { - let Some(host) = self.host_api_v4() else { - return NemoRelayStatus::InvalidArg; - }; - self.with_name(name, |_, name| unsafe { - (host.plugin_context_register_async_llm_execution_v2)( - self.raw, name, priority, cb, user_data, free_fn, - ) - }) - } - - /// Registers a native API v2 incremental LLM stream execution callback. - /// - /// # Safety - /// The host consumes `user_data` as soon as the registration function is - /// invoked and calls `free_fn` exactly once on registration failure or - /// eventual deregistration. Callback-owned `next` and stream handles must - /// each be released exactly once. - pub unsafe fn register_async_llm_stream_execution_v2_raw( - &mut self, - name: &str, - priority: i32, - cb: NemoRelayNativeAsyncStreamMiddlewareCb, - user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, - ) -> NemoRelayStatus { - let Some(host) = self.host_api_v4() else { - return NemoRelayStatus::InvalidArg; - }; - self.with_name(name, |_, name| unsafe { - (host.plugin_context_register_async_llm_stream_execution_v2)( - self.raw, name, priority, cb, user_data, free_fn, - ) - }) - } - fn with_name( &self, name: &str, diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index 1b0eaf16f..6598b6f28 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -7,12 +7,11 @@ use std::ffi::c_void; use std::future::Future; use std::pin::Pin; use std::ptr; -use std::sync::{Arc, Condvar, Mutex}; +use std::sync::Arc; use std::task::{Context, Poll}; -use std::thread; -use std::time::Duration; use futures::channel::oneshot; +use futures::future::FutureExt; use futures::task::{ArcWake, waker_ref}; use futures::{Stream, StreamExt}; use serde::Deserialize; @@ -21,14 +20,12 @@ use super::{ HostString, Json, LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationStreamEventV2, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, NemoRelayNativeAsyncCallbackState, - NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, - NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV4, NemoRelayNativeLlmStreamV2, - NemoRelayNativeString, NemoRelayStatus, PluginContext, Result, read_json_value, - read_required_host_string, set_last_error, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncStream, NemoRelayNativeAsyncTaskV2, NemoRelayNativeHostApiV1, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmStreamV2, NemoRelayNativeString, NemoRelayStatus, + PluginContext, Result, read_json_value, read_required_host_string, set_last_error, }; -const BACKPRESSURE_RETRY_DELAY: Duration = Duration::from_millis(1); -const CANCELLATION_POLL_DELAY: Duration = Duration::from_millis(10); const MAX_SDK_ERROR_BYTES: usize = 4 * 1024; /// Asynchronous JSON stream returned by a safe native API v2 stream callback. @@ -86,62 +83,228 @@ struct ContinuationInner { next: *const NemoRelayNativeAsyncNext, } -// Do not use `thread::current()` here. Safe callbacks run on reusable host -// threads, and Rust's thread handle installs a TLS destructor in this plugin -// library. Relay may unload the library before that host thread exits. -struct BlockingCallbackWaker { - notified: Mutex, - ready: Condvar, +enum HostTaskCancellation { + Completion(*const NemoRelayNativeAsyncCompletion), + Stream(*const NemoRelayNativeAsyncStream), } -impl BlockingCallbackWaker { - fn wait_timeout(&self, timeout: Duration) { - let notified = self - .notified - .lock() - .unwrap_or_else(|error| error.into_inner()); - let (mut notified, _) = self - .ready - .wait_timeout_while(notified, timeout, |notified| !*notified) - .unwrap_or_else(|error| error.into_inner()); - *notified = false; +#[derive(Clone, Copy)] +struct CompletionHandle(*const NemoRelayNativeAsyncCompletion); + +#[derive(Clone, Copy)] +struct OutputHandle(*const NemoRelayNativeAsyncStream); + +struct TaskHostString { + host: NemoRelayNativeHostApiV1, + raw: usize, +} + +// 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 TaskHostString { + fn new(host: &NemoRelayNativeHostApiV1, value: &str) -> Option { + let mut raw = ptr::null_mut(); + let status = unsafe { (host.string_new)(value.as_ptr(), value.len(), &mut raw) }; + if status != NemoRelayStatus::Ok || raw.is_null() { + return None; + } + Some(Self { + host: *host, + raw: raw as usize, + }) + } + + fn from_json(host: &NemoRelayNativeHostApiV1, value: &Json) -> Option { + Self::new(host, &serde_json::to_string(value).ok()?) + } + + fn as_ptr(&self) -> *const NemoRelayNativeString { + self.raw as *const NemoRelayNativeString } } -impl ArcWake for BlockingCallbackWaker { - fn wake_by_ref(arc_self: &Arc) { - let mut notified = arc_self - .notified - .lock() - .unwrap_or_else(|error| error.into_inner()); - *notified = true; - arc_self.ready.notify_one(); +impl Drop for TaskHostString { + fn drop(&mut self) { + unsafe { (self.host.string_free)(self.raw as *mut NemoRelayNativeString) }; } } -fn block_on_complete_callback(future: F, is_cancelled: C) -> Option -where - F: Future, - C: Fn() -> bool, -{ - let mut future = Box::pin(future); - let callback_waker = Arc::new(BlockingCallbackWaker { - notified: Mutex::new(false), - ready: Condvar::new(), - }); - let waker = waker_ref(&callback_waker); - let mut context = Context::from_waker(&waker); - loop { - if is_cancelled() { +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, + ) -> Option> { + if raw.is_null() { return None; } - match future.as_mut().poll(&mut context) { - Poll::Ready(output) => return Some(output), - Poll::Pending => callback_waker.wait_timeout(CANCELLATION_POLL_DELAY), + unsafe { (host.async_task_retain_v2)(raw) }; + Some(Arc::new(Self { host, raw })) + } +} + +impl ArcWake for HostTaskWaker { + fn wake_by_ref(arc_self: &Arc) { + let status = unsafe { (arc_self.host.async_task_wake_v2)(arc_self.raw) }; + if status != NemoRelayStatus::Ok { + set_last_error( + &arc_self.host.v3.v1, + &format!("native API v2 host task wake failed: {status:?}"), + ); + } + } +} + +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() { + let waker = match unsafe { HostTaskWaker::new(state.host, task) } { + Some(waker) => waker, + None => { + set_last_error(&state.host.v3.v1, "native API v2 host task was null"); + return NemoRelayNativeAsyncCallbackState::Complete as u32; + } + }; + state.waker = Some(waker); + } + 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 {} @@ -435,7 +598,7 @@ struct LlmCallbackInvocation { struct SafeV2Callback { host: NemoRelayNativeHostApiV4, - callback: F, + callback: Arc, } unsafe extern "C" fn drop_safe_v2_callback(user_data: *mut c_void) { @@ -453,9 +616,8 @@ impl PluginContext<'_> { /// Registers a safe asynchronous native API v2 buffered LLM execution callback. /// /// The callback receives owned Rust values and a cloneable continuation. - /// Its future is driven by a cancellation-aware local executor on Relay's - /// reusable blocking callback lane. It is not polled inside an entered - /// Tokio runtime; runtime-specific APIs should not be assumed. + /// Relay cooperatively polls its future on the host runtime with the + /// invocation's active scope stack restored for every poll. pub fn register_async_llm_execution_v2( &mut self, name: &str, @@ -473,16 +635,21 @@ impl PluginContext<'_> { 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 }); + 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.plugin_context_register_async_llm_execution_v2)( + (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::), @@ -502,8 +669,8 @@ impl PluginContext<'_> { /// Registers a safe asynchronous native API v2 streaming LLM callback. /// /// The callback may return a Rust stream or request host-owned direct - /// pass-through. Its future and returned stream are driven on Relay's - /// reusable blocking callback lane, outside an entered Tokio runtime. + /// pass-through. Relay cooperatively polls its future and returned stream + /// and wakes them when bounded output backpressure clears. pub fn register_async_llm_stream_execution_v2( &mut self, name: &str, @@ -521,12 +688,13 @@ impl PluginContext<'_> { 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 }); + let state = Box::new(SafeV2Callback { + host, + callback: Arc::new(callback), + }); let user_data = Box::into_raw(state).cast::(); let status = unsafe { - // The V4 host has the same consume-on-invocation ownership - // contract as buffered registration. - (host.plugin_context_register_async_llm_stream_execution_v2)( + (host.v3.plugin_context_register_async_stream_middleware)( self.raw, name.as_ptr(), priority, @@ -578,30 +746,72 @@ where return NemoRelayNativeAsyncCallbackState::Complete as u32; } }; - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let invocation: LlmCallbackInvocation = read_json_value( - &state.host.v3.v1, - invocation_json, - "native API v2 LLM invocation", - ) - .map_err(|status| format!("invalid native API v2 LLM invocation: {status:?}"))?; - let execution = (state.callback)(invocation.name, invocation.request, continuation.clone()); - block_on_complete_callback(execution, || unsafe { - (state.host.v3.async_completion_is_cancelled)(completion) + 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 callback_continuation = continuation.clone(); + let future = async move { + let result = std::panic::AssertUnwindSafe(async move { + callback(invocation.name, invocation.request, callback_continuation).await }) - .unwrap_or_else(|| Err("native API v2 callback was cancelled".into())) - })); - match result { - Ok(Ok(value)) => resolve_completion(&state.host, completion, &value), - Ok(Err(error)) => reject_completion(&state.host, completion, &error), - Err(_) => reject_completion( - &state.host, + .catch_unwind() + .await; + drop(continuation); + 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, - "native API v2 buffered LLM callback panicked", - ), + 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 } - drop(continuation); - NemoRelayNativeAsyncCallbackState::Complete as u32 } unsafe extern "C" fn safe_streaming_trampoline( @@ -628,7 +838,7 @@ where { Ok(continuation) => continuation, Err(status) => { - reject_output( + reject_output_once( &state.host, output, &format!("invalid native API v2 stream continuation: {status:?}"), @@ -637,16 +847,30 @@ where return NemoRelayNativeAsyncCallbackState::Complete as u32; } }; - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let invocation: LlmCallbackInvocation = read_json_value( - &state.host.v3.v1, - invocation_json, - "native API v2 streaming LLM invocation", - ) - .map_err(|status| format!("invalid native API v2 stream invocation: {status:?}"))?; - let execution = async { + 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 = - (state.callback)(invocation.name, invocation.request, continuation.clone()).await?; + callback(invocation.name, invocation.request, callback_continuation).await?; match outcome { LlmStreamExecutionOutcomeV2::Stream(stream) => { pump_output_stream(&continuation, stream).await @@ -655,23 +879,49 @@ where continuation.forward_passthrough(request).await } } - }; - block_on_complete_callback(execution, || unsafe { - (state.host.v3.async_stream_is_cancelled)(continuation.inner.output) }) - .unwrap_or_else(|| Err("native API v2 output stream was cancelled".into())) - })); - match result { - Ok(Ok(())) => {} - Ok(Err(error)) => reject_output(&state.host, output, &error), - Err(_) => reject_output( - &state.host, + .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, - "native API v2 streaming LLM callback panicked", - ), + 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 } - drop(continuation); - NemoRelayNativeAsyncCallbackState::Complete as u32 } async fn pump_output_stream( @@ -680,30 +930,31 @@ async fn pump_output_stream( ) -> Result<()> { while let Some(item) = stream.next().await { let chunk = item?; - push_output_json(continuation, &chunk)?; + push_output_json(continuation, &chunk).await?; } finish_output(continuation) } -fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { +async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { let host = &continuation.inner.continuation.host; - let chunk = HostString::from_json(&host.v3.v1, chunk) + let chunk = TaskHostString::from_json(&host.v3.v1, chunk) .ok_or_else(|| "failed to serialize native API v2 output chunk".to_string())?; - loop { + std::future::poll_fn(move |_context| { if unsafe { (host.v3.async_stream_is_cancelled)(continuation.inner.output) } { - return Err("native API v2 output stream was cancelled".into()); + 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 => return Ok(()), + NemoRelayStatus::Ok => Poll::Ready(Ok(())), // For this operation the V3 ABI contract reserves `Internal` for // a full bounded queue. Serialization and lifecycle faults use // distinct statuses, so retrying cannot mask another ABI error. - NemoRelayStatus::Internal => thread::sleep(BACKPRESSURE_RETRY_DELAY), - status => return Err(format!("native API v2 output push failed: {status:?}")), + NemoRelayStatus::Internal => Poll::Pending, + status => Poll::Ready(Err(format!("native API v2 output push failed: {status:?}"))), } - } + }) + .await } fn finish_output(continuation: &LlmStreamContinuationV2) -> Result<()> { @@ -719,39 +970,61 @@ fn finish_output(continuation: &LlmStreamContinuationV2) -> Result<()> { } } -fn reject_output( - host: &NemoRelayNativeHostApiV4, - output: *const NemoRelayNativeAsyncStream, - error: &str, -) { - if unsafe { (host.v3.async_stream_is_cancelled)(output) } { +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 { + let Some(message) = TaskHostString::new(&host.v3.v1, &error) else { set_last_error( &host.v3.v1, "failed to allocate native API v2 stream rejection", ); return; }; - loop { - if unsafe { (host.v3.async_stream_is_cancelled)(output) } { - 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, message.as_ptr()) }; + let status = unsafe { (host.v3.async_stream_reject)(output.as_ptr(), message.as_ptr()) }; match status { - NemoRelayStatus::Ok => return, + NemoRelayStatus::Ok => Poll::Ready(()), // `Internal` has the same queue-full-only contract for rejection. - NemoRelayStatus::Internal => thread::sleep(BACKPRESSURE_RETRY_DELAY), + NemoRelayStatus::Internal => Poll::Pending, status => { set_last_error( &host.v3.v1, &format!("native API v2 output rejection failed: {status:?}"), ); - return; + 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:?}"), + ); } } diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 7ad20e349..9880f9251 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -27,10 +27,11 @@ use nemo_relay_plugin::{ NemoRelayNativeAsyncLlmStreamNextCbV2, NemoRelayNativeAsyncLlmStreamOpenCbV2, NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, - NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, - NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, - NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeAsyncTaskPollCbV2, + NemoRelayNativeAsyncTaskV2, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, + NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativeHostApiV4, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, + NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, @@ -389,6 +390,7 @@ 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_CANCELS: AtomicUsize = AtomicUsize::new(0); static SAFE_V2_PROVIDER_RELEASES: AtomicUsize = AtomicUsize::new(0); @@ -409,6 +411,24 @@ static SAFE_V2_OUTPUT_PUSH_STATUSES: Mutex> = Mutex::n 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() { @@ -449,10 +469,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 504); + assert_eq!(size_of::(), 528); assert_eq!( host_api_v4_offsets(), - [0, 440, 448, 456, 464, 472, 480, 488, 496] + [0, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512, 520] ); assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 56); @@ -483,10 +503,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 248); + assert_eq!(size_of::(), 260); assert_eq!( host_api_v4_offsets(), - [0, 216, 220, 224, 228, 232, 236, 240, 244] + [0, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256] ); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); @@ -497,7 +517,7 @@ fn native_abi_v3_struct_sizes_are_self_describing() { } } -fn host_api_v4_offsets() -> [usize; 9] { +fn host_api_v4_offsets() -> [usize; 12] { [ offset_of!(NemoRelayNativeHostApiV4, v3), offset_of!(NemoRelayNativeHostApiV4, async_llm_next_invoke_result_v2), @@ -505,14 +525,11 @@ fn host_api_v4_offsets() -> [usize; 9] { offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_next_v2), offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_cancel_v2), offset_of!(NemoRelayNativeHostApiV4, async_llm_stream_release_v2), - offset_of!( - NemoRelayNativeHostApiV4, - plugin_context_register_async_llm_execution_v2 - ), - offset_of!( - NemoRelayNativeHostApiV4, - plugin_context_register_async_llm_stream_execution_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), ] } @@ -1438,6 +1455,10 @@ fn reset_state() { 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, @@ -1475,6 +1496,7 @@ fn reset_state() { 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_CANCELS.store(0, Ordering::SeqCst); SAFE_V2_PROVIDER_RELEASES.store(0, Ordering::SeqCst); @@ -1494,6 +1516,10 @@ fn reset_state() { 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<'_> { @@ -6140,6 +6166,7 @@ unsafe extern "C" fn safe_v2_completion_is_cancelled( 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( @@ -6156,15 +6183,48 @@ unsafe extern "C" fn safe_v2_next_release(_next: *const NemoRelayNativeAsyncNext unsafe extern "C" fn safe_v2_register_generic_async( _ctx: *mut NemoRelayNativePluginContext, - _kind: u32, - _name: *const NemoRelayNativeString, - _priority: i32, + kind: u32, + name: *const NemoRelayNativeString, + priority: i32, _break_chain: bool, - _cb: NemoRelayNativeAsyncMiddlewareCb, - _user_data: *mut c_void, - _free_fn: NemoRelayNativeFreeFn, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { - NemoRelayStatus::InvalidArg + if kind != NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept as u32 { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return NemoRelayStatus::InvalidArg; + } + let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + }; + 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( @@ -6186,6 +6246,13 @@ unsafe extern "C" fn safe_v2_stream_push( .unwrap_or(NemoRelayStatus::Ok); if status == NemoRelayStatus::Ok { SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); + } else if status == NemoRelayStatus::Internal { + SAFE_V2_CURRENT_TASK.with(|current| { + let task = current.get(); + if task != 0 { + let _ = unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + } + }); } status } @@ -6211,6 +6278,13 @@ unsafe extern "C" fn safe_v2_stream_reject( .unwrap_or(NemoRelayStatus::Ok); if status == NemoRelayStatus::Ok { SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); + } else if status == NemoRelayStatus::Internal { + SAFE_V2_CURRENT_TASK.with(|current| { + let task = current.get(); + if task != 0 { + let _ = unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + } + }); } status } @@ -6237,13 +6311,40 @@ unsafe extern "C" fn safe_v2_async_next_invoke_stream( 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, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeAsyncStreamMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { - NemoRelayStatus::InvalidArg + let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); + if status != NemoRelayStatus::Ok { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return status; + } + }; + 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( @@ -6386,72 +6487,123 @@ unsafe extern "C" fn safe_v2_provider_release(_stream: *const NemoRelayNativeLlm SAFE_V2_PROVIDER_RELEASES.fetch_add(1, Ordering::SeqCst); } -unsafe extern "C" fn safe_v2_register_buffered( - _ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncMiddlewareCb, +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_REGISTRATION_STATUS.lock().unwrap(); + let status = *SAFE_V2_TASK_SPAWN_STATUS.lock().unwrap(); if status != NemoRelayStatus::Ok { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } return status; } - let host = test_host(); - let name = match required_host_string(&host, name) { - Ok(name) => name, - Err(status) => return status, - }; - replace_registration( - &ASYNC_V2_REGISTRATION, - RegisteredAsyncV2 { - name, - priority, - cb, - user_data: user_data as usize, - free_fn, - }, - ); + 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_register_streaming( - _ctx: *mut NemoRelayNativePluginContext, - name: *const NemoRelayNativeString, - priority: i32, - cb: NemoRelayNativeAsyncStreamMiddlewareCb, +unsafe extern "C" fn safe_v2_completion_spawn_task( + completion: *const NemoRelayNativeAsyncCompletion, + cb: NemoRelayNativeAsyncTaskPollCbV2, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { - let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); - if status != NemoRelayStatus::Ok { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return status; + 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); } - let host = test_host(); - let name = match required_host_string(&host, name) { - Ok(name) => name, - Err(status) => return status, +} + +unsafe extern "C" fn safe_v2_task_wake(task: *const NemoRelayNativeAsyncTaskV2) -> NemoRelayStatus { + let Some(task) = (unsafe { task.cast::().as_ref() }) else { + return NemoRelayStatus::NullPointer; }; - replace_registration( - &ASYNC_STREAM_V2_REGISTRATION, - RegisteredAsyncStreamV2 { - name, - priority, - cb, - user_data: user_data as usize, - free_fn, - }, - ); + if task.completed.load(Ordering::Acquire) { + return NemoRelayStatus::Ok; + } + task.woken.store(true, Ordering::Release); NemoRelayStatus::Ok } +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() { + let _ = 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, @@ -6511,8 +6663,11 @@ fn test_host_v4() -> NemoRelayNativeHostApiV4 { async_llm_stream_next_v2: safe_v2_provider_next, async_llm_stream_cancel_v2: safe_v2_provider_cancel, async_llm_stream_release_v2: safe_v2_provider_release, - plugin_context_register_async_llm_execution_v2: safe_v2_register_buffered, - plugin_context_register_async_llm_stream_execution_v2: safe_v2_register_streaming, + 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, } } @@ -6593,6 +6748,11 @@ where NonNull::::dangling().as_ptr(), ); unsafe { registration.free() }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); state } @@ -6612,6 +6772,11 @@ where NonNull::::dangling().as_ptr(), ); unsafe { registration.free() }; + assert_eq!( + NemoRelayNativeAsyncCallbackState::try_from(state), + Ok(NemoRelayNativeAsyncCallbackState::Pending) + ); + drive_safe_v2_tasks(); state } @@ -6650,16 +6815,18 @@ unsafe extern "C" fn raw_v2_streaming_probe( } #[test] -fn native_api_v2_raw_registration_remains_an_advanced_escape_hatch() { +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_llm_execution_v2_raw( + ctx.register_async_middleware_raw( + NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept, "raw-buffered", 11, + false, raw_v2_buffered_probe, ptr::null_mut(), None, @@ -6676,7 +6843,7 @@ fn native_api_v2_raw_registration_remains_an_advanced_escape_hatch() { assert_eq!( unsafe { - ctx.register_async_llm_stream_execution_v2_raw( + ctx.register_async_stream_middleware_raw( "raw-streaming", 12, raw_v2_streaming_probe, @@ -6697,9 +6864,11 @@ fn native_api_v2_raw_registration_remains_an_advanced_escape_hatch() { let mut v1_ctx = test_context(&v1); assert_eq!( unsafe { - v1_ctx.register_async_llm_execution_v2_raw( + v1_ctx.register_async_middleware_raw( + NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept, "unsupported-buffered", 0, + false, raw_v2_buffered_probe, ptr::null_mut(), None, @@ -6709,7 +6878,7 @@ fn native_api_v2_raw_registration_remains_an_advanced_escape_hatch() { ); assert_eq!( unsafe { - v1_ctx.register_async_llm_stream_execution_v2_raw( + v1_ctx.register_async_stream_middleware_raw( "unsupported-streaming", 0, raw_v2_streaming_probe, @@ -6760,8 +6929,9 @@ fn safe_v2_buffered_registration_wraps_targeted_and_passthrough_calls() { unsafe { (host.v3.v1.string_free)(invocation) }; assert_eq!( NemoRelayNativeAsyncCallbackState::try_from(state), - Ok(NemoRelayNativeAsyncCallbackState::Complete) + Ok(NemoRelayNativeAsyncCallbackState::Pending) ); + drive_safe_v2_tasks(); assert_eq!( SAFE_V2_COMPLETION.lock().unwrap().take(), Some(Ok(json!({ @@ -6770,6 +6940,71 @@ fn safe_v2_buffered_registration_wraps_targeted_and_passthrough_calls() { }))) ); 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_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); } @@ -6787,36 +7022,84 @@ fn safe_v2_buffered_callback_stops_when_the_caller_cancels() { let invocation = json_host_string( &host.v3.v1, json!({ "name": "managed", "request": test_llm_request() }), - ) as usize; - let user_data = registration.user_data; - let callback = registration.cb; - let (done_tx, done_rx) = std::sync::mpsc::channel(); - let callback_thread = std::thread::spawn(move || { - let state = unsafe { - callback( - user_data as *mut c_void, - invocation as *const NemoRelayNativeString, - NonNull::::dangling().as_ptr(), - NonNull::::dangling().as_ptr(), - ) - }; - done_tx.send(state).unwrap(); - }); - - std::thread::sleep(std::time::Duration::from_millis(20)); + ); + 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); - let state = done_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .expect("caller cancellation must wake a pending safe buffered callback"); - callback_thread.join().unwrap(); - unsafe { (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString) }; + 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::Complete) + 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!(SAFE_V2_COMPLETION.lock().unwrap().is_none()); + 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); } @@ -6840,34 +7123,27 @@ fn safe_v2_cancelled_targeted_call_releases_next_before_the_host_callback() { let invocation = json_host_string( &host.v3.v1, json!({ "name": "managed", "request": test_llm_request() }), - ) as usize; - let user_data = registration.user_data; - let callback = registration.cb; - let (done_tx, done_rx) = std::sync::mpsc::channel(); - let callback_thread = std::thread::spawn(move || { - let state = unsafe { - callback( - user_data as *mut c_void, - invocation as *const NemoRelayNativeString, - NonNull::::dangling().as_ptr(), - NonNull::::dangling().as_ptr(), - ) - }; - done_tx.send(state).unwrap(); - }); + ); + 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); - let state = done_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .expect("cancelling the managed call must stop the targeted callback"); - callback_thread.join().unwrap(); - - assert_eq!( - NemoRelayNativeAsyncCallbackState::try_from(state), - Ok(NemoRelayNativeAsyncCallbackState::Complete) - ); + wake_safe_v2_tasks(); + drive_safe_v2_tasks(); assert_eq!( SAFE_V2_NEXT_RELEASES.load(Ordering::SeqCst), 1, @@ -6893,7 +7169,6 @@ fn safe_v2_cancelled_targeted_call_releases_next_before_the_host_callback() { unsafe { targeted_callback(targeted_user_data as *mut c_void, outcome) }; unsafe { (host.v3.v1.string_free)(outcome); - (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString); registration.free(); } assert_eq!(live_host_strings(), 0); @@ -6954,8 +7229,9 @@ fn safe_v2_stream_registration_pumps_provider_stream_and_releases_handles() { unsafe { (host.v3.v1.string_free)(invocation) }; assert_eq!( NemoRelayNativeAsyncCallbackState::try_from(state), - Ok(NemoRelayNativeAsyncCallbackState::Complete) + Ok(NemoRelayNativeAsyncCallbackState::Pending) ); + drive_safe_v2_tasks(); assert_eq!( *SAFE_V2_OUTPUT.lock().unwrap(), vec![Ok(json!({ "delta": "hello" }))] @@ -6999,8 +7275,9 @@ fn safe_v2_stream_passthrough_uses_host_owned_forwarding() { unsafe { (host.v3.v1.string_free)(invocation) }; assert_eq!( NemoRelayNativeAsyncCallbackState::try_from(state), - Ok(NemoRelayNativeAsyncCallbackState::Complete) + 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); @@ -7049,6 +7326,7 @@ fn safe_v2_provider_stream_drop_cancels_unfinished_production() { ) }; unsafe { (host.v3.v1.string_free)(invocation) }; + drive_safe_v2_tasks(); assert_eq!(SAFE_V2_PROVIDER_CANCELS.load(Ordering::SeqCst), 1); assert_eq!(SAFE_V2_PROVIDER_RELEASES.load(Ordering::SeqCst), 1); unsafe { registration.free() }; @@ -7068,34 +7346,24 @@ fn safe_v2_stream_callback_stops_when_the_caller_cancels() { let invocation = json_host_string( &host.v3.v1, json!({ "name": "managed", "request": test_llm_request() }), - ) as usize; - let user_data = registration.user_data; - let callback = registration.cb; - let (done_tx, done_rx) = std::sync::mpsc::channel(); - let callback_thread = std::thread::spawn(move || { - let state = unsafe { - callback( - user_data as *mut c_void, - invocation as *const NemoRelayNativeString, - NonNull::::dangling().as_ptr(), - NonNull::::dangling().as_ptr(), - ) - }; - done_tx.send(state).unwrap(); - }); - - std::thread::sleep(std::time::Duration::from_millis(20)); - SAFE_V2_OUTPUT_CANCELLED.store(true, Ordering::SeqCst); - let state = done_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .expect("caller cancellation must wake a pending safe stream callback"); - callback_thread.join().unwrap(); - unsafe { (host.v3.v1.string_free)(invocation as *mut NemoRelayNativeString) }; - + ); + 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::Complete) + 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()); @@ -7152,6 +7420,7 @@ fn safe_v2_stream_open_preserves_structured_failure() { ) }; 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" @@ -7207,6 +7476,7 @@ fn safe_v2_malformed_stream_open_cancels_and_releases_the_owned_stream() { ) }; unsafe { (host.v3.v1.string_free)(invocation) }; + drive_safe_v2_tasks(); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), @@ -7289,6 +7559,63 @@ fn safe_v2_failed_host_registration_frees_callback_state_exactly_once() { 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(); @@ -7343,7 +7670,7 @@ fn safe_v2_malformed_invocations_settle_and_release_callback_handles() { } #[test] -fn safe_v2_executor_honors_wakes_without_plugin_thread_local_state() { +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 { @@ -7366,6 +7693,29 @@ fn safe_v2_executor_honors_wakes_without_plugin_thread_local_state() { 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(); diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 6fc9e1def..3ecc4e838 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -168,10 +168,19 @@ each use bounded 32-event queues. Dropping or cancelling a stream stops provider production, and the library stays loaded until all callbacks and streams release their handles. -V2 callbacks run on Relay's blocking executor with the active scope stack -restored. Continuations run on Relay's existing Tokio runtime; Relay does not -create an OS thread per continuation. Separate callback invocations can run -concurrently and have no stable OS-thread affinity. +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 completion API and return +`Pending`. 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` @@ -219,10 +228,11 @@ an unmanaged stream, return original downstream stream directly to the caller through its bounded queue, so pass-through events do not cross the plugin boundary. -The SDK drives safe callback futures and returned streams to completion on the -blocking callback lane. It does not enter Relay's Tokio runtime on that thread. -`LlmProviderStreamV2` implements `Stream`, enforces one pending pull, and -cancels unfinished provider production on drop. +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: @@ -252,13 +262,15 @@ extern "C" fn nemo_relay_register_plugin( ) -> NemoRelayStatus ``` -`PluginContext::host_api_v4` and the `_raw` v2 registration methods are the -advanced escape hatch for raw ABI consumers and non-Rust bindings. They expose -opaque callbacks and handles intentionally: the caller must own host strings, -completion settlement, stream backpressure, cancellation, panic fencing, and -every release operation. The safe Rust facade is implemented on top of these -same C functions; Rust futures, streams, trait objects, and allocator-owned -strings never cross the shared-library boundary. +`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 registers callbacks through V3's completion-based `Pending` +contract. 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 From 6b470b04eb88ad5d60e0fe02b6de887909246243 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 10:46:01 -0600 Subject: [PATCH 16/32] docs(plugin): require executor-neutral native futures Signed-off-by: Bryan Bednarski --- crates/plugin/README.md | 7 +++++++ crates/plugin/src/lib.rs | 4 ++++ crates/plugin/src/native_v2.rs | 5 +++++ 3 files changed, 16 insertions(+) diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 35b20fc2f..8c75bb473 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -175,6 +175,13 @@ 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 diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index e88d27a48..1736078dd 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -888,6 +888,10 @@ pub struct NemoRelayNativeAsyncStream { /// 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], diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index 6598b6f28..97eaf59c9 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -618,6 +618,9 @@ impl PluginContext<'_> { /// 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, @@ -671,6 +674,8 @@ impl PluginContext<'_> { /// 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, From c14057b979dc0b2c980e60c01926f769ce1b6ba4 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 14:10:20 -0600 Subject: [PATCH 17/32] refactor(plugin): simplify targeted LLM continuations Signed-off-by: Bryan Bednarski --- crates/cli/src/gateway/mod.rs | 32 +- .../src/api/runtime/continuation_context.rs | 2 +- .../src/api/runtime/llm_dispatch_context.rs | 118 +-- crates/core/src/error.rs | 50 + crates/core/src/plugin/dynamic/native.rs | 960 +++++++++-------- crates/core/tests/coverage/error_tests.rs | 33 + .../tests/fixtures/native_plugin/src/lib.rs | 11 +- .../tests/unit/llm_dispatch_context_tests.rs | 174 ++- crates/core/tests/unit/native_plugin_tests.rs | 994 +++++++++++++----- crates/plugin/src/lib.rs | 178 ++-- crates/plugin/src/native_v2.rs | 230 ++-- crates/plugin/tests/typed_callbacks.rs | 337 ++++-- 12 files changed, 1891 insertions(+), 1228 deletions(-) diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index ad9893d45..76feccd4c 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -797,17 +797,16 @@ async fn forward_upstream_request( .provider_credential_present(), crate::provider_auth::has_provider_credential(headers) ); - let effective = build_effective_dispatch_request( + let effective = effective_dispatch_request( body_bytes, headers, effective_request, url, - method, forwarding.source_route, ); let configured_auth_header = forwarding.configured_auth_header(effective.target_route); let mut upstream = http - .request(effective.method.clone(), &effective.url) + .request(method.clone(), &effective.url) .body(effective.body_bytes.clone()); for (name, value) in &effective.headers { if should_forward_request_header(name, &effective.headers) { @@ -831,7 +830,6 @@ async fn forward_upstream_request( struct EffectiveUpstreamRequest { body_bytes: Bytes, headers: HeaderMap, - method: Method, url: String, target_route: ProviderRoute, credential_policy: TargetCredentialPolicy, @@ -859,39 +857,20 @@ fn effective_upstream_request( (effective.body_bytes, effective.headers) } -#[cfg(test)] fn effective_dispatch_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, url: &str, route: ProviderRoute, -) -> EffectiveUpstreamRequest { - build_effective_dispatch_request( - body_bytes, - headers, - effective_request, - url, - &Method::POST, - route, - ) -} - -fn build_effective_dispatch_request( - body_bytes: &Bytes, - headers: &HeaderMap, - effective_request: Option<&LlmRequest>, - url: &str, - method: &Method, - route: ProviderRoute, ) -> EffectiveUpstreamRequest { let mut headers = headers.clone(); strip_internal_dispatch_headers(&mut headers); let Some(request) = effective_request else { - return source_request(body_bytes, headers, method, url, route); + return source_request(body_bytes, headers, url, route); }; let Some((body_bytes, body_reencoded)) = reencode_request_body(request, body_bytes) else { - return source_request(body_bytes, headers, method, url, route); + return source_request(body_bytes, headers, url, route); }; let overrides = dispatch_overrides(&request.headers); let credential_policy = if overrides.is_explicit_target() { @@ -910,7 +889,6 @@ fn build_effective_dispatch_request( EffectiveUpstreamRequest { body_bytes, headers, - method: method.clone(), url: overrides.resolve_url(url), target_route: overrides.route.unwrap_or(route), credential_policy, @@ -920,14 +898,12 @@ fn build_effective_dispatch_request( fn source_request( body_bytes: &Bytes, headers: HeaderMap, - method: &Method, url: &str, route: ProviderRoute, ) -> EffectiveUpstreamRequest { EffectiveUpstreamRequest { body_bytes: body_bytes.clone(), headers, - method: method.clone(), url: url.to_string(), target_route: route, credential_policy: TargetCredentialPolicy::SourceOrEnvironment, diff --git a/crates/core/src/api/runtime/continuation_context.rs b/crates/core/src/api/runtime/continuation_context.rs index 2d0aba777..d21c2a4fa 100644 --- a/crates/core/src/api/runtime/continuation_context.rs +++ b/crates/core/src/api/runtime/continuation_context.rs @@ -117,7 +117,7 @@ impl MiddlewareContinuationContext { C: FnOnce() -> F, F: Future, { - scope_llm_dispatch_target(target, self.invoke(callback)).await + scope_llm_dispatch_target(self.active_event_uuid, target, self.invoke(callback)).await } /// Invoke a callback and poll its future with the captured Relay context. diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 02bdf480c..dc68e3a70 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -12,22 +12,31 @@ use std::time::Duration; use async_stream::stream; use futures_util::StreamExt; use reqwest::header::{self, HeaderMap, HeaderName, HeaderValue}; -use reqwest::{Client, Method, StatusCode, Url}; +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, Result, UpstreamFailure, UpstreamFailureClass}; +use crate::error::{ + FlowError, MAX_UPSTREAM_FAILURE_BODY_BYTES, Result, UpstreamFailure, UpstreamFailureClass, + 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_UPSTREAM_ERROR_BODY_BYTES: usize = 16 * 1024; -const MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES: usize = 1024; - tokio::task_local! { - static TASK_LLM_DISPATCH_TARGET: LlmDispatchTargetContext; + 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. @@ -37,26 +46,13 @@ tokio::task_local! { #[doc(hidden)] #[derive(Clone)] pub struct LlmDispatchTargetContext { - method: Method, url: Url, headers: HeaderMap, } impl LlmDispatchTargetContext { /// Validate and construct a target for one continuation invocation. - pub(crate) fn try_new( - method: String, - url: String, - headers: BTreeMap, - ) -> Result { - let method = Method::from_bytes(method.as_bytes()).map_err(|_| { - FlowError::InvalidArgument("LLM continuation method was invalid or prohibited".into()) - })?; - if matches!(method, Method::CONNECT | Method::TRACE) { - return Err(FlowError::InvalidArgument( - "LLM continuation method was invalid or prohibited".into(), - )); - } + 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() @@ -93,19 +89,11 @@ impl LlmDispatchTargetContext { .entry(header::CONTENT_TYPE) .or_insert(HeaderValue::from_static("application/json")); Ok(Self { - method, url, headers: validated_headers, }) } - /// HTTP method selected for this invocation. - #[doc(hidden)] - #[must_use] - pub(crate) fn method(&self) -> &Method { - &self.method - } - /// Absolute provider URL selected for this invocation. #[doc(hidden)] #[must_use] @@ -128,7 +116,6 @@ impl fmt::Debug for LlmDispatchTargetContext { redacted_url.set_fragment(None); formatter .debug_struct("LlmDispatchTargetContext") - .field("method", &self.method) .field("url", &redacted_url) .field( "header_names", @@ -166,15 +153,29 @@ fn prohibited_target_header(name: &HeaderName) -> bool { } pub(crate) fn current_llm_dispatch_target() -> Option { - TASK_LLM_DISPATCH_TARGET.try_with(Clone::clone).ok() + 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(target, future).await + 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. @@ -270,9 +271,7 @@ async fn send( ) -> Result { let body = serde_json::to_vec(&request.content) .map_err(|error| FlowError::InvalidArgument(error.to_string()))?; - let mut outbound = targeted_http_client() - .request(target.method().clone(), target.url().clone()) - .body(body); + let mut outbound = targeted_http_client().post(target.url().clone()).body(body); for (name, value) in target.headers() { outbound = outbound.header(name, value); } @@ -303,12 +302,12 @@ async fn bounded_response_body( ) -> Result> { let mut body = Vec::new(); let mut stream = response.bytes_stream(); - while body.len() < MAX_UPSTREAM_ERROR_BODY_BYTES { + 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_ERROR_BODY_BYTES - body.len(); + let remaining = MAX_UPSTREAM_FAILURE_BODY_BYTES - body.len(); body.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } Ok(body) @@ -341,7 +340,7 @@ fn transport_error(target: &LlmDispatchTargetContext, error: reqwest::Error) -> } fn http_error(status: StatusCode, headers: BTreeMap, body: &[u8]) -> FlowError { - let body = String::from_utf8_lossy(&body[..body.len().min(MAX_UPSTREAM_ERROR_BODY_BYTES)]); + 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(), @@ -355,45 +354,12 @@ fn http_error(status: StatusCode, headers: BTreeMap, body: &[u8] } fn safe_failure_headers(headers: &HeaderMap) -> BTreeMap { - headers - .iter() - .filter_map(|(name, value)| { - let name = name.as_str(); - matches!( - name, - "retry-after" - | "request-id" - | "traceparent" - | "x-request-id" - | "x-ratelimit-limit" - | "x-ratelimit-remaining" - | "x-ratelimit-reset" - | "ratelimit-limit" - | "ratelimit-remaining" - | "ratelimit-reset" - ) - .then(|| { - ( - name.to_owned(), - bounded_utf8( - String::from_utf8_lossy(value.as_bytes()).into_owned(), - MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES, - ), - ) - }) - }) - .collect() -} - -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() + 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)] diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 0768388f7..29585b2ec 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -12,6 +12,56 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use thiserror::Error; +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")] diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 7b44b234e..08833342e 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -6,9 +6,9 @@ #[cfg(test)] use std::cell::Cell; use std::cell::RefCell; +use std::collections::HashMap; #[cfg(test)] use std::collections::HashSet; -use std::collections::{BTreeMap, HashMap}; use std::ffi::c_void; use std::future::Future; use std::panic::{AssertUnwindSafe, catch_unwind}; @@ -43,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, UpstreamFailureClass}; +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, @@ -51,8 +53,7 @@ use crate::plugin::{ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ - LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmContinuationOutcomeV2, - LlmContinuationStreamEventV2, LlmHttpFailureV2, LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, + LlmContinuationFailureV2, LlmContinuationInvocationV2, LlmNonHttpFailureKindV2, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NEMO_RELAY_NATIVE_ABI_VERSION_TARGETED_LLM_CONTINUATIONS, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmResultCbV2, @@ -575,7 +576,6 @@ struct NativeHostScopeStackBinding(ThreadScopeStackBinding); thread_local! { static NATIVE_LAST_ERROR: RefCell> = const { RefCell::new(None) }; - static CURRENT_NATIVE_ASYNC_TASK_V2: RefCell>> = const { RefCell::new(None) }; #[cfg(test)] static NATIVE_STRING_LIVE_ALLOCATIONS: RefCell> = RefCell::new(HashSet::new()); #[cfg(test)] @@ -950,12 +950,13 @@ 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_cancel_v2: native_async_llm_stream_cancel_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, @@ -1634,7 +1635,6 @@ struct NativeAsyncStream { runtime: tokio::runtime::Handle, context: MiddlewareContinuationContext, task: Mutex>>, - backpressure_waiter: Mutex>>, #[cfg(test)] before_settlement_lock: Option>, _callback_user_data: Option>, @@ -1662,13 +1662,6 @@ impl NativeAsyncTaskOwnerV2 { } } - fn is_stream(&self, stream: *const NativeAsyncStream) -> bool { - match self { - Self::Stream(owner) => owner.as_ptr() == stream, - Self::Completion(_) => false, - } - } - fn detach(&self, task: &Arc) { match self { Self::Completion(completion) => { @@ -1679,7 +1672,6 @@ impl NativeAsyncTaskOwnerV2 { Self::Stream(stream) => { if let Some(owner) = stream.upgrade() { clear_native_async_task_owner_slot(&owner.task, task); - clear_native_async_task_waiter_slot(&owner.backpressure_waiter, task); } } } @@ -1699,24 +1691,10 @@ fn clear_native_async_task_owner_slot( } } -fn clear_native_async_task_waiter_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, - _library_guard: Option>, } impl Drop for NativeAsyncTaskPluginStateV2 { @@ -1746,25 +1724,6 @@ struct NativeAsyncTaskV2 { _library_guard: Option>, } -struct CurrentNativeAsyncTaskBindingV2(Option>); - -impl CurrentNativeAsyncTaskBindingV2 { - fn bind(task: &Arc) -> Self { - let previous = CURRENT_NATIVE_ASYNC_TASK_V2 - .with(|current| current.replace(Some(Arc::downgrade(task)))); - Self(previous) - } -} - -impl Drop for CurrentNativeAsyncTaskBindingV2 { - fn drop(&mut self) { - let previous = self.0.take(); - CURRENT_NATIVE_ASYNC_TASK_V2.with(|current| { - current.replace(previous); - }); - } -} - impl NativeAsyncTaskV2 { fn wake(task: &Arc) { if !task @@ -1826,7 +1785,6 @@ impl NativeAsyncTaskV2 { let callback_result = self .context .run(async move { - let _binding = CurrentNativeAsyncTaskBindingV2::bind(&task); catch_unwind(AssertUnwindSafe(|| unsafe { cb( user_data as *mut c_void, @@ -1933,28 +1891,228 @@ struct NativeAsyncStreamCallbackGuard { } struct NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex>, - producer_abort: Mutex>, + stream: tokio::sync::Mutex>, runtime: tokio::runtime::Handle, - next_in_flight: AtomicBool, - cancelled: AtomicBool, - _library_guard: Option>, + context: MiddlewareContinuationContext, + target: LlmDispatchTargetContext, + output: Arc, + lifecycle: Mutex, } -impl Drop for NativeLlmProviderStreamV2 { - fn drop(&mut self) { - self.cancelled.store(true, Ordering::Release); - if let Some(abort) = self - .producer_abort - .lock() - .unwrap_or_else(|error| error.into_inner()) - .take() +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(); } } } +impl Drop for NativeLlmProviderStreamV2 { + fn drop(&mut self) { + self.cancel(); + } +} + +struct NativeLlmProviderNextCallbackGuardV2 { + cb: NemoRelayNativeAsyncLlmStreamNextCbV2, + user_data: usize, + provider: Arc, + active: bool, +} + +enum NativeLlmProviderNextOutcomeV2 { + Chunk(Json), + Done, + Failure(LlmContinuationFailureV2), +} + +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); + } + 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); + } + } + + 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); + } + } + } + + fn complete(&mut self, outcome: NativeLlmProviderNextOutcomeV2) { + if !self.active { + return; + } + self.active = false; + match outcome { + NativeLlmProviderNextOutcomeV2::Chunk(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); + } + } + } + } + NativeLlmProviderNextOutcomeV2::Done => { + 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(), + )); + } + } + NativeLlmProviderNextOutcomeV2::Failure(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 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, @@ -2024,12 +2182,32 @@ impl Drop for NativeAsyncResultCallbackGuard { } impl NativeLlmResultCallbackGuardV2 { - fn complete(&mut self, outcome: &LlmContinuationOutcomeV2) { + 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 { - invoke_typed_llm_result_callback(self.cb, self.user_data as *mut c_void, outcome); + (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; } @@ -2038,12 +2216,10 @@ impl NativeLlmResultCallbackGuardV2 { impl Drop for NativeLlmResultCallbackGuardV2 { fn drop(&mut self) { if self.active { - self.complete(&LlmContinuationOutcomeV2::Failure { - error: non_http_llm_failure( - LlmNonHttpFailureKindV2::Cancelled, - "typed native LLM continuation was cancelled".into(), - ), - }); + self.complete(Err(non_http_llm_failure( + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM continuation was cancelled".into(), + ))); } } } @@ -2062,12 +2238,17 @@ impl NativeLlmStreamOpenCallbackGuardV2 { if !self.active { return; } - if let Some(error) = native_string_from_json( + 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); + ); + 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); } } @@ -2091,36 +2272,15 @@ struct NativeLlmStreamForwardCallbackGuardV2 { user_data: usize, stream: Arc, active: bool, - _library_guard: Option>, } impl NativeLlmStreamForwardCallbackGuardV2 { - fn complete(&mut self) { + fn settle(&mut self) { if !self.active { return; } self.active = false; - unsafe { (self.cb)(self.user_data as *mut c_void, ptr::null()) }; - } - - fn fail(&mut self, message: &str) { - if !self.active { - return; - } - self.active = false; - if let Some(message) = native_string_from_str(message) - .or_else(|| native_string_from_str("native LLM stream forwarding failed")) - { - unsafe { - (self.cb)(self.user_data as *mut c_void, message); - native_string_free(message); - } - } else { - // Allocation failure must not strand the SDK trampoline. The - // output is already failed or cancelled, so null only acts as the - // terminal wake-up in this exceptional case. - unsafe { (self.cb)(self.user_data as *mut c_void, ptr::null()) }; - } + unsafe { (self.cb)(self.user_data as *mut c_void) }; } fn cancel_unsettled_output(&self) { @@ -2154,7 +2314,7 @@ impl Drop for NativeLlmStreamForwardCallbackGuardV2 { return; } self.cancel_unsettled_output(); - self.fail("downstream LLM stream forwarding was cancelled"); + self.settle(); } } @@ -2224,7 +2384,7 @@ impl Stream for NativeAsyncStreamReceiver { 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_backpressure_waiter(&self.stream); + wake_native_async_stream_task(&self.stream); } result } @@ -2253,64 +2413,22 @@ impl Drop for NativeAsyncStreamReceiver { .unwrap_or_else(|error| error.into_inner()) .take(); drop(_settlement); - wake_native_async_stream_backpressure_waiter(&self.stream); cancel_native_async_task_slot(&self.stream.task); } } -fn wake_native_async_stream_backpressure_waiter(stream: &NativeAsyncStream) { +fn wake_native_async_stream_task(stream: &NativeAsyncStream) { let task = stream - .backpressure_waiter + .task .lock() .unwrap_or_else(|error| error.into_inner()) - .take() + .as_ref() .and_then(|task| task.upgrade()); if let Some(task) = task { NativeAsyncTaskV2::wake(&task); } } -fn register_current_native_async_stream_backpressure_waiter( - stream: &NativeAsyncStream, - sender: &tokio::sync::mpsc::Sender>, -) { - let task = CURRENT_NATIVE_ASYNC_TASK_V2.with(|current| { - current - .borrow() - .as_ref() - .and_then(Weak::upgrade) - .filter(|task| task.owner.is_stream(stream as *const NativeAsyncStream)) - }); - let Some(task) = task else { - return; - }; - *stream - .backpressure_waiter - .lock() - .unwrap_or_else(|error| error.into_inner()) = Some(Arc::downgrade(&task)); - // The consumer can free capacity between `try_send` and waiter - // registration. Recheck after publishing the waiter to avoid a lost wake. - if sender.capacity() > 0 || sender.is_closed() { - wake_native_async_stream_backpressure_waiter(stream); - } -} - -fn clear_current_native_async_stream_backpressure_waiter(stream: &NativeAsyncStream) { - let current = - CURRENT_NATIVE_ASYNC_TASK_V2.with(|task| task.borrow().as_ref().map(Weak::as_ptr)); - let mut waiter = stream - .backpressure_waiter - .lock() - .unwrap_or_else(|error| error.into_inner()); - if current.is_some_and(|current| { - waiter - .as_ref() - .is_some_and(|waiter| waiter.as_ptr() == current) - }) { - waiter.take(); - } -} - async fn invoke_native_async_callback( cb: NemoRelayNativeAsyncMiddlewareCb, user_data: Arc, @@ -2580,39 +2698,17 @@ unsafe extern "C" fn native_async_completion_spawn_task_v2( set_native_last_error("cannot spawn a task for a settled async completion"); return NemoRelayStatus::InvalidArg; } - let mut slot = completion - .task - .lock() - .unwrap_or_else(|error| error.into_inner()); - if slot.as_ref().and_then(Weak::upgrade).is_some() { - set_native_last_error("an async completion task is already active"); - return NemoRelayStatus::InvalidArg; - } - let task = Arc::new(NativeAsyncTaskV2 { - runtime: completion.runtime.clone(), - context: completion.context.clone(), - owner: NativeAsyncTaskOwnerV2::Completion(unsafe { + spawn_native_async_task_v2( + completion.runtime.clone(), + completion.context.clone(), + NativeAsyncTaskOwnerV2::Completion(unsafe { weak_from_arc_raw(completion as *const NativeAsyncCompletion) }), - wake: tokio::sync::Notify::new(), - state: Mutex::new(NativeAsyncTaskStateV2 { - polling: false, - cancel_requested: false, - complete: false, - plugin: Some(NativeAsyncTaskPluginStateV2 { - cb, - user_data: user_data as usize, - free_fn, - _library_guard: completion._callback_user_data.clone(), - }), - }), - _library_guard: completion._callback_user_data.clone(), - }); - *slot = Some(Arc::downgrade(&task)); - drop(slot); - task.runtime.spawn(Arc::clone(&task).run()); - NativeAsyncTaskV2::wake(&task); - NemoRelayStatus::Ok + &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( @@ -2629,20 +2725,42 @@ unsafe extern "C" fn native_async_stream_spawn_task_v2( set_native_last_error("cannot spawn a task for a settled async stream"); return NemoRelayStatus::InvalidArg; } - let mut slot = stream - .task - .lock() - .unwrap_or_else(|error| error.into_inner()); + 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("an async stream task is already active"); + 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 { - runtime: stream.runtime.clone(), - context: stream.context.clone(), - owner: NativeAsyncTaskOwnerV2::Stream(unsafe { - weak_from_arc_raw(stream as *const NativeAsyncStream) - }), + runtime, + context, + owner, wake: tokio::sync::Notify::new(), state: Mutex::new(NativeAsyncTaskStateV2 { polling: false, @@ -2650,12 +2768,11 @@ unsafe extern "C" fn native_async_stream_spawn_task_v2( complete: false, plugin: Some(NativeAsyncTaskPluginStateV2 { cb, - user_data: user_data as usize, + user_data, free_fn, - _library_guard: stream._callback_user_data.clone(), }), }), - _library_guard: stream._callback_user_data.clone(), + _library_guard: library_guard, }); *slot = Some(Arc::downgrade(&task)); drop(slot); @@ -2697,6 +2814,21 @@ unsafe extern "C" fn native_async_next_release(next: *const NemoRelayNativeAsync 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 { @@ -2732,16 +2864,12 @@ unsafe extern "C" fn native_async_stream_push_json( return NemoRelayStatus::InvalidArg; }; match sender.try_send(Ok(chunk)) { - Ok(()) => { - clear_current_native_async_stream_backpressure_waiter(stream); - NemoRelayStatus::Ok - } + Ok(()) => NemoRelayStatus::Ok, Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - register_current_native_async_stream_backpressure_waiter(stream, &sender); 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, } @@ -2784,7 +2912,6 @@ unsafe extern "C" fn native_async_stream_finish( } drop(downstream_aborts); drop(_settlement); - wake_native_async_stream_backpressure_waiter(stream); cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } else { @@ -2795,6 +2922,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 { @@ -2837,16 +2979,14 @@ unsafe extern "C" fn native_async_stream_reject( drop(downstream_aborts); drop(sender_guard); drop(_settlement); - wake_native_async_stream_backpressure_waiter(stream); cancel_native_async_task_slot(&stream.task); NemoRelayStatus::Ok } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - register_current_native_async_stream_backpressure_waiter(stream, sender); 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, } @@ -3066,22 +3206,17 @@ unsafe extern "C" fn native_async_next_invoke_result( NemoRelayStatus::Ok } -const NATIVE_API_V2_MAX_FAILURE_BODY_BYTES: usize = 16 * 1024; -const NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES: usize = 1024; 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.method, - invocation.target.url, - invocation.target.headers, - ) - .map_err(|error| { - set_native_last_error(error.to_string()); - NemoRelayStatus::InvalidArg - })?; + 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)) } @@ -3090,32 +3225,31 @@ fn non_http_llm_failure( message: String, ) -> LlmContinuationFailureV2 { LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind, - message: bounded_utf8(message, NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES), - }, + kind, + message: bounded_utf8(message, NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES), } } fn typed_llm_failure(error: FlowError) -> LlmContinuationFailureV2 { match error { - FlowError::Upstream(failure) => match failure.status { - Some(status) => LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { + FlowError::Upstream(failure) => { + let failure = sanitize_upstream_failure(failure); + match failure.status { + Some(status) => LlmContinuationFailureV2::Http { status, - body: bounded_utf8(failure.body, NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), - headers: safe_native_api_v2_failure_headers(failure.headers), + 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) + 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) } @@ -3126,63 +3260,6 @@ fn typed_llm_failure(error: FlowError) -> LlmContinuationFailureV2 { } } -fn safe_native_api_v2_failure_headers( - headers: BTreeMap, -) -> BTreeMap { - headers - .into_iter() - .filter_map(|(name, value)| { - let normalized = name.to_ascii_lowercase(); - let safe = 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" - ); - safe.then(|| { - ( - normalized, - bounded_utf8(value, NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES), - ) - }) - }) - .collect() -} - -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_string() -} - -unsafe fn invoke_typed_llm_result_callback( - cb: NemoRelayNativeAsyncLlmResultCbV2, - user_data: *mut c_void, - outcome: &LlmContinuationOutcomeV2, -) { - if let Some(outcome) = native_string_from_json( - &serde_json::to_value(outcome) - .expect("native API v2 LLM outcomes contain only serializable Relay DTOs"), - ) { - unsafe { - cb(user_data, outcome); - native_string_free(outcome); - } - } -} - /// Invokes a unary LLM continuation through native API v2. unsafe extern "C" fn native_async_llm_next_invoke_result_v2( next: *const NemoRelayNativeAsyncNext, @@ -3243,13 +3320,7 @@ unsafe extern "C" fn native_async_llm_next_invoke_result_v2( panic_payload_message(payload.as_ref()) ))) }); - let outcome = match result { - Ok(response) => LlmContinuationOutcomeV2::Success { response }, - Err(error) => LlmContinuationOutcomeV2::Failure { - error: typed_llm_failure(error), - }, - }; - callback_guard.complete(&outcome); + callback_guard.complete(result.map_err(typed_llm_failure)); callbacks .lock() .unwrap_or_else(|error| error.into_inner()) @@ -3458,14 +3529,14 @@ async fn push_forwarded_native_stream_chunk( .map_err(|_| FlowError::Internal("native LLM pass-through output was cancelled".into())) } -fn finish_forwarded_native_stream(stream: &NativeAsyncStream) -> bool { +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 false; + return; } let sender = stream .sender @@ -3473,24 +3544,23 @@ fn finish_forwarded_native_stream(stream: &NativeAsyncStream) -> bool { .unwrap_or_else(|error| error.into_inner()) .take(); if sender.is_none() { - return false; + return; } stream.settled.store(true, Ordering::Release); sender }; abort_native_stream_downstream_tasks(stream); drop(sender); - true } -async fn reject_forwarded_native_stream(stream: &NativeAsyncStream, error: FlowError) -> bool { +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 false; + return; } let sender = stream .sender @@ -3498,14 +3568,13 @@ async fn reject_forwarded_native_stream(stream: &NativeAsyncStream, error: FlowE .unwrap_or_else(|error| error.into_inner()) .take(); let Some(sender) = sender else { - return false; + return; }; stream.settled.store(true, Ordering::Release); sender }; abort_native_stream_downstream_tasks(stream); let _ = sender.send(Err(error)).await; - true } /// Forwards the ordinary downstream LLM continuation into the host-owned @@ -3564,7 +3633,6 @@ unsafe extern "C" fn native_async_llm_next_forward_stream_v2( user_data: user_data as usize, stream: Arc::clone(&output_stream), 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 { @@ -3593,20 +3661,10 @@ unsafe extern "C" fn native_async_llm_next_forward_stream_v2( .await; remove_current_native_stream_task(&callback_guard.stream); match result { - Ok(()) if finish_forwarded_native_stream(&callback_guard.stream) => { - callback_guard.complete(); - } - Ok(()) => { - callback_guard.fail("downstream LLM stream forwarding was cancelled"); - } - Err(error) => { - if reject_forwarded_native_stream(&callback_guard.stream, error).await { - callback_guard.fail("downstream LLM stream failed"); - } else { - callback_guard.fail("downstream LLM stream forwarding was cancelled"); - } - } + Ok(()) => finish_forwarded_native_stream(&callback_guard.stream), + Err(error) => reject_forwarded_native_stream(&callback_guard.stream, error).await, } + callback_guard.settle(); }); let abort = task.abort_handle(); downstream_aborts.insert(task.id(), abort); @@ -3676,13 +3734,13 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( .unwrap_or_else(|error| error.into_inner()); let next_fn = next_fn.clone(); let provider_runtime = next.runtime.clone(); - let provider_library_guard = next._callback_user_data.clone(); - let user_data = user_data as usize; + 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 = NativeLlmStreamOpenCallbackGuardV2 { cb, - user_data, + user_data: user_data as usize, _library_guard: next._callback_user_data.clone(), active: true, }; @@ -3691,99 +3749,46 @@ unsafe extern "C" fn native_async_llm_next_open_stream_v2( if start_rx.await.is_err() { return; } - continuation_context - .invoke_with_llm_dispatch_target(target, move || async move { - let mut callback_guard = callback_guard; - let result = AssertUnwindSafe(async { - match next_fn(request).await { - Ok(mut provider_stream) => { - let (sender, receiver) = - tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); - let provider = Arc::new(NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex::new(receiver), - producer_abort: Mutex::new(None), - runtime: provider_runtime.clone(), - next_in_flight: AtomicBool::new(false), - cancelled: AtomicBool::new(false), - _library_guard: provider_library_guard.clone(), - }); - let output_for_producer = Arc::clone(&output_stream_for_task); - let producer = provider_runtime.spawn(async move { - let mut failed = false; - while let Some(item) = provider_stream.next().await { - let terminal = item.is_err(); - let event = match item { - Ok(chunk) => LlmContinuationStreamEventV2::Chunk { chunk }, - Err(error) => LlmContinuationStreamEventV2::Failure { - error: typed_llm_failure(error), - }, - }; - if sender.send(event).await.is_err() || terminal { - failed = terminal; - break; - } - } - if !failed - && !output_for_producer.cancelled.load(Ordering::Acquire) - && !output_for_producer.settled.load(Ordering::Acquire) - { - let _ = sender.send(LlmContinuationStreamEventV2::Done).await; - } - output_for_producer - .downstream_aborts - .lock() - .unwrap_or_else(|error| error.into_inner()) - .remove(&tokio::task::id()); - }); - let producer_abort = producer.abort_handle(); - let producer_id = producer.id(); - *provider - .producer_abort - .lock() - .unwrap_or_else(|error| error.into_inner()) = - Some(producer_abort.clone()); - { - 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) - { - producer_abort.abort(); - callback_guard.failure(&non_http_llm_failure( - LlmNonHttpFailureKindV2::Cancelled, - "typed native LLM stream output settled during setup" - .into(), - )); - return; - } - output_stream_for_task - .downstream_aborts - .lock() - .unwrap_or_else(|error| error.into_inner()) - .insert(producer_id, producer_abort); - } - callback_guard.success(provider); - } - Err(error) => { - callback_guard.failure(&typed_llm_failure(error)); - } - } - }) - .catch_unwind() - .await; - if let Err(payload) = result { + 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::Internal, - format!( - "typed native LLM stream continuation panicked: {}", - panic_payload_message(payload.as_ref()) - ), + LlmNonHttpFailureKindV2::Cancelled, + "typed native LLM stream output settled during setup".into(), )); + } else { + callback_guard.success(provider); } - }) - .await; + } + 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() @@ -3806,71 +3811,111 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( let Some(stream) = (unsafe { (stream as *const NativeLlmProviderStreamV2).as_ref() }) else { return NemoRelayStatus::NullPointer; }; - if stream.cancelled.load(Ordering::Acquire) { - set_native_last_error("native API v2 provider stream is cancelled"); + 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; } - if stream.next_in_flight.swap(true, Ordering::AcqRel) { - set_native_last_error("native API v2 provider stream already has a pending next operation"); - 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; + } } unsafe { Arc::increment_strong_count(stream as *const NativeLlmProviderStreamV2) }; - let stream = unsafe { Arc::from_raw(stream as *const NativeLlmProviderStreamV2) }; - let runtime = stream.runtime.clone(); - let user_data = user_data as usize; - runtime.spawn(async move { - let event = stream - .receiver - .lock() - .await - .recv() - .await - .unwrap_or_else(|| LlmContinuationStreamEventV2::Failure { - error: non_http_llm_failure( - LlmNonHttpFailureKindV2::Cancelled, - "native API v2 provider stream closed without a terminal event".into(), - ), - }); - // The pull operation is complete before callback delivery. Clearing - // the guard first lets a callback wake plugin code that immediately - // requests the next event without racing this task's epilogue. - stream.next_in_flight.store(false, Ordering::Release); - if let Some(event) = native_string_from_json( - &serde_json::to_value(&event) - .expect("native API v2 stream events contain serializable Relay DTOs"), - ) { - unsafe { - cb(user_data as *mut c_void, event); - native_string_free(event); - } + 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 NativeLlmProviderNextOutcomeV2::Done; + }; + match provider_stream.next().await { + Some(Ok(chunk)) => NativeLlmProviderNextOutcomeV2::Chunk(chunk), + Some(Err(error)) => { + stream.take(); + NativeLlmProviderNextOutcomeV2::Failure(typed_llm_failure(error)) + } + None => { + stream.take(); + NativeLlmProviderNextOutcomeV2::Done + } + } + })) + .catch_unwind() + .await; + let outcome = match result { + Ok(outcome) => outcome, + Err(payload) => { + provider_for_task.stream.lock().await.take(); + NativeLlmProviderNextOutcomeV2::Failure(non_http_llm_failure( + LlmNonHttpFailureKindV2::Internal, + format!( + "typed native LLM provider stream panicked: {}", + panic_payload_message(payload.as_ref()) + ), + )) + } + }; + callback_guard.complete(outcome); }); - NemoRelayStatus::Ok -} - -/// Cancels a native API v2 provider stream. -unsafe extern "C" fn native_async_llm_stream_cancel_v2( - stream: *const NemoRelayNativeLlmStreamV2, -) -> NemoRelayStatus { - let Some(stream) = (unsafe { (stream as *const NativeLlmProviderStreamV2).as_ref() }) else { - return NemoRelayStatus::NullPointer; + let task_id_value = task.id(); + let abort = task.abort_handle(); + *lifecycle = NativeLlmProviderStreamLifecycleV2::Pulling { + task_id: task_id_value, }; - stream.cancelled.store(true, Ordering::Release); - if let Some(abort) = stream - .producer_abort + provider + .output + .downstream_aborts .lock() .unwrap_or_else(|error| error.into_inner()) - .take() - { - abort.abort(); - } + .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() { - unsafe { drop(Arc::from_raw(stream as *const NativeLlmProviderStreamV2)) }; + let stream = unsafe { Arc::from_raw(stream as *const NativeLlmProviderStreamV2) }; + stream.cancel(); } } @@ -4187,7 +4232,6 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( runtime: runtime.clone(), context, task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), #[cfg(test)] before_settlement_lock: None, _callback_user_data: Some(user_data.clone()), @@ -4271,7 +4315,7 @@ async fn settle_native_async_stream_error(stream: Arc, messag abort.abort(); } drop(downstream_aborts); - wake_native_async_stream_backpressure_waiter(&stream); + wake_native_async_stream_task(&stream); cancel_native_async_task_slot(&stream.task); } 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/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 050b2c7fa..6d38e73c7 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -12,11 +12,11 @@ use nemo_relay_plugin::{ 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, + NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, + NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, + NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, + NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, + ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, }; use serde_json::{Map, json}; @@ -433,7 +433,6 @@ fn targeted_fixture_invocation(url: String, request: LlmRequest) -> LlmContinuat LlmContinuationInvocationV2 { request, target: LlmContinuationTargetV2 { - method: "POST".into(), url, headers: std::collections::BTreeMap::from([( "authorization".into(), diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index 29763aba4..86985b278 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -2,6 +2,7 @@ // 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}; @@ -11,6 +12,12 @@ 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 { @@ -131,8 +138,7 @@ fn response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { } fn target(url: String, headers: BTreeMap) -> LlmDispatchTargetContext { - LlmDispatchTargetContext::try_new("POST".into(), url, headers) - .expect("test target should be valid") + LlmDispatchTargetContext::try_new(url, headers).expect("test target should be valid") } fn request() -> LlmRequest { @@ -148,6 +154,15 @@ fn request() -> LlmRequest { } } +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( @@ -180,7 +195,7 @@ async fn buffered_target_runs_after_downstream_middleware_and_ignores_host_callb terminal(request) }); - let result = scope_llm_dispatch_target(target, downstream(request())) + let result = scope_test_target(target, downstream(request())) .await .expect("targeted request should succeed"); @@ -198,7 +213,7 @@ async fn buffered_target_runs_after_downstream_middleware_and_ignores_host_callb #[tokio::test] async fn buffered_http_failure_is_bounded_and_filters_headers() { - let body = vec![b'x'; MAX_UPSTREAM_ERROR_BODY_BYTES + 1024]; + let body = vec![b'x'; MAX_UPSTREAM_FAILURE_BODY_BYTES + 1024]; let provider = FakeProvider::spawn(response( "429 Too Many Requests", &[ @@ -217,7 +232,7 @@ async fn buffered_http_failure_is_bounded_and_filters_headers() { }; assert_eq!(failure.status, Some(429)); assert_eq!(failure.class, UpstreamFailureClass::RetryableStatus); - assert_eq!(failure.body.len(), MAX_UPSTREAM_ERROR_BODY_BYTES); + assert_eq!(failure.body.len(), MAX_UPSTREAM_FAILURE_BODY_BYTES); assert_eq!( failure.headers.get("retry-after").map(String::as_str), Some("2") @@ -287,7 +302,7 @@ async fn streaming_target_decodes_events_empty_streams_and_late_errors() { Box::pin(async { Ok(LlmJsonStream::new(futures_util::stream::empty())) }) })); let stream_target = target(provider.url.clone(), BTreeMap::new()); - let mut stream = scope_llm_dispatch_target(stream_target, terminal(request())) + let mut stream = scope_test_target(stream_target, terminal(request())) .await .expect("stream should open"); assert_eq!( @@ -337,21 +352,14 @@ async fn streaming_target_decodes_events_empty_streams_and_late_errors() { #[test] fn target_validation_rejects_unsafe_transport_inputs() { - for (method, url, headers) in [ - ("TRACE", "https://provider.example/v1", BTreeMap::new()), - ("POST", "ftp://provider.example/v1", BTreeMap::new()), - ( - "POST", - "https://user:secret@provider.example/v1", - BTreeMap::new(), - ), + for (url, headers) in [ + ("ftp://provider.example/v1", BTreeMap::new()), + ("https://user:secret@provider.example/v1", BTreeMap::new()), ( - "POST", "https://provider.example/v1", BTreeMap::from([("host".into(), "attacker.invalid".into())]), ), ( - "POST", "https://provider.example/v1", BTreeMap::from([( "x-nemo-relay-internal-dispatch-url".into(), @@ -359,7 +367,6 @@ fn target_validation_rejects_unsafe_transport_inputs() { )]), ), ( - "POST", "https://provider.example/v1", BTreeMap::from([ ("Authorization".into(), "Bearer first".into()), @@ -367,25 +374,13 @@ fn target_validation_rejects_unsafe_transport_inputs() { ]), ), ] { - assert!(LlmDispatchTargetContext::try_new(method.into(), url.into(), headers,).is_err()); + assert!(LlmDispatchTargetContext::try_new(url.into(), headers).is_err()); } } #[test] -fn target_validation_reports_malformed_method_and_headers() { - let error = LlmDispatchTargetContext::try_new( - "P OST".into(), - "https://provider.example/v1".into(), - BTreeMap::new(), - ) - .expect_err("method token containing a space should be rejected"); - let FlowError::InvalidArgument(message) = error else { - panic!("expected invalid method argument"); - }; - assert_eq!(message, "LLM continuation method was invalid or prohibited"); - +fn target_validation_reports_malformed_headers() { let error = LlmDispatchTargetContext::try_new( - "POST".into(), "https://provider.example/v1".into(), BTreeMap::from([("bad header".into(), "value".into())]), ) @@ -399,7 +394,6 @@ fn target_validation_reports_malformed_method_and_headers() { ); let error = LlmDispatchTargetContext::try_new( - "POST".into(), "https://provider.example/v1".into(), BTreeMap::from([("x-target".into(), "line one\nline two".into())]), ) @@ -415,14 +409,14 @@ fn target_validation_reports_malformed_method_and_headers() { #[test] fn bounded_utf8_truncates_long_multibyte_value_at_character_boundary() { - let expected = "a".repeat(MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES - 1); + let expected = "a".repeat(MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1); let value = format!("{expected}\u{e9}"); - assert_eq!(value.len(), MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES + 1); + assert_eq!(value.len(), MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES + 1); - let bounded = bounded_utf8(value, MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES); + let bounded = bounded_utf8(value, MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES); assert_eq!(bounded, expected); - assert_eq!(bounded.len(), MAX_UPSTREAM_ERROR_HEADER_VALUE_BYTES - 1); + assert_eq!(bounded.len(), MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1); } #[tokio::test] @@ -441,7 +435,7 @@ async fn transport_failures_do_not_fall_back_to_the_host_callback() { BTreeMap::new(), ); - let error = scope_llm_dispatch_target(target, terminal(request())) + let error = scope_test_target(target, terminal(request())) .await .expect_err("connection should fail"); let FlowError::Upstream(failure) = error else { @@ -452,3 +446,109 @@ async fn transport_failures_do_not_fall_back_to_the_host_callback() { 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 6ff8a5a25..3f01cc15d 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -34,6 +34,19 @@ use crate::api::runtime::{ 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 { @@ -113,20 +126,47 @@ unsafe extern "C" fn complete_native_next_result( unsafe extern "C" fn complete_typed_llm_result( user_data: *mut c_void, - outcome_json: *const NemoRelayNativeString, + response_json: *const NemoRelayNativeString, + error_json: *const NemoRelayNativeString, ) { let sender = unsafe { Box::from_raw(user_data as *mut tokio::sync::oneshot::Sender) }; - let outcome = parse_json_arg(outcome_json, "typed LLM result") - .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) - .expect("host emitted a valid typed LLM outcome"); + 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, - _outcome_json: *const NemoRelayNativeString, + _response_json: *const NemoRelayNativeString, + _error_json: *const NemoRelayNativeString, ) { panic!("invalid dispatch unexpectedly invoked its callback"); } @@ -141,9 +181,17 @@ unsafe extern "C" fn record_typed_llm_stream_open( user_data as *mut tokio::sync::oneshot::Sender>, ) }; - let result = if error_json.is_null() { + 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 { + } 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) }) { @@ -153,13 +201,16 @@ unsafe extern "C" fn record_typed_llm_stream_open( format!("invalid typed stream open error: {status:?}"), )), } - }; - let _ = sender.send(result); + } 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 { - method: "POST".into(), url: url.into(), headers: BTreeMap::new(), } @@ -167,36 +218,123 @@ fn test_dispatch_target(url: &str) -> nemo_relay_plugin::LlmContinuationTargetV2 unsafe extern "C" fn record_typed_llm_stream_next( user_data: *mut c_void, - event_json: *const NemoRelayNativeString, + 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: LlmContinuationStreamEventV2 = parse_json_arg(event_json, "typed LLM stream result") - .and_then(|value| serde_json::from_value(value).map_err(|_| NemoRelayStatus::InvalidJson)) - .expect("host emitted a valid typed LLM stream event"); + 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(), + ), + } + } +} + #[derive(Default)] -struct NativeForwardTerminalState { +struct TypedLlmResultCallbackState { callbacks: AtomicUsize, - error: Mutex>, + outcome: Mutex>, notified: tokio::sync::Notify, } -unsafe extern "C" fn record_native_forward_terminal( +unsafe extern "C" fn record_typed_llm_result_state( user_data: *mut c_void, - error: *const NemoRelayNativeString, + 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); - if !error.is_null() { - *state - .error - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = read_native_string(error).ok(); - } state.notified.notify_one(); } @@ -1769,7 +1907,6 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { let target = current_llm_dispatch_target().expect("typed target is bound"); Box::pin(async move { assert!(request.headers.is_empty()); - assert_eq!(target.method(), "POST"); assert_eq!( target.url().as_str(), "https://provider.example/v1/chat/completions" @@ -1824,11 +1961,9 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { outcome, LlmContinuationOutcomeV2::Failure { error: LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { - status: 429, - body: "rate limited".into(), - headers: BTreeMap::from([("retry-after".into(), "1".into())]), - }, + status: 429, + body: "rate limited".into(), + headers: BTreeMap::from([("retry-after".into(), "1".into())]), }, } ); @@ -1903,10 +2038,8 @@ fn native_api_v2_releasing_next_cancels_a_pending_targeted_call() { outcome, LlmContinuationOutcomeV2::Failure { error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - .. - } + kind: LlmNonHttpFailureKindV2::Cancelled, + .. } } )); @@ -1967,23 +2100,11 @@ fn native_api_v2_rejects_non_absolute_dispatch_targets_before_continuation() { } #[test] -fn native_api_v2_rejects_prohibited_target_methods_and_headers() { +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.method = "CONNECT".into(); - assert_eq!( - prepare_llm_continuation_invocation(LlmContinuationInvocationV2 { - request: request.clone(), - target, - }) - .unwrap_err(), - NemoRelayStatus::InvalidArg - ); - assert_last_error_contains("method was invalid or prohibited"); - let mut target = test_dispatch_target("https://provider.example/v1/chat/completions"); target .headers @@ -2031,7 +2152,7 @@ fn native_api_v2_rejects_target_url_credentials() { 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(NATIVE_API_V2_MAX_FAILURE_BODY_BYTES), + body: "é".repeat(crate::error::MAX_UPSTREAM_FAILURE_BODY_BYTES), headers: BTreeMap::from([ ("Authorization".into(), "Bearer secret".into()), ("Set-Cookie".into(), "session=secret".into()), @@ -2040,23 +2161,22 @@ fn native_api_v2_bounds_failure_data_and_removes_sensitive_headers() { "X-Request-ID".into(), format!( "{}é", - "x".repeat(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES - 1) + "x".repeat(crate::error::MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) ), ), ]), class: UpstreamFailureClass::RetryableStatus, })); - let LlmContinuationFailureV2::Http { failure } = error else { + let LlmContinuationFailureV2::Http { body, headers, .. } = error else { panic!("expected an upstream failure"); }; - let LlmHttpFailureV2 { body, headers, .. } = failure; - assert_eq!(body.len(), NATIVE_API_V2_MAX_FAILURE_BODY_BYTES); + 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(NATIVE_API_V2_MAX_FAILURE_HEADER_VALUE_BYTES - 1) + Some(crate::error::MAX_UPSTREAM_FAILURE_HEADER_VALUE_BYTES - 1) ); assert!(!headers.contains_key("authorization")); assert!(!headers.contains_key("set-cookie")); @@ -2129,28 +2249,39 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { .enable_all() .build() .unwrap(); + let provider_polls = Arc::new(AtomicUsize::new(0)); let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(Arc::new(|request| { - assert!(request.headers.is_empty()); - Box::pin(async move { - assert_eq!( - current_llm_dispatch_target() - .expect("typed target is bound") - .url() - .as_str(), - "https://provider.example/v1/messages" - ); - Ok(LlmJsonStream::new(tokio_stream::iter(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, - })), - ]))) + 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") + .url() + .as_str(), + "https://provider.example/v1/messages" + ); + 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, )); @@ -2160,7 +2291,6 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2204,6 +2334,11 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { .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 (sender, receiver) = tokio::sync::oneshot::channel(); @@ -2228,6 +2363,11 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { 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; } @@ -2243,11 +2383,9 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { }, LlmContinuationStreamEventV2::Failure { error: LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { - status: 503, - body: "unavailable".into(), - headers: BTreeMap::new(), - }, + status: 503, + body: "unavailable".into(), + headers: BTreeMap::new(), }, }, ] @@ -2265,6 +2403,98 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { } } +#[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; + 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") + .url() + .as_str() + .to_owned(); + 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 (sender, receiver) = tokio::sync::oneshot::channel::(); + assert_eq!( + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok + ); + 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, target_url.into()), + (captured_scope, target_url.into()), + ] + ); + unsafe { native_async_llm_stream_release_v2(provider_ref) }; +} + #[test] fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { let runtime = tokio::runtime::Builder::new_current_thread() @@ -2294,7 +2524,6 @@ fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2353,13 +2582,6 @@ fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { }); assert!(stream.settled.load(Ordering::Acquire)); assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); - assert!( - terminal - .error - .lock() - .unwrap_or_else(|error| error.into_inner()) - .is_none() - ); unsafe { native_string_free(request); @@ -2397,7 +2619,6 @@ fn native_api_v2_direct_stream_forwarding_preserves_downstream_failure() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2453,14 +2674,6 @@ fn native_api_v2_direct_stream_forwarding_preserves_downstream_failure() { .expect("failure should settle the terminal callback"); }); assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); - assert_eq!( - terminal - .error - .lock() - .unwrap_or_else(|error| error.into_inner()) - .as_deref(), - Some("downstream LLM stream failed") - ); unsafe { native_string_free(request); @@ -2500,7 +2713,6 @@ fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2545,14 +2757,6 @@ fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { }); assert!(stream.cancelled.load(Ordering::Acquire)); assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); - assert!( - terminal - .error - .lock() - .unwrap_or_else(|error| error.into_inner()) - .as_deref() - .is_some_and(|error| error.contains("cancelled")) - ); unsafe { native_string_free(request); @@ -2563,33 +2767,49 @@ fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { #[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 (provider_sender, provider_receiver) = - tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); - let provider = Arc::new(NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex::new(provider_receiver), - producer_abort: Mutex::new(None), - runtime: runtime.handle().clone(), - next_in_flight: AtomicBool::new(false), - cancelled: AtomicBool::new(false), - _library_guard: None, - }); + 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 (first_sender, first_receiver) = tokio::sync::oneshot::channel(); + let callback = Arc::new(TypedLlmStreamNextCallbackState::default()); assert_eq!( unsafe { native_async_llm_stream_next_v2( provider_ref, - record_typed_llm_stream_next, - Box::into_raw(Box::new(first_sender)).cast(), + 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::(); @@ -2607,34 +2827,113 @@ fn native_api_v2_provider_stream_rejects_overlapping_next_and_releases_in_flight 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_cancel_v2(provider_ref) }, - NemoRelayStatus::Ok - ); - assert_eq!( - unsafe { native_async_llm_stream_cancel_v2(provider_ref) }, + unsafe { + native_async_llm_stream_next_v2( + provider_ref, + record_typed_llm_stream_next_state, + Arc::as_ptr(&callback).cast_mut().cast(), + ) + }, NemoRelayStatus::Ok ); - drop(provider_sender); - // Releasing the plugin's reference while `next` is pending is safe because - // the callback task retains its own provider-stream reference. - unsafe { native_async_llm_stream_release_v2(provider_ref) }; + runtime + .block_on(async { tokio::task::spawn_blocking(move || started_rx.recv()).await }) + .unwrap() + .unwrap(); - let event = runtime - .block_on(async { tokio::time::timeout(Duration::from_secs(1), first_receiver).await }) - .expect("cancelled provider next should settle") - .expect("provider next callback should be delivered"); + 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!( - event, - LlmContinuationStreamEventV2::Failure { + callback + .event + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref(), + Some(LlmContinuationStreamEventV2::Failure { error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - .. - } + 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] @@ -2799,11 +3098,9 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { } #[test] -fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { +fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_deadlock() { const STREAM_COUNT: usize = 64; const EVENT_COUNT: usize = 100; - assert_eq!(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1, 64); - assert_eq!(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2, 32); let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) @@ -2828,7 +3125,6 @@ fn native_api_v2_handles_64_concurrent_100_event_provider_streams() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(output_sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2971,7 +3267,6 @@ fn test_native_output_stream( runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -2984,6 +3279,31 @@ fn test_native_output_stream( ) } +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() @@ -3198,12 +3518,8 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { active: true, _library_guard: None, }; - typed_guard.complete(&LlmContinuationOutcomeV2::Success { - response: json!({"ok": true}), - }); - typed_guard.complete(&LlmContinuationOutcomeV2::Success { - response: Json::Null, - }); + typed_guard.complete(Ok(json!({"ok": true}))); + typed_guard.complete(Ok(Json::Null)); assert_eq!( receiver.blocking_recv().unwrap(), LlmContinuationOutcomeV2::Success { @@ -3222,10 +3538,8 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { receiver.blocking_recv().unwrap(), LlmContinuationOutcomeV2::Failure { error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - .. - } + kind: LlmNonHttpFailureKindV2::Cancelled, + .. } } )); @@ -3234,15 +3548,8 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { .enable_all() .build() .unwrap(); - let (provider_sender, provider_receiver) = tokio::sync::mpsc::channel(1); - let provider = Arc::new(NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex::new(provider_receiver), - producer_abort: Mutex::new(None), - runtime: runtime.handle().clone(), - next_in_flight: AtomicBool::new(false), - cancelled: AtomicBool::new(false), - _library_guard: None, - }); + 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 { @@ -3259,10 +3566,8 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { assert!(matches!( receiver.blocking_recv().unwrap(), Err(LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Transport, - .. - } + kind: LlmNonHttpFailureKindV2::Transport, + .. }) )); @@ -3277,13 +3582,10 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { assert!(matches!( receiver.blocking_recv().unwrap(), Err(LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - .. - } + kind: LlmNonHttpFailureKindV2::Cancelled, + .. }) )); - drop(provider_sender); drop(provider); let (stream, _receiver) = test_native_output_stream(1); @@ -3293,45 +3595,178 @@ fn native_api_v2_callback_guards_settle_once_and_cancel_on_drop() { user_data: (&terminal as *const NativeForwardTerminalState) as usize, stream: Arc::clone(&stream), active: true, - _library_guard: None, }; - forward_guard.complete(); - forward_guard.complete(); + 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(); - fail_native_string_allocation_after(0); - let mut forward_guard = NativeLlmStreamForwardCallbackGuardV2 { + drop(NativeLlmStreamForwardCallbackGuardV2 { cb: record_native_forward_terminal, user_data: (&terminal as *const NativeForwardTerminalState) as usize, stream: Arc::clone(&stream), active: true, - _library_guard: None, - }; - forward_guard.fail("allocation failure"); - forward_guard.fail("ignored"); + }); + 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(); + + 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) }; + } + + 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!( - terminal - .error + 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_deref(), - Some("native LLM stream forwarding failed") - ); + .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); + } - 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, - _library_guard: None, - }); - assert!(stream.cancelled.load(Ordering::Acquire)); - assert_eq!(terminal.callbacks.load(Ordering::SeqCst), 1); + 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); } #[test] @@ -3369,23 +3804,20 @@ fn native_api_v2_failure_mapping_covers_http_and_non_http_kinds() { ), ]; for (error, expected_kind) in cases { - let LlmContinuationFailureV2::NonHttp { failure } = typed_llm_failure(error) else { + let LlmContinuationFailureV2::NonHttp { kind, .. } = typed_llm_failure(error) else { panic!("expected non-HTTP failure") }; - assert_eq!(failure.kind, expected_kind); + assert_eq!(kind, expected_kind); } let message = format!("{}é", "x".repeat(NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES)); - let LlmContinuationFailureV2::NonHttp { failure } = + let LlmContinuationFailureV2::NonHttp { message, .. } = non_http_llm_failure(LlmNonHttpFailureKindV2::Internal, message) else { panic!("expected non-HTTP failure") }; - assert_eq!( - failure.message.len(), - NATIVE_API_V2_MAX_FAILURE_MESSAGE_BYTES - ); - assert!(failure.message.is_char_boundary(failure.message.len())); + 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 { @@ -3520,11 +3952,11 @@ fn native_api_v2_unary_entrypoint_reports_validation_panics_and_typed_failures() let outcome = runtime.block_on(receiver).unwrap(); match outcome { LlmContinuationOutcomeV2::Failure { - error: LlmContinuationFailureV2::Http { failure }, - } => assert_eq!(Some(failure.status), http_status), + error: LlmContinuationFailureV2::Http { status, .. }, + } => assert_eq!(Some(status), http_status), LlmContinuationOutcomeV2::Failure { - error: LlmContinuationFailureV2::NonHttp { failure }, - } => assert_eq!(Some(failure.kind), non_http_kind), + error: LlmContinuationFailureV2::NonHttp { kind, .. }, + } => assert_eq!(Some(kind), non_http_kind), other => panic!("unexpected continuation outcome: {other:?}"), } unsafe { native_async_next_release(next_ref) }; @@ -3556,10 +3988,8 @@ fn native_api_v2_unary_entrypoint_reports_validation_panics_and_typed_failures() runtime.block_on(receiver).unwrap(), LlmContinuationOutcomeV2::Failure { error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Internal, - .. - } + kind: LlmNonHttpFailureKindV2::Internal, + .. } } )); @@ -3734,10 +4164,8 @@ fn native_api_v2_stream_open_reports_validation_setup_and_provider_failures() { assert!(matches!( runtime.block_on(open_receiver).unwrap(), Err(LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Internal, - .. - } + kind: LlmNonHttpFailureKindV2::Internal, + .. }) )); drop(panic_receiver); @@ -3842,17 +4270,17 @@ fn native_api_v2_stream_forwarding_validates_handles_and_contains_panics() { drop(receiver); let (settled, _receiver) = test_native_output_stream(1); - assert!(finish_forwarded_native_stream(&settled)); - assert!(!finish_forwarded_native_stream(&settled)); + finish_forwarded_native_stream(&settled); + finish_forwarded_native_stream(&settled); assert!( runtime .block_on(push_forwarded_native_stream_chunk(&settled, Json::Null)) .is_err() ); - assert!(!runtime.block_on(reject_forwarded_native_stream( + runtime.block_on(reject_forwarded_native_stream( &settled, FlowError::Internal("late".into()), - ))); + )); let panicking = Arc::new(NativeAsyncNext::new( NativeAsyncNextInner::LlmStream(Arc::new(|_| { @@ -3912,7 +4340,7 @@ fn native_api_v2_stream_forwarding_validates_handles_and_contains_panics() { } #[test] -fn native_api_v2_provider_stream_reports_closed_and_cancelled_states() { +fn native_api_v2_provider_stream_reports_terminal_and_release_states() { assert_eq!( unsafe { native_async_llm_stream_next_v2( @@ -3923,25 +4351,14 @@ fn native_api_v2_provider_stream_reports_closed_and_cancelled_states() { }, NemoRelayStatus::NullPointer ); - assert_eq!( - unsafe { native_async_llm_stream_cancel_v2(ptr::null()) }, - NemoRelayStatus::NullPointer - ); + unsafe { native_async_llm_stream_release_v2(ptr::null()) }; let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); - let (sender, receiver) = tokio::sync::mpsc::channel(1); - drop(sender); - let stream = Arc::new(NativeLlmProviderStreamV2 { - receiver: tokio::sync::Mutex::new(receiver), - producer_abort: Mutex::new(None), - runtime: runtime.handle().clone(), - next_in_flight: AtomicBool::new(false), - cancelled: AtomicBool::new(false), - _library_guard: None, - }); + 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 (sender, receiver) = tokio::sync::oneshot::channel(); assert_eq!( @@ -3956,19 +4373,8 @@ fn native_api_v2_provider_stream_reports_closed_and_cancelled_states() { ); assert!(matches!( runtime.block_on(receiver).unwrap(), - LlmContinuationStreamEventV2::Failure { - error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - .. - } - } - } + LlmContinuationStreamEventV2::Done )); - assert_eq!( - unsafe { native_async_llm_stream_cancel_v2(stream_ref) }, - NemoRelayStatus::Ok - ); let callback_state = Box::into_raw(Box::new( tokio::sync::oneshot::channel::().0, )); @@ -3982,12 +4388,32 @@ fn native_api_v2_provider_stream_reports_closed_and_cancelled_states() { }, NemoRelayStatus::InvalidArg ); - assert_last_error_contains("cancelled"); + assert_last_error_contains("terminal"); unsafe { drop(Box::from_raw(callback_state)); native_async_llm_stream_release_v2(stream_ref); - native_async_llm_stream_release_v2(ptr::null()); } + + 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] @@ -4463,7 +4889,6 @@ fn native_async_next_preserves_runtime_context_for_unary_and_stream_continuation runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -4696,7 +5121,6 @@ fn native_async_next_panics_settle_unary_and_stream_errors() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -4863,7 +5287,6 @@ fn cancelled_native_async_next_does_not_start_unary_or_stream_continuations() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(true), settled: AtomicBool::new(false), @@ -4980,7 +5403,6 @@ fn native_async_stream_next_supports_repeated_concurrent_calls() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5100,7 +5522,6 @@ fn native_async_stream_settlement_rejects_late_next_and_aborts_in_flight_next() runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5224,7 +5645,6 @@ fn native_async_stream_next_stops_callbacks_after_false() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5301,7 +5721,6 @@ fn native_async_stream_in_flight_cancellation_releases_callback_state() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5386,7 +5805,6 @@ fn native_async_stream_cancellation_before_first_poll_releases_callback_state() runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -5593,7 +6011,6 @@ fn native_async_task_cancellation_reclaims_never_woken_completion_and_stream_tas runtime: runtime.handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -5783,7 +6200,6 @@ fn native_async_task_stream_backpressure_wakes_after_consumer_drain() { runtime: runtime.handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), before_settlement_lock: None, _callback_user_data: None, }); @@ -6330,7 +6746,6 @@ fn native_async_stream_settlement_cannot_succeed_after_cancellation() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -6392,7 +6807,7 @@ 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() @@ -6402,7 +6817,6 @@ fn native_async_stream_push_is_bounded_retryable_and_incremental() { runtime: native_runtime().handle().clone(), context: MiddlewareContinuationContext::capture(), task: Mutex::new(None), - backpressure_waiter: Mutex::new(None), sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), settled: AtomicBool::new(false), @@ -6414,15 +6828,21 @@ fn native_async_stream_push_is_bounded_retryable_and_incremental() { 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 { native_async_stream_push_json(stream_ref, first_chunk) }, + unsafe { (host_v3.async_stream_push_json)(stream_ref, first_chunk) }, NemoRelayStatus::Ok ); assert_eq!( - unsafe { native_async_stream_push_json(stream_ref, second_chunk) }, + 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), @@ -6432,7 +6852,7 @@ 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_push_json)(stream_ref, second_chunk) }, NemoRelayStatus::Ok ); assert_eq!( @@ -6455,6 +6875,50 @@ fn native_async_stream_push_is_bounded_retryable_and_incremental() { 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 { (host_v3.async_stream_reject)(stream_ref, message) }, + NemoRelayStatus::Internal + ); + assert_eq!( + unsafe { (host_v4.v3.async_stream_reject)(stream_ref, message) }, + 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_reject)(stream_ref, message) }, + 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); + unsafe { + native_string_free(chunk); + native_string_free(message); + native_async_stream_release(stream_ref); + } } #[test] diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 1736078dd..6f28439fd 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -127,6 +127,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. @@ -909,10 +911,9 @@ pub type NemoRelayNativeAsyncTaskPollCbV2 = /// 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 cancels -/// or releases the handle exactly once. Relay pumps provider output into a -/// bounded queue so a plugin that consumes slowly applies backpressure without -/// blocking a runtime worker. +/// 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], @@ -946,10 +947,10 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( ); /// Explicit provider target for one native API v2 LLM continuation. +/// +/// Relay sends LLM continuation requests with HTTP `POST`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] pub struct LlmContinuationTargetV2 { - /// HTTP method used for the provider call. - pub method: String, /// Absolute HTTP(S) provider URL including the selected endpoint. pub url: String, /// Explicit outbound provider headers, including target credentials. @@ -968,17 +969,6 @@ pub struct LlmContinuationInvocationV2 { pub target: LlmContinuationTargetV2, } -/// Bounded HTTP failure returned by a provider. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] -pub struct LlmHttpFailureV2 { - /// Provider HTTP status. - pub status: u16, - /// Bounded provider response body. - pub body: String, - /// Safe response headers with credential-bearing fields removed. - pub headers: BTreeMap, -} - /// Stable non-HTTP failure classification exposed through native API v2. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] #[serde(rename_all = "snake_case")] @@ -997,71 +987,38 @@ pub enum LlmNonHttpFailureKindV2 { Internal, } -/// Bounded failure for an operation that produced no provider HTTP response. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] -pub struct LlmNonHttpFailureV2 { - /// Stable failure kind. - pub kind: LlmNonHttpFailureKindV2, - /// Bounded human-readable context. - pub message: String, -} - /// Structured LLM continuation failure exposed through native API v2. #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] +#[serde(tag = "failure_type", rename_all = "snake_case")] pub enum LlmContinuationFailureV2 { /// A provider returned a non-success HTTP response. Http { - /// Bounded HTTP failure details. - failure: LlmHttpFailureV2, + /// 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 { - /// Bounded non-HTTP failure details. - failure: LlmNonHttpFailureV2, - }, -} - -/// Unary LLM continuation outcome delivered through native API v2. -#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmContinuationOutcomeV2 { - /// Provider call completed successfully. - Success { - /// Provider response JSON. - response: Json, - }, - /// Provider call failed before producing a response. - Failure { - /// Structured Relay/provider failure. - error: LlmContinuationFailureV2, - }, -} - -/// Streaming LLM continuation event delivered through native API v2. -#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum LlmContinuationStreamEventV2 { - /// One provider stream event. - Chunk { - /// Provider event JSON. - chunk: Json, - }, - /// Provider stream completed successfully. - Done, - /// Provider stream failed before clean completion. - Failure { - /// Structured Relay/provider failure. - error: LlmContinuationFailureV2, + /// Stable failure kind. + kind: LlmNonHttpFailureKindV2, + /// Bounded human-readable context. + message: String, }, } /// Receives one typed unary LLM continuation outcome. /// -/// `outcome_json` contains one serialized [`LlmContinuationOutcomeV2`] and is borrowed -/// for the callback. -pub type NemoRelayNativeAsyncLlmResultCbV2 = - unsafe extern "C" fn(user_data: *mut c_void, outcome_json: *const NemoRelayNativeString); +/// 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. /// @@ -1075,18 +1032,23 @@ pub type NemoRelayNativeAsyncLlmStreamOpenCbV2 = unsafe extern "C" fn( /// Receives one item from a native API v2 provider stream. /// -/// `event_json` contains one serialized [`LlmContinuationStreamEventV2`] and is 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, event_json: *const NemoRelayNativeString); +/// 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. /// -/// `error` is null after clean completion and contains a borrowed UTF-8 error -/// message after downstream failure or cancellation. Relay settles the output -/// stream before invoking this callback. -pub type NemoRelayNativeAsyncLlmStreamForwardCbV2 = - unsafe extern "C" fn(user_data: *mut c_void, error: *const NemoRelayNativeString); +/// 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. /// @@ -1194,9 +1156,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, @@ -1206,8 +1169,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, @@ -1288,10 +1253,8 @@ pub struct NemoRelayNativeHostApiV4 { cb: NemoRelayNativeAsyncLlmStreamNextCbV2, user_data: *mut c_void, ) -> NemoRelayStatus, - /// Cancels provider production for a native API v2 stream. - pub async_llm_stream_cancel_v2: - unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2) -> NemoRelayStatus, - /// Releases the plugin-owned provider-stream reference. + /// Releases the plugin-owned provider-stream reference, cancelling provider + /// production first when the stream has not reached a terminal item. pub async_llm_stream_release_v2: unsafe extern "C" fn(stream: *const NemoRelayNativeLlmStreamV2), /// Starts a cooperative task associated with an async completion. @@ -2879,12 +2842,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, @@ -3886,31 +3849,6 @@ pub unsafe fn export_plugin( ) } -/// Initializes a native API v2 plugin descriptor for a Rust SDK plugin value. -/// -/// # 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. -pub unsafe fn export_plugin_v2( - host: *const NemoRelayNativeHostApiV1, - out: *mut NemoRelayNativePluginV1, - plugin: P, -) -> NemoRelayStatus { - 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::(), - || plugin, - ) -} - /// Initializes a native plugin descriptor from a constructor callback. /// /// # Safety @@ -3944,7 +3882,9 @@ where /// Initializes a native API v2 plugin descriptor from a constructor callback. /// /// # Safety -/// `host` and `out` must satisfy [`export_plugin_v2`]'s requirements. +/// `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, diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index 97eaf59c9..b3a25d746 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -18,8 +18,7 @@ use serde::Deserialize; use super::{ HostString, Json, LlmContinuationFailureV2, LlmContinuationInvocationV2, - LlmContinuationOutcomeV2, LlmContinuationStreamEventV2, LlmNonHttpFailureKindV2, - LlmNonHttpFailureV2, LlmRequest, NemoRelayNativeAsyncCallbackState, + LlmNonHttpFailureKindV2, LlmRequest, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncTaskV2, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV4, NemoRelayNativeLlmStreamV2, NemoRelayNativeString, NemoRelayStatus, @@ -68,14 +67,8 @@ pub struct LlmStreamContinuationV2 { pub struct LlmProviderStreamV2 { host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeLlmStreamV2, - pending: Option>, + pending: Option, LlmContinuationFailureV2>>>, finished: bool, - terminal: bool, -} - -struct ProviderItem { - value: std::result::Result, LlmContinuationFailureV2>, - terminal: bool, } struct ContinuationInner { @@ -317,7 +310,8 @@ impl Drop for ContinuationInner { } struct StreamContinuationInner { - continuation: Arc, + host: NemoRelayNativeHostApiV4, + next: *const NemoRelayNativeAsyncNext, output: *const NemoRelayNativeAsyncStream, } @@ -328,7 +322,10 @@ unsafe impl Sync for StreamContinuationInner {} impl Drop for StreamContinuationInner { fn drop(&mut self) { - unsafe { (self.continuation.host.v3.async_stream_release)(self.output) }; + unsafe { + (self.host.v3.async_stream_release)(self.output); + (self.host.v3.async_next_release)(self.next); + } } } @@ -424,12 +421,8 @@ impl LlmStreamContinuationV2 { if next.is_null() || output.is_null() { return Err(NemoRelayStatus::NullPointer); } - let continuation = unsafe { LlmContinuationV2::from_raw(host, next) }?; Ok(Self { - inner: Arc::new(StreamContinuationInner { - continuation: continuation.inner, - output, - }), + inner: Arc::new(StreamContinuationInner { host, next, output }), }) } @@ -438,14 +431,14 @@ impl LlmStreamContinuationV2 { &self, invocation: LlmContinuationInvocationV2, ) -> std::result::Result { - let host = self.inner.continuation.host; + let host = self.inner.host; let (sender, receiver) = oneshot::channel(); let state = Box::new(StreamOpenCallback { host, sender }); let state = Box::into_raw(state).cast::(); let status = match HostString::from_json(&host.v3.v1, &invocation) { Some(invocation) => unsafe { (host.async_llm_next_open_stream_v2)( - self.inner.continuation.next, + self.inner.next, invocation.as_ptr(), self.inner.output, stream_open_callback, @@ -463,37 +456,28 @@ impl LlmStreamContinuationV2 { unsafe { drop(Box::from_raw(state.cast::())) }; return Err(status_failure("targeted LLM stream setup", status)); } - let raw = receiver.await.unwrap_or_else(|_| { + let provider = receiver.await.unwrap_or_else(|_| { Err(internal_failure( "targeted LLM stream setup callback closed without an outcome", )) - })? as *const NemoRelayNativeLlmStreamV2; - if raw.is_null() { - return Err(internal_failure( - "targeted LLM stream setup returned a null stream", - )); - } + })?; Ok(LlmProviderStreamV2 { host, - raw, + raw: provider.into_raw(), pending: None, finished: false, - terminal: false, }) } async fn forward_passthrough(&self, request: LlmRequest) -> Result<()> { - let host = self.inner.continuation.host; + let host = self.inner.host; let (sender, receiver) = oneshot::channel(); - let state = Box::new(ForwardStreamCallback { - host: host.v3.v1, - sender, - }); + let state = Box::new(ForwardStreamCallback { sender }); let state = Box::into_raw(state).cast::(); let status = match HostString::from_json(&host.v3.v1, &request) { Some(request) => unsafe { (host.async_llm_next_forward_stream_v2)( - self.inner.continuation.next, + self.inner.next, request.as_ptr(), self.inner.output, forward_stream_callback, @@ -511,9 +495,9 @@ impl LlmStreamContinuationV2 { "streaming pass-through continuation failed: {status:?}" )); } - receiver.await.unwrap_or_else(|_| { - Err("streaming pass-through callback closed before settlement".into()) - }) + receiver + .await + .map_err(|_| "streaming pass-through callback closed before settlement".into()) } } @@ -557,14 +541,12 @@ impl Stream for LlmProviderStreamV2 { Poll::Pending => Poll::Pending, Poll::Ready(result) => { self.pending = None; - let item = result.unwrap_or_else(|_| ProviderItem { - value: Err(internal_failure( + let item = result.unwrap_or_else(|_| { + Err(internal_failure( "provider stream callback closed without an event", - )), - terminal: false, + )) }); - self.terminal = item.terminal; - match item.value { + match item { Ok(Some(chunk)) => Poll::Ready(Some(Ok(chunk))), Ok(None) => { self.finished = true; @@ -582,11 +564,7 @@ impl Stream for LlmProviderStreamV2 { impl Drop for LlmProviderStreamV2 { fn drop(&mut self) { - if !self.terminal { - let _ = unsafe { (self.host.async_llm_stream_cancel_v2)(self.raw) }; - } unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; - self.raw = ptr::null(); } } @@ -770,14 +748,12 @@ where let host = state.host; let completion_handle = CompletionHandle(completion); let callback = Arc::clone(&state.callback); - let callback_continuation = continuation.clone(); let future = async move { let result = std::panic::AssertUnwindSafe(async move { - callback(invocation.name, invocation.request, callback_continuation).await + callback(invocation.name, invocation.request, continuation).await }) .catch_unwind() .await; - drop(continuation); match result { Ok(Ok(value)) => resolve_completion(&host, completion_handle.as_ptr(), &value), Ok(Err(error)) => reject_completion(&host, completion_handle.as_ptr(), &error), @@ -941,7 +917,7 @@ async fn pump_output_stream( } async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { - let host = &continuation.inner.continuation.host; + let host = &continuation.inner.host; let chunk = TaskHostString::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| { @@ -952,10 +928,7 @@ async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) unsafe { (host.v3.async_stream_push_json)(continuation.inner.output, chunk.as_ptr()) }; match status { NemoRelayStatus::Ok => Poll::Ready(Ok(())), - // For this operation the V3 ABI contract reserves `Internal` for - // a full bounded queue. Serialization and lifecycle faults use - // distinct statuses, so retrying cannot mask another ABI error. - NemoRelayStatus::Internal => Poll::Pending, + NemoRelayStatus::WouldBlock => Poll::Pending, status => Poll::Ready(Err(format!("native API v2 output push failed: {status:?}"))), } }) @@ -963,7 +936,7 @@ async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) } fn finish_output(continuation: &LlmStreamContinuationV2) -> Result<()> { - let host = &continuation.inner.continuation.host; + 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()); } @@ -994,8 +967,7 @@ async fn reject_output(host: &NemoRelayNativeHostApiV4, output: OutputHandle, er let status = unsafe { (host.v3.async_stream_reject)(output.as_ptr(), message.as_ptr()) }; match status { NemoRelayStatus::Ok => Poll::Ready(()), - // `Internal` has the same queue-full-only contract for rejection. - NemoRelayStatus::Internal => Poll::Pending, + NemoRelayStatus::WouldBlock => Poll::Pending, status => { set_last_error( &host.v3.v1, @@ -1090,22 +1062,31 @@ struct TargetedResultCallback { unsafe extern "C" fn targeted_result_callback( user_data: *mut c_void, - outcome_json: *const NemoRelayNativeString, + 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(|| { - let outcome: LlmContinuationOutcomeV2 = read_json_value( - &state.host, - outcome_json, - "targeted LLM continuation outcome", - ) - .map_err(|status| status_failure("targeted LLM continuation outcome", status))?; - match outcome { - LlmContinuationOutcomeV2::Success { response } => Ok(response), - LlmContinuationOutcomeV2::Failure { error } => Err(error), + 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"))); @@ -1143,39 +1124,37 @@ unsafe extern "C" fn passthrough_result_callback( struct StreamOpenCallback { host: NemoRelayNativeHostApiV4, - sender: oneshot::Sender>, + sender: oneshot::Sender>, } struct OwnedProviderStream { host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeLlmStreamV2, - armed: bool, } impl OwnedProviderStream { fn new(host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeLlmStreamV2) -> Self { - Self { - host, - raw, - armed: !raw.is_null(), - } + debug_assert!(!raw.is_null()); + Self { host, raw } } - fn disarm(&mut self) { - self.armed = false; + fn into_raw(mut self) -> *const NemoRelayNativeLlmStreamV2 { + std::mem::replace(&mut self.raw, ptr::null()) } } impl Drop for OwnedProviderStream { fn drop(&mut self) { - if !self.armed { - return; + if !self.raw.is_null() { + unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; } - let _ = unsafe { (self.host.async_llm_stream_cancel_v2)(self.raw) }; - unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; } } +// The host owns the opaque stream and documents release as thread-safe for a +// plugin-owned reference. +unsafe impl Send for OwnedProviderStream {} + unsafe extern "C" fn stream_open_callback( user_data: *mut c_void, stream: *const NemoRelayNativeLlmStreamV2, @@ -1185,10 +1164,11 @@ unsafe extern "C" fn stream_open_callback( return; } let state = unsafe { Box::from_raw(user_data.cast::()) }; - let mut owned_stream = OwnedProviderStream::new(state.host, stream); + let mut owned_stream = + (!stream.is_null()).then(|| OwnedProviderStream::new(state.host, stream)); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { match (stream.is_null(), error_json.is_null()) { - (false, true) => Ok(stream as usize), + (false, true) => Ok(()), (true, false) => read_json_value( &state.host.v3.v1, error_json, @@ -1208,87 +1188,65 @@ unsafe extern "C" fn stream_open_callback( "targeted LLM stream setup callback panicked", )) }); - let transfers_stream = result.is_ok(); - if state.sender.send(result).is_ok() && transfers_stream { - owned_stream.disarm(); - } + let result = result.and_then(|()| { + owned_stream + .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, + sender: oneshot::Sender, LlmContinuationFailureV2>>, } unsafe extern "C" fn provider_next_callback( user_data: *mut c_void, - event_json: *const NemoRelayNativeString, + 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(|| { - let event: LlmContinuationStreamEventV2 = - read_json_value(&state.host, event_json, "targeted provider stream event") - .map_err(|status| status_failure("targeted provider stream event", status))?; - Ok::<_, LlmContinuationFailureV2>(match event { - LlmContinuationStreamEventV2::Chunk { chunk } => ProviderItem { - value: Ok(Some(chunk)), - terminal: false, - }, - LlmContinuationStreamEventV2::Done => ProviderItem { - value: Ok(None), - terminal: true, - }, - LlmContinuationStreamEventV2::Failure { error } => ProviderItem { - value: Err(error), - terminal: true, - }, - }) + 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", )) - }) - .unwrap_or_else(|error| ProviderItem { - value: Err(error), - terminal: false, }); let _ = state.sender.send(result); } struct ForwardStreamCallback { - host: NemoRelayNativeHostApiV1, - sender: oneshot::Sender>, + sender: oneshot::Sender<()>, } -unsafe extern "C" fn forward_stream_callback( - user_data: *mut c_void, - error: *const NemoRelayNativeString, -) { +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 result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - if error.is_null() { - Ok(()) - } else { - // Relay has already settled the output before this wake-up. Keep - // the terminal context available for diagnostics, but report - // completion to the trampoline so it cannot reject a second time. - let message = - read_required_host_string(&state.host, error, "streaming pass-through error") - .unwrap_or_else(|status| { - format!("invalid streaming pass-through error: {status:?}") - }); - set_last_error(&state.host, &message); - Ok(()) - } - })) - .unwrap_or_else(|_| Err("streaming pass-through terminal callback panicked".into())); - let _ = state.sender.send(result); + let _ = state.sender.send(()); } fn status_failure(label: &str, status: NemoRelayStatus) -> LlmContinuationFailureV2 { @@ -1297,10 +1255,8 @@ fn status_failure(label: &str, status: NemoRelayStatus) -> LlmContinuationFailur fn internal_failure(message: impl Into) -> LlmContinuationFailureV2 { LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Internal, - message: bounded_error(&message.into()), - }, + kind: LlmNonHttpFailureKindV2::Internal, + message: bounded_error(&message.into()), } } diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 9880f9251..9c5bd2a2f 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -17,10 +17,9 @@ use futures::{StreamExt, stream}; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmContinuationFailureV2, - LlmContinuationInvocationV2, LlmContinuationOutcomeV2, LlmContinuationStreamEventV2, - LlmContinuationTargetV2, LlmContinuationV2, LlmHttpFailureV2, LlmJsonStream, LlmNext, - LlmNonHttpFailureKindV2, LlmNonHttpFailureV2, LlmRequest, LlmRequestInterceptOutcome, - LlmStream, LlmStreamContinuationV2, LlmStreamExecutionOutcomeV2, LlmStreamNext, + 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, @@ -381,10 +380,20 @@ static ASYNC_STREAM_V2_REGISTRATION: Mutex> = Mu 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()); -static SAFE_V2_PROVIDER_EVENTS: Mutex> = - Mutex::new(VecDeque::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> = @@ -392,19 +401,20 @@ static SAFE_V2_HELD_TARGETED_CALLBACK: 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_FORWARD_ERROR: Mutex> = Mutex::new(None); 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()); @@ -431,7 +441,7 @@ struct SafeV2Task { } #[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!( @@ -447,6 +457,7 @@ 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")] { @@ -469,10 +480,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 528); + assert_eq!(size_of::(), 520); assert_eq!( host_api_v4_offsets(), - [0, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512, 520] + [0, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512] ); assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 56); @@ -503,10 +514,10 @@ fn native_abi_v3_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 260); + assert_eq!(size_of::(), 256); assert_eq!( host_api_v4_offsets(), - [0, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256] + [0, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252] ); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); @@ -517,13 +528,12 @@ fn native_abi_v3_struct_sizes_are_self_describing() { } } -fn host_api_v4_offsets() -> [usize; 12] { +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_cancel_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), @@ -1492,25 +1502,28 @@ fn reset_state() { 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_CANCELS.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_FORWARD_ERROR.lock().unwrap() = None; *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(); @@ -6246,7 +6259,7 @@ unsafe extern "C" fn safe_v2_stream_push( .unwrap_or(NemoRelayStatus::Ok); if status == NemoRelayStatus::Ok { SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); - } else if status == NemoRelayStatus::Internal { + } else if status == NemoRelayStatus::WouldBlock { SAFE_V2_CURRENT_TASK.with(|current| { let task = current.get(); if task != 0 { @@ -6278,7 +6291,7 @@ unsafe extern "C" fn safe_v2_stream_reject( .unwrap_or(NemoRelayStatus::Ok); if status == NemoRelayStatus::Ok { SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); - } else if status == NemoRelayStatus::Internal { + } else if status == NemoRelayStatus::WouldBlock { SAFE_V2_CURRENT_TASK.with(|current| { let task = current.get(); if task != 0 { @@ -6399,13 +6412,19 @@ unsafe extern "C" fn safe_v2_targeted_result( *SAFE_V2_HELD_TARGETED_CALLBACK.lock().unwrap() = Some((cb, user_data as usize)); return NemoRelayStatus::Ok; } - let outcome = serde_json::to_value(LlmContinuationOutcomeV2::Success { - response: json!({ "targeted": true }), - }) - .unwrap(); - let outcome = json_host_string(&host, outcome); - unsafe { cb(user_data, outcome) }; - unsafe { (host.string_free)(outcome) }; + 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 } @@ -6428,6 +6447,10 @@ unsafe extern "C" fn safe_v2_stream_open( { return NemoRelayStatus::InvalidJson; } + 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) { @@ -6460,26 +6483,33 @@ unsafe extern "C" fn safe_v2_provider_next( } let host = test_host(); if let Some(event) = SAFE_V2_PROVIDER_EVENT_JSON.lock().unwrap().take() { - let event = host_string(&host, &event); - unsafe { cb(user_data, event) }; - unsafe { (host.string_free)(event) }; + 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(LlmContinuationStreamEventV2::Done); - let event = json_host_string(&host, serde_json::to_value(event).unwrap()); - unsafe { cb(user_data, event) }; - unsafe { (host.string_free)(event) }; - NemoRelayStatus::Ok -} - -unsafe extern "C" fn safe_v2_provider_cancel( - _stream: *const NemoRelayNativeLlmStreamV2, -) -> NemoRelayStatus { - SAFE_V2_PROVIDER_CANCELS.fetch_add(1, Ordering::SeqCst); + .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 } @@ -6625,13 +6655,7 @@ unsafe extern "C" fn safe_v2_forward_stream( }; SAFE_V2_FORWARDED_REQUESTS.lock().unwrap().push(request); SAFE_V2_OUTPUT_FINISHES.fetch_add(1, Ordering::SeqCst); - if let Some(error) = SAFE_V2_FORWARD_ERROR.lock().unwrap().clone() { - let error = host_string(&host, &error); - unsafe { cb(user_data, error) }; - unsafe { (host.string_free)(error) }; - } else { - unsafe { cb(user_data, ptr::null()) }; - } + unsafe { cb(user_data) }; NemoRelayStatus::Ok } @@ -6661,7 +6685,6 @@ fn test_host_v4() -> NemoRelayNativeHostApiV4 { 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_cancel_v2: safe_v2_provider_cancel, 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, @@ -6674,7 +6697,6 @@ fn test_host_v4() -> NemoRelayNativeHostApiV4 { fn safe_v2_target() -> LlmContinuationTargetV2 { LlmContinuationTargetV2 { - method: "POST".into(), url: "https://provider.example/v1/chat/completions".into(), headers: Default::default(), } @@ -6945,6 +6967,64 @@ fn safe_v2_buffered_registration_wraps_targeted_and_passthrough_calls() { 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(); @@ -7154,21 +7234,17 @@ fn safe_v2_cancelled_targeted_call_releases_next_before_the_host_callback() { .unwrap() .take() .unwrap(); - let outcome = json_host_string( + let error = json_host_string( &host.v3.v1, - serde_json::to_value(LlmContinuationOutcomeV2::Failure { - error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Cancelled, - message: "cancelled by host".into(), - }, - }, + 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, outcome) }; + unsafe { targeted_callback(targeted_user_data as *mut c_void, ptr::null(), error) }; unsafe { - (host.v3.v1.string_free)(outcome); + (host.v3.v1.string_free)(error); registration.free(); } assert_eq!(live_host_strings(), 0); @@ -7178,10 +7254,8 @@ fn safe_v2_cancelled_targeted_call_releases_next_before_the_host_callback() { fn safe_v2_stream_registration_pumps_provider_stream_and_releases_handles() { let _guard = begin_test(); SAFE_V2_PROVIDER_EVENTS.lock().unwrap().extend([ - LlmContinuationStreamEventV2::Chunk { - chunk: json!({ "delta": "hello" }), - }, - LlmContinuationStreamEventV2::Done, + SafeV2ProviderEvent::Chunk(json!({ "delta": "hello" })), + SafeV2ProviderEvent::Done, ]); let host = test_host_v4(); let mut ctx = test_context(&host.v3.v1); @@ -7237,7 +7311,6 @@ fn safe_v2_stream_registration_pumps_provider_stream_and_releases_handles() { vec![Ok(json!({ "delta": "hello" }))] ); assert_eq!(SAFE_V2_OUTPUT_FINISHES.load(Ordering::SeqCst), 1); - assert_eq!(SAFE_V2_PROVIDER_CANCELS.load(Ordering::SeqCst), 0); 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); @@ -7288,14 +7361,12 @@ fn safe_v2_stream_passthrough_uses_host_owned_forwarding() { } #[test] -fn safe_v2_provider_stream_drop_cancels_unfinished_production() { +fn safe_v2_provider_stream_drop_releases_unfinished_production() { let _guard = begin_test(); SAFE_V2_PROVIDER_EVENTS .lock() .unwrap() - .push_back(LlmContinuationStreamEventV2::Chunk { - chunk: json!({ "unused": true }), - }); + .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 { @@ -7327,12 +7398,69 @@ fn safe_v2_provider_stream_drop_cancels_unfinished_production() { }; unsafe { (host.v3.v1.string_free)(invocation) }; drive_safe_v2_tasks(); - assert_eq!(SAFE_V2_PROVIDER_CANCELS.load(Ordering::SeqCst), 1); 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(); @@ -7375,11 +7503,9 @@ fn safe_v2_stream_callback_stops_when_the_caller_cancels() { fn safe_v2_stream_open_preserves_structured_failure() { let _guard = begin_test(); let expected = LlmContinuationFailureV2::Http { - failure: LlmHttpFailureV2 { - status: 429, - body: "bounded".into(), - headers: Default::default(), - }, + status: 429, + body: "bounded".into(), + headers: Default::default(), }; *SAFE_V2_OPEN_FAILURE.lock().unwrap() = Some(expected.clone()); let host = test_host_v4(); @@ -7433,13 +7559,11 @@ fn safe_v2_stream_open_preserves_structured_failure() { } #[test] -fn safe_v2_malformed_stream_open_cancels_and_releases_the_owned_stream() { +fn safe_v2_malformed_stream_open_releases_the_owned_stream() { let _guard = begin_test(); *SAFE_V2_OPEN_FAILURE.lock().unwrap() = Some(LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Internal, - message: "must not accompany a stream".into(), - }, + 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(); @@ -7482,7 +7606,6 @@ fn safe_v2_malformed_stream_open_cancels_and_releases_the_owned_stream() { SAFE_V2_OUTPUT.lock().unwrap().as_slice(), [Err(error)] if error.contains("invalid outcome") )); - assert_eq!(SAFE_V2_PROVIDER_CANCELS.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); @@ -7973,14 +8096,12 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { SAFE_V2_PROVIDER_EVENTS .lock() .unwrap() - .push_back(LlmContinuationStreamEventV2::Failure { - error: LlmContinuationFailureV2::NonHttp { - failure: LlmNonHttpFailureV2 { - kind: LlmNonHttpFailureKindV2::Transport, - message: "late provider failure".into(), - }, + .push_back(SafeV2ProviderEvent::Failure( + LlmContinuationFailureV2::NonHttp { + kind: LlmNonHttpFailureKindV2::Transport, + message: "late provider failure".into(), }, - }); + )); run_safe_v2_streaming(&host, "late-failure", |_, request, next| { safe_v2_targeted_provider_stream(request, next) }); @@ -7996,7 +8117,17 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { }); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), - [Err(error)] if error.contains("stream event") + [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_streaming(&host, "invalid-event", |_, request, next| { + safe_v2_targeted_provider_stream(request, next) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("invalid event") )); assert_eq!(live_host_strings(), 0); } @@ -8009,7 +8140,7 @@ fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors SAFE_V2_OUTPUT_PUSH_STATUSES .lock() .unwrap() - .extend([NemoRelayStatus::Internal, NemoRelayStatus::Ok]); + .extend([NemoRelayStatus::WouldBlock, NemoRelayStatus::Ok]); run_safe_v2_streaming(&host, "push-backpressure", |_, _, _| async { Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( async { Ok(json!({ "chunk": true })) }, @@ -8021,6 +8152,21 @@ fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors ); SAFE_V2_OUTPUT.lock().unwrap().clear(); + SAFE_V2_OUTPUT_PUSH_STATUSES + .lock() + .unwrap() + .push_back(NemoRelayStatus::Internal); + run_safe_v2_streaming(&host, "push-internal", |_, _, _| async { + Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( + async { Ok(json!({ "chunk": true })) }, + )))) + }); + assert!(matches!( + SAFE_V2_OUTPUT.lock().unwrap().as_slice(), + [Err(error)] if error.contains("output push failed: Internal") + )); + SAFE_V2_OUTPUT.lock().unwrap().clear(); + SAFE_V2_OUTPUT_PUSH_STATUSES .lock() .unwrap() @@ -8052,7 +8198,7 @@ fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors SAFE_V2_OUTPUT_REJECT_STATUSES .lock() .unwrap() - .extend([NemoRelayStatus::Internal, NemoRelayStatus::Ok]); + .extend([NemoRelayStatus::WouldBlock, NemoRelayStatus::Ok]); run_safe_v2_streaming(&host, "reject-backpressure", |_, _, _| async { Err("stream rejected".into()) }); @@ -8118,17 +8264,6 @@ fn safe_v2_stream_cancellation_and_pass_through_failures_settle_once() { *SAFE_V2_FORWARD_STATUS.lock().unwrap() = NemoRelayStatus::Ok; SAFE_V2_OUTPUT.lock().unwrap().clear(); - *SAFE_V2_FORWARD_ERROR.lock().unwrap() = Some("downstream stream failed".into()); - *LAST_ERROR.lock().unwrap() = None; - run_safe_v2_streaming(&host, "forward-error", |_, request, _| async move { - Ok(LlmStreamExecutionOutcomeV2::Passthrough(request)) - }); - assert_eq!( - LAST_ERROR.lock().unwrap().as_deref(), - Some("downstream stream failed") - ); - assert!(SAFE_V2_OUTPUT.lock().unwrap().is_empty()); - run_safe_v2_streaming(&host, "stream-panic", |_, _, _| async move { panic!("stream callback panic") }); From fc027436073e273cbd8fb80531b610f0001cc965 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 15:02:22 -0600 Subject: [PATCH 18/32] refactor(plugin): remove native v2 continuation duplication Signed-off-by: Bryan Bednarski --- .../src/api/runtime/llm_dispatch_context.rs | 6 +- crates/core/src/plugin/dynamic/native.rs | 24 +- .../tests/unit/llm_dispatch_context_tests.rs | 20 + crates/core/tests/unit/native_plugin_tests.rs | 433 ++++-------------- crates/plugin/src/lib.rs | 55 ++- crates/plugin/src/native_v2.rs | 86 +--- crates/plugin/tests/typed_callbacks.rs | 304 ++++++------ 7 files changed, 355 insertions(+), 573 deletions(-) diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index dc68e3a70..219c11ba6 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -209,8 +209,8 @@ pub(crate) fn targeted_llm_stream_execution( async fn dispatch_buffered(target: &LlmDispatchTargetContext, request: LlmRequest) -> Result { let response = send(target, request, Some(HTTP_REQUEST_TIMEOUT)).await?; let status = response.status(); - let headers = safe_failure_headers(response.headers()); 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)); } @@ -218,7 +218,9 @@ async fn dispatch_buffered(target: &LlmDispatchTargetContext, request: LlmReques .bytes() .await .map_err(|error| transport_error(target, error))?; - serde_json::from_slice(&bytes).map_err(|_| http_error(status, headers, &bytes)) + serde_json::from_slice(&bytes).map_err(|_| { + FlowError::Internal("targeted LLM provider returned malformed response JSON".into()) + }) } async fn dispatch_stream( diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 08833342e..fe17b1c34 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -1952,12 +1952,6 @@ struct NativeLlmProviderNextCallbackGuardV2 { active: bool, } -enum NativeLlmProviderNextOutcomeV2 { - Chunk(Json), - Done, - Failure(LlmContinuationFailureV2), -} - impl NativeLlmProviderNextCallbackGuardV2 { fn complete_pull(&self, terminal: bool) -> bool { let output_live = !self.provider.output.cancelled.load(Ordering::Acquire) @@ -2040,13 +2034,13 @@ impl NativeLlmProviderNextCallbackGuardV2 { } } - fn complete(&mut self, outcome: NativeLlmProviderNextOutcomeV2) { + fn complete(&mut self, outcome: std::result::Result, LlmContinuationFailureV2>) { if !self.active { return; } self.active = false; match outcome { - NativeLlmProviderNextOutcomeV2::Chunk(chunk) => { + Ok(Some(chunk)) => { let chunk = native_string_from_json(&chunk); let allowed = self.complete_pull(chunk.is_none()); if !allowed { @@ -2068,7 +2062,7 @@ impl NativeLlmProviderNextCallbackGuardV2 { } } } - NativeLlmProviderNextOutcomeV2::Done => { + Ok(None) => { if self.complete_pull(true) { unsafe { (self.cb)( @@ -2085,7 +2079,7 @@ impl NativeLlmProviderNextCallbackGuardV2 { )); } } - NativeLlmProviderNextOutcomeV2::Failure(error) => { + Err(error) => { if self.complete_pull(true) { self.invoke_failure(&error); } else { @@ -3863,17 +3857,17 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( 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 NativeLlmProviderNextOutcomeV2::Done; + return Ok(None); }; match provider_stream.next().await { - Some(Ok(chunk)) => NativeLlmProviderNextOutcomeV2::Chunk(chunk), + Some(Ok(chunk)) => Ok(Some(chunk)), Some(Err(error)) => { stream.take(); - NativeLlmProviderNextOutcomeV2::Failure(typed_llm_failure(error)) + Err(typed_llm_failure(error)) } None => { stream.take(); - NativeLlmProviderNextOutcomeV2::Done + Ok(None) } } })) @@ -3883,7 +3877,7 @@ unsafe extern "C" fn native_async_llm_stream_next_v2( Ok(outcome) => outcome, Err(payload) => { provider_for_task.stream.lock().await.take(); - NativeLlmProviderNextOutcomeV2::Failure(non_http_llm_failure( + Err(non_http_llm_failure( LlmNonHttpFailureKindV2::Internal, format!( "typed native LLM provider stream panicked: {}", diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index 86985b278..b3e3735d8 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -211,6 +211,26 @@ async fn buffered_target_runs_after_downstream_middleware_and_ignores_host_callb 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_http_failure_is_bounded_and_filters_headers() { let body = vec![b'x'; MAX_UPSTREAM_FAILURE_BODY_BYTES + 1024]; diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 3f01cc15d..2b92c4489 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -259,6 +259,63 @@ fn typed_llm_stream_event( } } +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, @@ -1944,18 +2001,7 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { .unwrap(), ) .unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - next_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = start_typed_llm_call(next_ref, dispatch); let outcome = runtime.block_on(receiver).unwrap(); assert_eq!( outcome, @@ -2015,18 +2061,7 @@ fn native_api_v2_releasing_next_cancels_a_pending_targeted_call() { .unwrap(), ) .unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - next_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = start_typed_llm_call(next_ref, dispatch); runtime .block_on(async { tokio::task::spawn_blocking(move || started_rx.recv()).await }) .unwrap() @@ -2286,19 +2321,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { 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 { - 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, - }); + 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( @@ -2313,20 +2336,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { ) .unwrap(); - let (open_sender, open_receiver) = - tokio::sync::oneshot::channel::>(); - assert_eq!( - unsafe { - native_async_llm_next_open_stream_v2( - next_ref, - dispatch, - output_stream_ref, - record_typed_llm_stream_open, - Box::into_raw(Box::new(open_sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + 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 @@ -2341,17 +2351,7 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { ); let mut events = Vec::new(); loop { - let (sender, receiver) = tokio::sync::oneshot::channel(); - assert_eq!( - unsafe { - native_async_llm_stream_next_v2( - provider_stream, - record_typed_llm_stream_next, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = pull_typed_llm_stream(provider_stream); let event = runtime.block_on(async { tokio::time::timeout(Duration::from_secs(1), receiver) .await @@ -2457,17 +2457,7 @@ fn native_api_v2_provider_pulls_restore_scope_and_target_after_pending_wakes() { let provider_ref = Arc::into_raw(provider) as *const NemoRelayNativeLlmStreamV2; for index in 0..2 { - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_stream_next_v2( - provider_ref, - record_typed_llm_stream_next, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = pull_typed_llm_stream(provider_ref); assert_eq!( runtime.block_on(async { tokio::time::timeout(Duration::from_secs(1), receiver) @@ -2519,19 +2509,7 @@ fn native_api_v2_direct_stream_forwarding_is_bounded_and_settles_once() { None, )); let next_ref = Arc::into_raw(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(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 mut output = NativeAsyncStreamReceiver { receiver, @@ -2614,19 +2592,7 @@ fn native_api_v2_direct_stream_forwarding_preserves_downstream_failure() { None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::mpsc::channel(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), - downstream_aborts: Mutex::new(HashMap::new()), - settlement: Mutex::new(()), - before_settlement_lock: None, - _callback_user_data: None, - }); + 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, @@ -2708,19 +2674,7 @@ fn native_api_v2_direct_stream_forwarding_cancels_with_the_consumer() { None, )); let next_ref = Arc::into_raw(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(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 output = NativeAsyncStreamReceiver { receiver, @@ -2974,19 +2928,7 @@ fn native_api_v2_handles_256_concurrent_buffered_dispatches() { let mut receivers = Vec::with_capacity(DISPATCH_COUNT); for _ in 0..DISPATCH_COUNT { - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - next_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); - receivers.push(receiver); + receivers.push(start_typed_llm_call(next_ref, dispatch)); } let outcomes = runtime.block_on(async { @@ -3061,18 +3003,7 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { .unwrap(), ) .unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - next_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = start_typed_llm_call(next_ref, dispatch); unsafe { native_string_free(dispatch) }; receivers.push((index, receiver)); } @@ -3119,20 +3050,8 @@ fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_d None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (output_sender, output_receiver) = - tokio::sync::mpsc::channel(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); - let output_stream = Arc::new(NativeAsyncStream { - runtime: native_runtime().handle().clone(), - context: MiddlewareContinuationContext::capture(), - task: Mutex::new(None), - sender: Mutex::new(Some(output_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, output_receiver) = + test_native_output_stream(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); let output_stream_ref = Arc::into_raw(Arc::clone(&output_stream)) as *const NemoRelayNativeAsyncStream; let dispatch = native_string_from_json( @@ -3149,21 +3068,11 @@ fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_d let mut open_receivers = Vec::with_capacity(STREAM_COUNT); for _ in 0..STREAM_COUNT { - let (sender, receiver) = - tokio::sync::oneshot::channel::>(); - assert_eq!( - unsafe { - native_async_llm_next_open_stream_v2( - next_ref, - dispatch, - output_stream_ref, - record_typed_llm_stream_open, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); - open_receivers.push(receiver); + open_receivers.push(start_typed_llm_stream( + next_ref, + dispatch, + output_stream_ref, + )); } let event_counts = runtime.block_on(async { @@ -3184,17 +3093,7 @@ fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_d let drains = streams.into_iter().map(|provider_stream| async move { let mut chunks = 0; loop { - let (sender, receiver) = tokio::sync::oneshot::channel(); - assert_eq!( - unsafe { - native_async_llm_stream_next_v2( - provider_stream, - record_typed_llm_stream_next, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = pull_typed_llm_stream(provider_stream); match receiver.await.expect("next callback should be delivered") { LlmContinuationStreamEventV2::Chunk { .. } => chunks += 1, LlmContinuationStreamEventV2::Done => break, @@ -3937,18 +3836,7 @@ fn native_api_v2_unary_entrypoint_reports_validation_panics_and_typed_failures() None, )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::oneshot::channel(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - next_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = start_typed_llm_call(next_ref, dispatch); let outcome = runtime.block_on(receiver).unwrap(); match outcome { LlmContinuationOutcomeV2::Failure { @@ -3972,18 +3860,7 @@ fn native_api_v2_unary_entrypoint_reports_validation_panics_and_typed_failures() None, )); let panicking_ref = Arc::into_raw(panicking) as *const NemoRelayNativeAsyncNext; - let (sender, receiver) = tokio::sync::oneshot::channel::(); - assert_eq!( - unsafe { - native_async_llm_next_invoke_result_v2( - panicking_ref, - dispatch, - complete_typed_llm_result, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = start_typed_llm_call(panicking_ref, dispatch); assert!(matches!( runtime.block_on(receiver).unwrap(), LlmContinuationOutcomeV2::Failure { @@ -4112,20 +3989,7 @@ fn native_api_v2_stream_open_reports_validation_setup_and_provider_failures() { 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 (sender, open_receiver) = - tokio::sync::oneshot::channel::>(); - assert_eq!( - unsafe { - native_async_llm_next_open_stream_v2( - next_ref, - dispatch, - output_ref, - record_typed_llm_stream_open, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let open_receiver = start_typed_llm_stream(next_ref, dispatch, output_ref); assert!(runtime.block_on(open_receiver).unwrap().is_err()); drop(receiver); unsafe { @@ -4147,20 +4011,7 @@ fn native_api_v2_stream_open_reports_validation_setup_and_provider_failures() { 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 (sender, open_receiver) = - tokio::sync::oneshot::channel::>(); - assert_eq!( - unsafe { - native_async_llm_next_open_stream_v2( - panicking_ref, - dispatch, - panic_output_ref, - record_typed_llm_stream_open, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let open_receiver = start_typed_llm_stream(panicking_ref, dispatch, panic_output_ref); assert!(matches!( runtime.block_on(open_receiver).unwrap(), Err(LlmContinuationFailureV2::NonHttp { @@ -4360,17 +4211,7 @@ fn native_api_v2_provider_stream_reports_terminal_and_release_states() { 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 (sender, receiver) = tokio::sync::oneshot::channel(); - assert_eq!( - unsafe { - native_async_llm_stream_next_v2( - stream_ref, - record_typed_llm_stream_next, - Box::into_raw(Box::new(sender)).cast(), - ) - }, - NemoRelayStatus::Ok - ); + let receiver = pull_typed_llm_stream(stream_ref); assert!(matches!( runtime.block_on(receiver).unwrap(), LlmContinuationStreamEventV2::Done @@ -4884,19 +4725,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 { - 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, - }); + 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( @@ -5116,19 +4945,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 { - 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, - }); + 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()); @@ -5398,19 +5215,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 { - 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, - }); + 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 { @@ -5517,19 +5322,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 { - 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, - }); + 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 { @@ -5640,19 +5433,7 @@ fn native_async_stream_next_stops_callbacks_after_false() { None, )); let next_ref = Arc::into_raw(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(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 { @@ -5716,19 +5497,7 @@ fn native_async_stream_in_flight_cancellation_releases_callback_state() { None, )); let next_ref = Arc::into_raw(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(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 { @@ -5800,19 +5569,7 @@ fn native_async_stream_cancellation_before_first_poll_releases_callback_state() None, )); let next_ref = Arc::into_raw(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(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 { @@ -6812,19 +6569,7 @@ fn native_async_stream_backpressure_preserves_v3_internal_and_v4_would_block() { .enable_all() .build() .unwrap(); - 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(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}"#); diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 6f28439fd..7895c785b 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -11,6 +11,7 @@ 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; @@ -949,7 +950,7 @@ pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn( /// Explicit provider target for one native API v2 LLM continuation. /// /// Relay sends LLM continuation requests with HTTP `POST`. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)] +#[derive(Clone, PartialEq, Eq, Serialize, serde::Deserialize)] pub struct LlmContinuationTargetV2 { /// Absolute HTTP(S) provider URL including the selected endpoint. pub url: String, @@ -960,8 +961,18 @@ pub struct LlmContinuationTargetV2 { 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(Debug, Clone, PartialEq, Serialize, serde::Deserialize)] +#[derive(Clone, PartialEq, Serialize, serde::Deserialize)] pub struct LlmContinuationInvocationV2 { /// Replacement request passed to the Relay execution continuation. pub request: LlmRequest, @@ -969,6 +980,16 @@ pub struct LlmContinuationInvocationV2 { 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")] @@ -987,6 +1008,14 @@ pub enum LlmNonHttpFailureKindV2 { Internal, } +impl LlmNonHttpFailureKindV2 { + /// Whether the provider-neutral retry policy retries this failure kind. + #[must_use] + pub const fn is_retryable(self) -> bool { + matches!(self, Self::Transport | Self::Timeout) + } +} + /// 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")] @@ -1009,6 +1038,23 @@ pub enum LlmContinuationFailureV2 { }, } +impl LlmContinuationFailureV2 { + /// Whether the provider-neutral retry policy retries an HTTP status. + #[must_use] + pub const fn http_status_is_retryable(status: u16) -> bool { + matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504) + } + + /// Whether a provider-neutral routing policy should retry this failure. + #[must_use] + pub fn is_retryable(&self) -> bool { + match self { + Self::Http { status, .. } => Self::http_status_is_retryable(*status), + Self::NonHttp { kind, .. } => kind.is_retryable(), + } + } +} + /// Receives one typed unary LLM continuation outcome. /// /// Exactly one of `response_json` and `error_json` is non-null. A response is @@ -3457,6 +3503,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, diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index b3a25d746..60b50b8c3 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -6,7 +6,6 @@ use std::ffi::c_void; use std::future::Future; use std::pin::Pin; -use std::ptr; use std::sync::Arc; use std::task::{Context, Poll}; @@ -87,11 +86,6 @@ struct CompletionHandle(*const NemoRelayNativeAsyncCompletion); #[derive(Clone, Copy)] struct OutputHandle(*const NemoRelayNativeAsyncStream); -struct TaskHostString { - host: NemoRelayNativeHostApiV1, - raw: usize, -} - // The host owns these opaque handles and documents their operations as // thread-safe for a pending callback's lifetime. unsafe impl Send for CompletionHandle {} @@ -99,34 +93,6 @@ unsafe impl Sync for CompletionHandle {} unsafe impl Send for OutputHandle {} unsafe impl Sync for OutputHandle {} -impl TaskHostString { - fn new(host: &NemoRelayNativeHostApiV1, value: &str) -> Option { - let mut raw = ptr::null_mut(); - let status = unsafe { (host.string_new)(value.as_ptr(), value.len(), &mut raw) }; - if status != NemoRelayStatus::Ok || raw.is_null() { - return None; - } - Some(Self { - host: *host, - raw: raw as usize, - }) - } - - fn from_json(host: &NemoRelayNativeHostApiV1, value: &Json) -> Option { - Self::new(host, &serde_json::to_string(value).ok()?) - } - - fn as_ptr(&self) -> *const NemoRelayNativeString { - self.raw as *const NemoRelayNativeString - } -} - -impl Drop for TaskHostString { - fn drop(&mut self) { - unsafe { (self.host.string_free)(self.raw as *mut NemoRelayNativeString) }; - } -} - impl CompletionHandle { fn as_ptr(&self) -> *const NemoRelayNativeAsyncCompletion { self.0 @@ -456,16 +422,10 @@ impl LlmStreamContinuationV2 { unsafe { drop(Box::from_raw(state.cast::())) }; return Err(status_failure("targeted LLM stream setup", status)); } - let provider = receiver.await.unwrap_or_else(|_| { + receiver.await.unwrap_or_else(|_| { Err(internal_failure( "targeted LLM stream setup callback closed without an outcome", )) - })?; - Ok(LlmProviderStreamV2 { - host, - raw: provider.into_raw(), - pending: None, - finished: false, }) } @@ -918,7 +878,7 @@ async fn pump_output_stream( async fn push_output_json(continuation: &LlmStreamContinuationV2, chunk: &Json) -> Result<()> { let host = &continuation.inner.host; - let chunk = TaskHostString::from_json(&host.v3.v1, chunk) + 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) } { @@ -953,7 +913,7 @@ async fn reject_output(host: &NemoRelayNativeHostApiV4, output: OutputHandle, er return; } let error = bounded_error(error); - let Some(message) = TaskHostString::new(&host.v3.v1, &error) else { + let Some(message) = HostString::new(&host.v3.v1, &error) else { set_last_error( &host.v3.v1, "failed to allocate native API v2 stream rejection", @@ -1124,37 +1084,9 @@ unsafe extern "C" fn passthrough_result_callback( struct StreamOpenCallback { host: NemoRelayNativeHostApiV4, - sender: oneshot::Sender>, -} - -struct OwnedProviderStream { - host: NemoRelayNativeHostApiV4, - raw: *const NemoRelayNativeLlmStreamV2, -} - -impl OwnedProviderStream { - fn new(host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeLlmStreamV2) -> Self { - debug_assert!(!raw.is_null()); - Self { host, raw } - } - - fn into_raw(mut self) -> *const NemoRelayNativeLlmStreamV2 { - std::mem::replace(&mut self.raw, ptr::null()) - } + sender: oneshot::Sender>, } -impl Drop for OwnedProviderStream { - fn drop(&mut self) { - if !self.raw.is_null() { - unsafe { (self.host.async_llm_stream_release_v2)(self.raw) }; - } - } -} - -// The host owns the opaque stream and documents release as thread-safe for a -// plugin-owned reference. -unsafe impl Send for OwnedProviderStream {} - unsafe extern "C" fn stream_open_callback( user_data: *mut c_void, stream: *const NemoRelayNativeLlmStreamV2, @@ -1164,8 +1096,12 @@ unsafe extern "C" fn stream_open_callback( return; } let state = unsafe { Box::from_raw(user_data.cast::()) }; - let mut owned_stream = - (!stream.is_null()).then(|| OwnedProviderStream::new(state.host, stream)); + 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(()), @@ -1189,7 +1125,7 @@ unsafe extern "C" fn stream_open_callback( )) }); let result = result.and_then(|()| { - owned_stream + provider .take() .ok_or_else(|| internal_failure("targeted LLM stream setup returned a null stream")) }); diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 9c5bd2a2f..788c154af 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -3,7 +3,7 @@ //! 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}; @@ -41,6 +41,7 @@ use nemo_relay_plugin::{ NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, ScopeType, ToolExecutionInterceptOutcome, ToolNext, }; +use serde::de::DeserializeOwned; use serde_json::{Map, json}; #[test] @@ -1574,6 +1575,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, @@ -6148,12 +6157,9 @@ unsafe extern "C" fn safe_v2_completion_resolve( value_json: *const NemoRelayNativeString, ) -> NemoRelayStatus { let host = test_host(); - let value = match required_host_string(&host, value_json) - .ok() - .and_then(|value| serde_json::from_str(&value).ok()) - { - Some(value) => value, - None => return NemoRelayStatus::InvalidJson, + 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() @@ -6194,6 +6200,27 @@ 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, @@ -6205,27 +6232,13 @@ unsafe extern "C" fn safe_v2_register_generic_async( free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { if kind != NemoRelayNativeAsyncMiddlewareKind::LlmExecutionIntercept as u32 { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return NemoRelayStatus::InvalidArg; - } - let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); - if status != NemoRelayStatus::Ok { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return status; + return unsafe { + safe_v2_reject_registration(NemoRelayStatus::InvalidArg, user_data, free_fn) + }; } - let host = test_host(); - let name = match required_host_string(&host, name) { + let name = match safe_v2_registration_name(name) { Ok(name) => name, - Err(status) => { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return status; - } + Err(status) => return unsafe { safe_v2_reject_registration(status, user_data, free_fn) }, }; replace_registration( &ASYNC_V2_REGISTRATION, @@ -6245,21 +6258,24 @@ unsafe extern "C" fn safe_v2_stream_push( chunk_json: *const NemoRelayNativeString, ) -> NemoRelayStatus { let host = test_host(); - let chunk = match required_host_string(&host, chunk_json) - .ok() - .and_then(|value| serde_json::from_str(&value).ok()) - { - Some(chunk) => chunk, - None => return NemoRelayStatus::InvalidJson, + let chunk = match required_host_json(&host, chunk_json) { + Ok(chunk) => chunk, + Err(status) => return status, }; - let status = SAFE_V2_OUTPUT_PUSH_STATUSES + 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::Ok { - SAFE_V2_OUTPUT.lock().unwrap().push(Ok(chunk)); - } else if status == NemoRelayStatus::WouldBlock { + if status == NemoRelayStatus::WouldBlock { SAFE_V2_CURRENT_TASK.with(|current| { let task = current.get(); if task != 0 { @@ -6284,20 +6300,9 @@ unsafe extern "C" fn safe_v2_stream_reject( let host = test_host(); let message = required_host_string(&host, message) .unwrap_or_else(|status| format!("invalid rejection: {status:?}")); - let status = SAFE_V2_OUTPUT_REJECT_STATUSES - .lock() - .unwrap() - .pop_front() - .unwrap_or(NemoRelayStatus::Ok); + let status = safe_v2_output_status(&SAFE_V2_OUTPUT_REJECT_STATUSES); if status == NemoRelayStatus::Ok { SAFE_V2_OUTPUT.lock().unwrap().push(Err(message)); - } else if status == NemoRelayStatus::WouldBlock { - SAFE_V2_CURRENT_TASK.with(|current| { - let task = current.get(); - if task != 0 { - let _ = unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; - } - }); } status } @@ -6330,22 +6335,9 @@ unsafe extern "C" fn safe_v2_register_generic_stream( user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { - let status = *SAFE_V2_REGISTRATION_STATUS.lock().unwrap(); - if status != NemoRelayStatus::Ok { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return status; - } - let host = test_host(); - let name = match required_host_string(&host, name) { + let name = match safe_v2_registration_name(name) { Ok(name) => name, - Err(status) => { - if let Some(free_fn) = free_fn { - unsafe { free_fn(user_data) }; - } - return status; - } + Err(status) => return unsafe { safe_v2_reject_registration(status, user_data, free_fn) }, }; replace_registration( &ASYNC_STREAM_V2_REGISTRATION, @@ -6371,12 +6363,8 @@ unsafe extern "C" fn safe_v2_passthrough_result( return status; } let host = test_host(); - if required_host_string(&host, invocation_json) - .ok() - .and_then(|value| serde_json::from_str::(&value).ok()) - .is_none() - { - return NemoRelayStatus::InvalidJson; + 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); @@ -6401,12 +6389,8 @@ unsafe extern "C" fn safe_v2_targeted_result( return status; } let host = test_host(); - if required_host_string(&host, invocation_json) - .ok() - .and_then(|value| serde_json::from_str::(&value).ok()) - .is_none() - { - return NemoRelayStatus::InvalidJson; + 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)); @@ -6440,12 +6424,8 @@ unsafe extern "C" fn safe_v2_stream_open( return status; } let host = test_host(); - if required_host_string(&host, invocation_json) - .ok() - .and_then(|value| serde_json::from_str::(&value).ok()) - .is_none() - { - return NemoRelayStatus::InvalidJson; + 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)); @@ -6646,12 +6626,9 @@ unsafe extern "C" fn safe_v2_forward_stream( return status; } let host = test_host(); - let request = match required_host_string(&host, request_json) - .ok() - .and_then(|value| serde_json::from_str::(&value).ok()) - { - Some(request) => request, - None => return NemoRelayStatus::InvalidJson, + 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); @@ -6702,6 +6679,67 @@ fn safe_v2_target() -> LlmContinuationTargetV2 { } } +#[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)); + } +} + +#[test] +fn native_v2_retry_helper_uses_provider_neutral_status_semantics() { + let http_failure = |status| LlmContinuationFailureV2::Http { + status, + body: String::new(), + headers: BTreeMap::new(), + }; + for status in [408, 425, 429, 500, 502, 503, 504] { + assert!(LlmContinuationFailureV2::http_status_is_retryable(status)); + assert!(http_failure(status).is_retryable(), "status={status}"); + } + for status in [400, 401, 404, 409, 422, 501] { + assert!(!LlmContinuationFailureV2::http_status_is_retryable(status)); + assert!(!http_failure(status).is_retryable(), "status={status}"); + } + + let non_http_failure = |kind| LlmContinuationFailureV2::NonHttp { + kind, + message: String::new(), + }; + assert!(LlmNonHttpFailureKindV2::Transport.is_retryable()); + assert!(non_http_failure(LlmNonHttpFailureKindV2::Transport).is_retryable()); + assert!(LlmNonHttpFailureKindV2::Timeout.is_retryable()); + assert!(non_http_failure(LlmNonHttpFailureKindV2::Timeout).is_retryable()); + for kind in [ + LlmNonHttpFailureKindV2::Cancelled, + LlmNonHttpFailureKindV2::InvalidRequest, + LlmNonHttpFailureKindV2::Guardrail, + LlmNonHttpFailureKindV2::Internal, + ] { + assert!(!kind.is_retryable(), "kind={kind:?}"); + assert!(!non_http_failure(kind).is_retryable(), "kind={kind:?}"); + } +} + fn take_safe_v2_buffered_registration() -> RegisteredAsyncV2 { ASYNC_V2_REGISTRATION.lock().unwrap().take().unwrap() } @@ -6818,6 +6856,18 @@ async fn safe_v2_targeted_provider_stream( ))) } +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, @@ -7559,7 +7609,7 @@ fn safe_v2_stream_open_preserves_structured_failure() { } #[test] -fn safe_v2_malformed_stream_open_releases_the_owned_stream() { +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, @@ -8072,9 +8122,7 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { let host = test_host_v4(); *SAFE_V2_STREAM_OPEN_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; - run_safe_v2_streaming(&host, "open-status", |_, request, next| { - safe_v2_targeted_provider_stream(request, next) - }); + run_safe_v2_targeted_stream(&host, "open-status"); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), [Err(error)] if error.contains("InvalidArg") @@ -8083,9 +8131,7 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { SAFE_V2_OUTPUT.lock().unwrap().clear(); *SAFE_V2_PROVIDER_NEXT_STATUS.lock().unwrap() = NemoRelayStatus::InvalidArg; - run_safe_v2_streaming(&host, "poll-status", |_, request, next| { - safe_v2_targeted_provider_stream(request, next) - }); + run_safe_v2_targeted_stream(&host, "poll-status"); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), [Err(error)] if error.contains("InvalidArg") @@ -8102,9 +8148,7 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { message: "late provider failure".into(), }, )); - run_safe_v2_streaming(&host, "late-failure", |_, request, next| { - safe_v2_targeted_provider_stream(request, next) - }); + 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") @@ -8112,9 +8156,7 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { SAFE_V2_OUTPUT.lock().unwrap().clear(); *SAFE_V2_PROVIDER_EVENT_JSON.lock().unwrap() = Some("not-json".into()); - run_safe_v2_streaming(&host, "malformed-event", |_, request, next| { - safe_v2_targeted_provider_stream(request, next) - }); + run_safe_v2_targeted_stream(&host, "malformed-event"); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), [Err(error)] if error.contains("stream chunk") @@ -8122,9 +8164,7 @@ fn safe_v2_provider_stream_preserves_late_and_boundary_failures() { SAFE_V2_OUTPUT.lock().unwrap().clear(); SAFE_V2_PROVIDER_INVALID_EVENT.store(true, Ordering::SeqCst); - run_safe_v2_streaming(&host, "invalid-event", |_, request, next| { - safe_v2_targeted_provider_stream(request, next) - }); + run_safe_v2_targeted_stream(&host, "invalid-event"); assert!(matches!( SAFE_V2_OUTPUT.lock().unwrap().as_slice(), [Err(error)] if error.contains("invalid event") @@ -8142,9 +8182,7 @@ fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors .unwrap() .extend([NemoRelayStatus::WouldBlock, NemoRelayStatus::Ok]); run_safe_v2_streaming(&host, "push-backpressure", |_, _, _| async { - Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( - async { Ok(json!({ "chunk": true })) }, - )))) + Ok(safe_v2_one_chunk_stream()) }); assert_eq!( *SAFE_V2_OUTPUT.lock().unwrap(), @@ -8152,35 +8190,31 @@ fn safe_v2_stream_output_handles_backpressure_cancellation_and_settlement_errors ); SAFE_V2_OUTPUT.lock().unwrap().clear(); - SAFE_V2_OUTPUT_PUSH_STATUSES - .lock() - .unwrap() - .push_back(NemoRelayStatus::Internal); - run_safe_v2_streaming(&host, "push-internal", |_, _, _| async { - Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( - async { Ok(json!({ "chunk": true })) }, - )))) - }); - assert!(matches!( - SAFE_V2_OUTPUT.lock().unwrap().as_slice(), - [Err(error)] if error.contains("output push failed: Internal") - )); - SAFE_V2_OUTPUT.lock().unwrap().clear(); - - SAFE_V2_OUTPUT_PUSH_STATUSES - .lock() - .unwrap() - .push_back(NemoRelayStatus::InvalidArg); - run_safe_v2_streaming(&host, "push-failure", |_, _, _| async { - Ok(LlmStreamExecutionOutcomeV2::Stream(Box::pin(stream::once( - async { Ok(json!({ "chunk": true })) }, - )))) - }); - assert!(matches!( - SAFE_V2_OUTPUT.lock().unwrap().as_slice(), - [Err(error)] if error.contains("output push failed") - )); - 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 { From 53ba86a90bb7f3406171fc8dd68a7e0f52774ce9 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 15:50:01 -0600 Subject: [PATCH 19/32] refactor(plugin): trim native v2 policy surface Signed-off-by: Bryan Bednarski --- .../src/api/runtime/continuation_context.rs | 2 +- .../src/api/runtime/llm_dispatch_context.rs | 23 ++----- crates/core/src/plugin/dynamic/native.rs | 11 +-- crates/core/tests/unit/native_plugin_tests.rs | 68 +++++++++++-------- crates/plugin/src/lib.rs | 31 +-------- crates/plugin/src/native_v2.rs | 24 ++----- crates/plugin/tests/typed_callbacks.rs | 46 ++----------- 7 files changed, 58 insertions(+), 147 deletions(-) diff --git a/crates/core/src/api/runtime/continuation_context.rs b/crates/core/src/api/runtime/continuation_context.rs index d21c2a4fa..02fb577ee 100644 --- a/crates/core/src/api/runtime/continuation_context.rs +++ b/crates/core/src/api/runtime/continuation_context.rs @@ -108,7 +108,7 @@ impl MiddlewareContinuationContext { /// Invoke a callback and poll its future with the captured Relay context and typed LLM target. #[doc(hidden)] - pub async fn invoke_with_llm_dispatch_target( + pub(crate) async fn invoke_with_llm_dispatch_target( &self, target: LlmDispatchTargetContext, callback: C, diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 219c11ba6..c05dc4306 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -45,7 +45,8 @@ struct LlmDispatchTargetBinding { /// transport routing cannot leak into provider JSON or observability payloads. #[doc(hidden)] #[derive(Clone)] -pub struct LlmDispatchTargetContext { +#[cfg_attr(test, derive(PartialEq, Eq))] +pub(crate) struct LlmDispatchTargetContext { url: Url, headers: HeaderMap, } @@ -93,20 +94,6 @@ impl LlmDispatchTargetContext { headers: validated_headers, }) } - - /// Absolute provider URL selected for this invocation. - #[doc(hidden)] - #[must_use] - pub(crate) fn url(&self) -> &Url { - &self.url - } - - /// Explicit provider headers selected for this invocation. - #[doc(hidden)] - #[must_use] - pub(crate) fn headers(&self) -> &HeaderMap { - &self.headers - } } impl fmt::Debug for LlmDispatchTargetContext { @@ -273,8 +260,8 @@ async fn send( ) -> 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() { + 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 { @@ -321,7 +308,7 @@ fn transport_error(target: &LlmDispatchTargetContext, error: reqwest::Error) -> log::warn!( target: "nemo_relay.runtime", event = "targeted_llm_transport_failed", - provider_host = target.url().host_str().unwrap_or(""), + provider_host = target.url.host_str().unwrap_or(""), failure_kind = if timeout { "timeout" } else { "transport" }; "Targeted LLM provider request failed: {diagnostic}" ); diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index fe17b1c34..4734146cc 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -1713,7 +1713,6 @@ struct NativeAsyncTaskStateV2 { } struct NativeAsyncTaskV2 { - runtime: tokio::runtime::Handle, context: MiddlewareContinuationContext, owner: NativeAsyncTaskOwnerV2, wake: tokio::sync::Notify, @@ -2752,7 +2751,6 @@ fn spawn_native_async_task_v2( } let (cb, user_data, free_fn) = callback; let task = Arc::new(NativeAsyncTaskV2 { - runtime, context, owner, wake: tokio::sync::Notify::new(), @@ -2770,7 +2768,7 @@ fn spawn_native_async_task_v2( }); *slot = Some(Arc::downgrade(&task)); drop(slot); - task.runtime.spawn(Arc::clone(&task).run()); + runtime.spawn(Arc::clone(&task).run()); NativeAsyncTaskV2::wake(&task); NemoRelayStatus::Ok } @@ -2781,16 +2779,13 @@ unsafe extern "C" fn native_async_task_retain_v2(task: *const NemoRelayNativeAsy } } -unsafe extern "C" fn native_async_task_wake_v2( - task: *const NemoRelayNativeAsyncTaskV2, -) -> NemoRelayStatus { +unsafe extern "C" fn native_async_task_wake_v2(task: *const NemoRelayNativeAsyncTaskV2) { if task.is_null() { - return NemoRelayStatus::NullPointer; + return; } unsafe { Arc::increment_strong_count(task as *const NativeAsyncTaskV2) }; let task = unsafe { Arc::from_raw(task as *const NativeAsyncTaskV2) }; NativeAsyncTaskV2::wake(&task); - NemoRelayStatus::Ok } unsafe extern "C" fn native_async_task_release_v2(task: *const NemoRelayNativeAsyncTaskV2) { diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 2b92c4489..49eeb0511 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -1965,15 +1965,12 @@ fn native_api_v2_unary_dispatch_uses_explicit_target_and_structured_failures() { Box::pin(async move { assert!(request.headers.is_empty()); assert_eq!( - target.url().as_str(), - "https://provider.example/v1/chat/completions" - ); - assert_eq!( - target - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some("Bearer target-secret") + 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), @@ -2293,11 +2290,12 @@ fn native_api_v2_stream_dispatch_reports_chunks_and_a_typed_late_failure() { assert!(request.headers.is_empty()); Box::pin(async move { assert_eq!( - current_llm_dispatch_target() - .expect("typed target is bound") - .url() - .as_str(), - "https://provider.example/v1/messages" + 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"}})), @@ -2433,10 +2431,7 @@ fn native_api_v2_provider_pulls_restore_scope_and_target_after_pending_wakes() { 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") - .url() - .as_str() - .to_owned(); + .expect("provider pull should restore its dispatch target"); observations_for_stream .lock() .unwrap_or_else(|error| error.into_inner()) @@ -2478,8 +2473,14 @@ fn native_api_v2_provider_pulls_restore_scope_and_target_after_pending_wakes() { .lock() .unwrap_or_else(|error| error.into_inner()), vec![ - (captured_scope, target_url.into()), - (captured_scope, target_url.into()), + ( + 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) }; @@ -2967,13 +2968,25 @@ fn native_api_v2_isolates_concurrent_dispatch_targets() { 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": target.url().as_str(), - "authorization": target - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), + "url": url, + "authorization": authorization, })) }) })), @@ -5690,10 +5703,7 @@ fn native_async_task_pending_wake_settles_and_frees_once() { .recv_timeout(Duration::from_secs(1)) .expect("first poll should retain a task waker") as *const NemoRelayNativeAsyncTaskV2; - assert_eq!( - unsafe { native_async_task_wake_v2(task) }, - NemoRelayStatus::Ok - ); + 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 }) diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 7895c785b..d5ee6b662 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -1008,14 +1008,6 @@ pub enum LlmNonHttpFailureKindV2 { Internal, } -impl LlmNonHttpFailureKindV2 { - /// Whether the provider-neutral retry policy retries this failure kind. - #[must_use] - pub const fn is_retryable(self) -> bool { - matches!(self, Self::Transport | Self::Timeout) - } -} - /// 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")] @@ -1038,23 +1030,6 @@ pub enum LlmContinuationFailureV2 { }, } -impl LlmContinuationFailureV2 { - /// Whether the provider-neutral retry policy retries an HTTP status. - #[must_use] - pub const fn http_status_is_retryable(status: u16) -> bool { - matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504) - } - - /// Whether a provider-neutral routing policy should retry this failure. - #[must_use] - pub fn is_retryable(&self) -> bool { - match self { - Self::Http { status, .. } => Self::http_status_is_retryable(*status), - Self::NonHttp { kind, .. } => kind.is_retryable(), - } - } -} - /// Receives one typed unary LLM continuation outcome. /// /// Exactly one of `response_json` and `error_json` is non-null. A response is @@ -1330,8 +1305,7 @@ pub struct NemoRelayNativeHostApiV4 { /// 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) -> NemoRelayStatus, + 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 @@ -1349,9 +1323,6 @@ pub struct NemoRelayNativeHostApiV4 { ) -> NemoRelayStatus, } -unsafe impl Send for NemoRelayNativeHostApiV4 {} -unsafe impl Sync for NemoRelayNativeHostApiV4 {} - unsafe impl Send for NemoRelayNativeHostApiV3 {} unsafe impl Sync for NemoRelayNativeHostApiV3 {} diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index 60b50b8c3..283795087 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -181,24 +181,15 @@ impl HostTaskWaker { unsafe fn new( host: NemoRelayNativeHostApiV4, raw: *const NemoRelayNativeAsyncTaskV2, - ) -> Option> { - if raw.is_null() { - return None; - } + ) -> Arc { unsafe { (host.async_task_retain_v2)(raw) }; - Some(Arc::new(Self { host, raw })) + Arc::new(Self { host, raw }) } } impl ArcWake for HostTaskWaker { fn wake_by_ref(arc_self: &Arc) { - let status = unsafe { (arc_self.host.async_task_wake_v2)(arc_self.raw) }; - if status != NemoRelayStatus::Ok { - set_last_error( - &arc_self.host.v3.v1, - &format!("native API v2 host task wake failed: {status:?}"), - ); - } + unsafe { (arc_self.host.async_task_wake_v2)(arc_self.raw) }; } } @@ -220,14 +211,7 @@ unsafe extern "C" fn poll_host_future_task( return NemoRelayNativeAsyncCallbackState::Complete as u32; } if state.waker.is_none() { - let waker = match unsafe { HostTaskWaker::new(state.host, task) } { - Some(waker) => waker, - None => { - set_last_error(&state.host.v3.v1, "native API v2 host task was null"); - return NemoRelayNativeAsyncCallbackState::Complete as u32; - } - }; - state.waker = Some(waker); + state.waker = Some(unsafe { HostTaskWaker::new(state.host, task) }); } let waker = waker_ref( state diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 788c154af..1aa9bb4fa 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -6279,7 +6279,7 @@ fn safe_v2_output_status(statuses: &Mutex>) -> NemoRel SAFE_V2_CURRENT_TASK.with(|current| { let task = current.get(); if task != 0 { - let _ = unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; } }); } @@ -6545,15 +6545,14 @@ unsafe extern "C" fn safe_v2_task_retain(task: *const NemoRelayNativeAsyncTaskV2 } } -unsafe extern "C" fn safe_v2_task_wake(task: *const NemoRelayNativeAsyncTaskV2) -> NemoRelayStatus { +unsafe extern "C" fn safe_v2_task_wake(task: *const NemoRelayNativeAsyncTaskV2) { let Some(task) = (unsafe { task.cast::().as_ref() }) else { - return NemoRelayStatus::NullPointer; + return; }; if task.completed.load(Ordering::Acquire) { - return NemoRelayStatus::Ok; + return; } task.woken.store(true, Ordering::Release); - NemoRelayStatus::Ok } unsafe extern "C" fn safe_v2_task_release(task: *const NemoRelayNativeAsyncTaskV2) { @@ -6568,7 +6567,7 @@ unsafe extern "C" fn safe_v2_task_release(task: *const NemoRelayNativeAsyncTaskV fn wake_safe_v2_tasks() { for task in SAFE_V2_TASKS.lock().unwrap().iter().copied() { - let _ = unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; + unsafe { safe_v2_task_wake(task as *const NemoRelayNativeAsyncTaskV2) }; } } @@ -6705,41 +6704,6 @@ fn native_v2_debug_output_redacts_requests_targets_and_credentials() { } } -#[test] -fn native_v2_retry_helper_uses_provider_neutral_status_semantics() { - let http_failure = |status| LlmContinuationFailureV2::Http { - status, - body: String::new(), - headers: BTreeMap::new(), - }; - for status in [408, 425, 429, 500, 502, 503, 504] { - assert!(LlmContinuationFailureV2::http_status_is_retryable(status)); - assert!(http_failure(status).is_retryable(), "status={status}"); - } - for status in [400, 401, 404, 409, 422, 501] { - assert!(!LlmContinuationFailureV2::http_status_is_retryable(status)); - assert!(!http_failure(status).is_retryable(), "status={status}"); - } - - let non_http_failure = |kind| LlmContinuationFailureV2::NonHttp { - kind, - message: String::new(), - }; - assert!(LlmNonHttpFailureKindV2::Transport.is_retryable()); - assert!(non_http_failure(LlmNonHttpFailureKindV2::Transport).is_retryable()); - assert!(LlmNonHttpFailureKindV2::Timeout.is_retryable()); - assert!(non_http_failure(LlmNonHttpFailureKindV2::Timeout).is_retryable()); - for kind in [ - LlmNonHttpFailureKindV2::Cancelled, - LlmNonHttpFailureKindV2::InvalidRequest, - LlmNonHttpFailureKindV2::Guardrail, - LlmNonHttpFailureKindV2::Internal, - ] { - assert!(!kind.is_retryable(), "kind={kind:?}"); - assert!(!non_http_failure(kind).is_retryable(), "kind={kind:?}"); - } -} - fn take_safe_v2_buffered_registration() -> RegisteredAsyncV2 { ASYNC_V2_REGISTRATION.lock().unwrap().take().unwrap() } From 5bbbce16fc2a16e2d747e42191fe506c8d984564 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 16:09:11 -0600 Subject: [PATCH 20/32] fix(plugin): preserve targets across continuation hops Signed-off-by: Bryan Bednarski --- .../src/api/runtime/continuation_context.rs | 23 +++++++-- .../tests/unit/continuation_context_tests.rs | 48 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/crates/core/src/api/runtime/continuation_context.rs b/crates/core/src/api/runtime/continuation_context.rs index 02fb577ee..5257bac53 100644 --- a/crates/core/src/api/runtime/continuation_context.rs +++ b/crates/core/src/api/runtime/continuation_context.rs @@ -9,7 +9,7 @@ use crate::api::optimization::{ LlmOptimizationRecorder, current_llm_optimization_recorder, scope_llm_optimization_recorder, }; use crate::api::runtime::llm_dispatch_context::{ - LlmDispatchTargetContext, scope_llm_dispatch_target, + 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, @@ -34,6 +34,7 @@ pub struct MiddlewareContinuationContext { publication_context: Option, publication_buffer: Option, optimization_recorder: Option, + llm_dispatch_target: Option, } impl MiddlewareContinuationContext { @@ -47,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(), } } @@ -72,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(), }) } @@ -100,9 +103,17 @@ 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, } } @@ -117,7 +128,9 @@ impl MiddlewareContinuationContext { C: FnOnce() -> F, F: Future, { - scope_llm_dispatch_target(self.active_event_uuid, target, self.invoke(callback)).await + 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. 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(); From ad2f60e82adef52d52c9e2336dd4a66e4b1f55ec Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 16:20:31 -0600 Subject: [PATCH 21/32] docs(plugin): clarify cooperative continuation semantics Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 7 ++- crates/plugin/README.md | 47 +++++++------- crates/plugin/src/lib.rs | 4 +- .../dynamic-plugins/native-dynamic/about.mdx | 63 +++++++++++-------- 4 files changed, 68 insertions(+), 53 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 4734146cc..938a53e7f 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -1517,9 +1517,10 @@ fn make_user_data( }) } -// Native API v1's V3 incremental stream contract shipped with a 64-event -// queue. Keep that observable backpressure boundary stable for existing -// plugins while native API v2 uses the documented tighter bound. +// Native API v1's V3 incremental plugin-output stream shipped with a 64-event +// queue. Keep that boundary stable for existing plugins. Native API v2 uses a +// 32-event queue for plugin output and direct pass-through; targeted provider +// streams are pull-based instead. const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1: usize = 64; const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2: usize = 32; diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 8c75bb473..61d705a04 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -49,7 +49,7 @@ the dynamic-library boundary on the stable C-compatible ABI. [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. @@ -147,18 +147,20 @@ ctx.register_async_llm_stream_execution_v2("route-stream", 0, move |_name, reque })?; ``` -For an unmanaged buffered request, call -`LlmContinuationV2::call_passthrough`. For an unmanaged streaming request, -return `LlmStreamExecutionOutcomeV2::Passthrough(request)`. Relay then pumps -the original downstream stream directly through its bounded queue; provider -events do not cross into the plugin merely to be forwarded. - -The plugin provides JSON plus an HTTP method, absolute target URL, and explicit -target headers. Relay binds that transport target to the current LLM -continuation without storing it in `LlmRequest.headers`. Successful calls return -provider JSON. 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. +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. + +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 calls return provider JSON. 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 @@ -166,14 +168,17 @@ 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. Safe callbacks register -through the generic V3 completion API 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. +`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 +32-event 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 diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index d5ee6b662..6b4616b02 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -2453,8 +2453,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, diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 3ecc4e838..29d4e43c4 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -132,11 +132,12 @@ 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 HTTP method, absolute -HTTP(S) URL, and explicit outbound headers. 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. +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 @@ -150,6 +151,11 @@ request middleware. Relay returns either: - 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. @@ -163,10 +169,10 @@ 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 stream production and plugin output -each use bounded 32-event queues. Dropping or cancelling a stream stops provider -production, and the library stays loaded until all callbacks and streams -release their handles. +`x-nemo-relay-internal-*` headers. Provider streams are pull-based and permit +one pending poll; plugin output and direct pass-through use bounded 32-event +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 @@ -175,12 +181,13 @@ 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 completion API and return -`Pending`. 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. +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` @@ -222,11 +229,12 @@ context.register_async_llm_stream_execution_v2( )?; ``` -Use `LlmContinuationV2::call_passthrough` for an unmanaged buffered call. For -an unmanaged stream, return -`LlmStreamExecutionOutcomeV2::Passthrough(request)`. Relay connects the -original downstream stream directly to the caller through its bounded queue, -so pass-through events do not cross the plugin boundary. +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. The SDK returns `Pending` to the host, which drives safe callback futures and returned streams as cooperative tasks. `LlmProviderStreamV2` implements @@ -264,13 +272,14 @@ extern "C" fn nemo_relay_register_plugin( `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 registers callbacks through V3's completion-based `Pending` -contract. 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. +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 From 3861d63e3ff9dce4e7e2612fb2493dc1b590d745 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 17:05:23 -0600 Subject: [PATCH 22/32] fix(runtime): bound targeted buffered responses Signed-off-by: Bryan Bednarski --- .../src/api/runtime/llm_dispatch_context.rs | 42 +++++++++++++++++-- .../tests/unit/llm_dispatch_context_tests.rs | 27 ++++++++++++ crates/plugin/README.md | 8 ++-- .../dynamic-plugins/native-dynamic/about.mdx | 2 +- 4 files changed, 70 insertions(+), 9 deletions(-) diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index c05dc4306..492c28a8e 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -27,6 +27,7 @@ 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; } @@ -201,10 +202,7 @@ async fn dispatch_buffered(target: &LlmDispatchTargetContext, request: LlmReques let bytes = bounded_response_body(target, response).await?; return Err(http_error(status, headers, &bytes)); } - let bytes = response - .bytes() - .await - .map_err(|error| transport_error(target, error))?; + 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()) }) @@ -285,6 +283,42 @@ fn targeted_http_client() -> &'static Client { }) } +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, diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index b3e3735d8..e7bc6b3ea 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -231,6 +231,33 @@ async fn malformed_success_json_is_an_internal_provider_failure() { 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]; diff --git a/crates/plugin/README.md b/crates/plugin/README.md index 61d705a04..b02353306 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -157,10 +157,10 @@ to be forwarded. 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 calls return provider JSON. 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. +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 diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 29d4e43c4..445c97b14 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -143,7 +143,7 @@ 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; +- 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 From 82141cf8c110b8b376fb7e303f0ac5c050eeefc7 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 17:32:42 -0600 Subject: [PATCH 23/32] refactor(plugin): share bounded native stream capacity Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/native.rs | 24 ++++--------------- crates/core/tests/unit/native_plugin_tests.rs | 8 +------ crates/plugin/README.md | 2 +- .../dynamic-plugins/native-dynamic/about.mdx | 4 ++-- 4 files changed, 9 insertions(+), 29 deletions(-) diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 938a53e7f..a474d6d47 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -307,7 +307,6 @@ struct NativePluginInstance { plugin_kind: String, relay_compat: String, allows_multiple_components: bool, - uses_native_api_v2: bool, plugin: Mutex, library: Option, retain_library_on_drop: AtomicBool, @@ -466,7 +465,6 @@ fn load_one_native_plugin( plugin_kind, relay_compat, allows_multiple_components: plugin.allows_multiple_components, - uses_native_api_v2: native_api == Some("2"), plugin: Mutex::new(plugin), library: Some(library), retain_library_on_drop: AtomicBool::new(false), @@ -1517,20 +1515,9 @@ fn make_user_data( }) } -// Native API v1's V3 incremental plugin-output stream shipped with a 64-event -// queue. Keep that boundary stable for existing plugins. Native API v2 uses a -// 32-event queue for plugin output and direct pass-through; targeted provider -// streams are pull-based instead. -const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1: usize = 64; -const NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2: usize = 32; - -fn native_async_stream_channel_capacity(uses_native_api_v2: bool) -> usize { - if uses_native_api_v2 { - NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2 - } else { - NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1 - } -} +// 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 { sender: Mutex>>>, @@ -4176,12 +4163,11 @@ fn wrap_native_incremental_llm_stream_execution( user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> LlmStreamExecutionFn { - let channel_capacity = native_async_stream_channel_capacity(instance.uses_native_api_v2); let user_data = make_user_data(instance, user_data, free_fn); wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( cb, user_data, - channel_capacity, + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY, ) } @@ -4193,7 +4179,7 @@ fn wrap_native_incremental_llm_stream_execution_with_user_data( wrap_native_incremental_llm_stream_execution_with_user_data_and_capacity( cb, user_data, - NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V1, + NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY, ) } diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 49eeb0511..8ecc6db58 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -3064,7 +3064,7 @@ fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_d )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; let (output_stream, output_receiver) = - test_native_output_stream(NATIVE_ASYNC_STREAM_CHANNEL_CAPACITY_V2); + 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( @@ -3138,12 +3138,6 @@ fn native_api_v2_direct_pulls_64_concurrent_100_event_provider_streams_without_d } } -#[test] -fn native_async_stream_capacity_preserves_v1_and_bounds_v2() { - assert_eq!(native_async_stream_channel_capacity(false), 64); - assert_eq!(native_async_stream_channel_capacity(true), 32); -} - fn test_v2_callback_user_data(ptr: *mut c_void) -> Arc { Arc::new(NativeCallbackUserData { ptr, diff --git a/crates/plugin/README.md b/crates/plugin/README.md index b02353306..d8fa749d6 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -172,7 +172,7 @@ Streaming dispatch returns `LlmProviderStreamV2`, which implements Rust 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 -32-event queues. Safe callbacks register through the generic V3 +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 diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 445c97b14..666abaff8 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -170,8 +170,8 @@ 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 32-event -queues. Dropping or cancelling a stream stops provider production, and the +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 From 85e72ed982477db518b6903a1f86324a836b809c Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 21:41:45 -0600 Subject: [PATCH 24/32] fix(plugin): address native API v2 review findings Signed-off-by: Bryan Bednarski --- .../src/api/runtime/llm_dispatch_context.rs | 34 +++++++++++++- .../tests/fixtures/native_plugin/src/lib.rs | 4 +- .../tests/integration/native_plugin_tests.rs | 47 +++++++++++++------ .../tests/unit/llm_dispatch_context_tests.rs | 32 ++++++++++++- crates/core/tests/unit/native_plugin_tests.rs | 15 ++---- crates/plugin/src/lib.rs | 23 ++++++++- crates/plugin/src/native_v2.rs | 12 +++++ docs/build-plugins/dynamic-plugins/about.mdx | 2 +- .../dynamic-plugins/native-dynamic/about.mdx | 4 +- docs/reference/migration-guides.mdx | 2 +- 10 files changed, 142 insertions(+), 33 deletions(-) diff --git a/crates/core/src/api/runtime/llm_dispatch_context.rs b/crates/core/src/api/runtime/llm_dispatch_context.rs index 492c28a8e..decdc75d4 100644 --- a/crates/core/src/api/runtime/llm_dispatch_context.rs +++ b/crates/core/src/api/runtime/llm_dispatch_context.rs @@ -20,7 +20,7 @@ use crate::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionN use crate::codec::streaming::SseEventDecoder; use crate::error::{ FlowError, MAX_UPSTREAM_FAILURE_BODY_BYTES, Result, UpstreamFailure, UpstreamFailureClass, - sanitize_upstream_failure_headers, + bounded_utf8, sanitize_upstream_failure_headers, }; use crate::json::Json; @@ -219,6 +219,16 @@ async fn dispatch_stream( 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(); @@ -251,6 +261,28 @@ async fn dispatch_stream( })) } +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, diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 6d38e73c7..b44d7aa72 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -16,7 +16,7 @@ use nemo_relay_plugin::{ NemoRelayNativeAsyncStream, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, - ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, + ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, NEMO_RELAY_NATIVE_ABI_VERSION, }; use serde_json::{Map, json}; @@ -454,7 +454,7 @@ pub unsafe extern "C" fn nemo_relay_fixture_native_api_v1_plugin( let Some(host_ref) = (unsafe { host.as_ref() }) else { return NemoRelayStatus::NullPointer; }; - if host_ref.abi_version != 3 + if host_ref.abi_version != NEMO_RELAY_NATIVE_ABI_VERSION || host_ref.struct_size != std::mem::size_of::() { return NemoRelayStatus::InvalidArg; diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index c830549cf..230b8843b 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -2614,7 +2614,7 @@ impl Drop for PendingDropStream { struct EmbeddedFakeProvider { url: String, request: std::sync::mpsc::Receiver>, - thread: Option>, + thread: Option>>, } impl EmbeddedFakeProvider { @@ -2629,8 +2629,8 @@ impl EmbeddedFakeProvider { .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 || { - let deadline = std::time::Instant::now() + Duration::from_secs(5); + 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, @@ -2640,19 +2640,30 @@ impl EmbeddedFakeProvider { { std::thread::sleep(Duration::from_millis(10)); } - Err(error) => panic!("embedded provider should accept: {error}"), + 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) - .expect("embedded provider socket should be blocking"); + 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))) - .expect("embedded provider timeout should configure"); + .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).expect("provider request read"); + let read = socket.read(&mut buffer).map_err(|error| { + format!("embedded provider failed to read request: {error}") + })?; if read == 0 { break; } @@ -2679,17 +2690,18 @@ impl EmbeddedFakeProvider { } request_tx .send(request_bytes) - .expect("embedded test should receive provider request"); + .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()) - .expect("embedded provider should write headers"); + .map_err(|error| format!("embedded provider failed to write headers: {error}"))?; socket .write_all(body) - .expect("embedded provider should write body"); + .map_err(|error| format!("embedded provider failed to write body: {error}"))?; + Ok(()) }); Self { url: format!("http://{address}/v1/chat/completions"), @@ -2708,9 +2720,14 @@ impl EmbeddedFakeProvider { impl Drop for EmbeddedFakeProvider { fn drop(&mut self) { if let Some(thread) = self.thread.take() { - thread - .join() - .expect("embedded provider thread should finish"); + 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), + } + } } } } diff --git a/crates/core/tests/unit/llm_dispatch_context_tests.rs b/crates/core/tests/unit/llm_dispatch_context_tests.rs index e7bc6b3ea..ed769ce22 100644 --- a/crates/core/tests/unit/llm_dispatch_context_tests.rs +++ b/crates/core/tests/unit/llm_dispatch_context_tests.rs @@ -339,7 +339,7 @@ async fn redirects_are_returned_without_following() { async fn streaming_target_decodes_events_empty_streams_and_late_errors() { let provider = FakeProvider::spawn(response( "200 OK", - &[("Content-Type", "text/event-stream")], + &[("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)); @@ -397,6 +397,36 @@ async fn streaming_target_decodes_events_empty_streams_and_late_errors() { 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 [ diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 8ecc6db58..6430e92de 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -522,24 +522,19 @@ unsafe extern "C" fn poll_cooperative_task( 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 { native_async_stream_push_json(stream, first) }, - NemoRelayStatus::Ok - ); - assert_eq!( - unsafe { native_async_stream_push_json(stream, second) }, - NemoRelayStatus::Internal + 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 { native_async_stream_push_json(stream, second) }, - NemoRelayStatus::Ok - ); + assert_eq!(unsafe { push_json(stream, second) }, NemoRelayStatus::Ok); assert_eq!( unsafe { native_async_stream_finish(stream) }, NemoRelayStatus::Ok diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 6b4616b02..18e0f1d32 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -1165,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, @@ -1218,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, @@ -1251,6 +1256,10 @@ pub struct NemoRelayNativeHostApiV4 { 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, @@ -1261,6 +1270,9 @@ pub struct NemoRelayNativeHostApiV4 { /// /// 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, @@ -1269,13 +1281,20 @@ pub struct NemoRelayNativeHostApiV4 { 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. + /// 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. @@ -1314,6 +1333,8 @@ pub struct NemoRelayNativeHostApiV4 { /// 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, diff --git a/crates/plugin/src/native_v2.rs b/crates/plugin/src/native_v2.rs index 283795087..03d66eec5 100644 --- a/crates/plugin/src/native_v2.rs +++ b/crates/plugin/src/native_v2.rs @@ -303,6 +303,8 @@ impl LlmContinuationV2 { 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)( @@ -385,6 +387,8 @@ impl LlmStreamContinuationV2 { 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)( @@ -418,6 +422,8 @@ impl LlmStreamContinuationV2 { 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)( @@ -463,6 +469,9 @@ impl Stream for LlmProviderStreamV2 { 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) }; @@ -621,6 +630,9 @@ impl PluginContext<'_> { }); 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(), diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 0df4516a1..b380e14d2 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -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"` for the established surface, or `"2"` for targeted LLM continuations | `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 666abaff8..a9700ca7a 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -191,7 +191,9 @@ 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: +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; diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index 1dead7a9e..fe65044f4 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -343,7 +343,7 @@ Recompile native plugins against the 0.7 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. From cbd994336e201cdfd16e68eda0160116c789fbe5 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 21:59:31 -0600 Subject: [PATCH 25/32] fix(core): bound and preserve targeted SSE frames Signed-off-by: Bryan Bednarski --- crates/core/src/codec/streaming.rs | 87 +++++++++++++++---- .../core/tests/unit/codec/streaming_tests.rs | 42 +++++++++ 2 files changed, 113 insertions(+), 16 deletions(-) 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/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 3a9856429..1df125e3e 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -31,6 +31,48 @@ 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 normalizes_crlf_terminator_split_across_pushes() { let mut decoder = SseEventDecoder::new(); From a83f142dfb105e5dc678e1c34db77d715e8acd02 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 22:03:51 -0600 Subject: [PATCH 26/32] fix(plugin): correct 32-bit host table assertions Signed-off-by: Bryan Bednarski --- crates/plugin/tests/typed_callbacks.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 1aa9bb4fa..f6b9df57a 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -507,18 +507,18 @@ fn native_abi_struct_sizes_are_self_describing() { ] ); assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 216); + 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 + 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::(), 256); + assert_eq!(size_of::(), 260); assert_eq!( host_api_v4_offsets(), - [0, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252] + [0, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256] ); assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 28); From 5918f8c99831ea97a2486d51044c10eefd6919c7 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Sun, 2 Aug 2026 22:04:02 -0600 Subject: [PATCH 27/32] docs(plugin): clarify native API v2 compatibility Signed-off-by: Bryan Bednarski --- crates/plugin/README.md | 5 +++++ .../dynamic-plugins/native-dynamic/about.mdx | 14 +++++++++++--- docs/reference/migration-guides.mdx | 4 +++- 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/crates/plugin/README.md b/crates/plugin/README.md index d8fa749d6..c666836df 100644 --- a/crates/plugin/README.md +++ b/crates/plugin/README.md @@ -154,6 +154,11 @@ 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`. diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index a9700ca7a..c62476f65 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -126,9 +126,12 @@ separate from the internal host-table `abi_version` field: | `"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 retains native API v1 unchanged. 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. +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 @@ -238,6 +241,11 @@ 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 diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index fe65044f4..ab170287c 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -351,7 +351,9 @@ separate from the synchronous `nemo-relay-ffi` middleware registration API. 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. +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 From ec8f9cd88aef722d3b273cac779acc65d6abe6da Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 3 Aug 2026 09:32:17 -0600 Subject: [PATCH 28/32] test(plugin): cover native stream error cleanup Signed-off-by: Bryan Bednarski --- .../core/tests/unit/codec/streaming_tests.rs | 38 ++++++++++++++ crates/core/tests/unit/native_plugin_tests.rs | 51 +++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/crates/core/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 1df125e3e..ad9b52b4d 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -73,6 +73,31 @@ fn preserves_completed_events_before_an_oversized_partial_frame() { 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(); @@ -111,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/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 6430e92de..659349a5a 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -3180,6 +3180,57 @@ fn test_native_output_stream( ) } +#[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, From a5acddf4f59ce37d9adca3c85e11de67c53df9de Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 3 Aug 2026 10:18:58 -0600 Subject: [PATCH 29/32] test(plugin): reconcile native lifecycle coverage Signed-off-by: Bryan Bednarski --- crates/core/tests/unit/native_plugin_tests.rs | 73 +++++++++++++------ 1 file changed, 51 insertions(+), 22 deletions(-) diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 659349a5a..22af0077e 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -763,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), }), } } @@ -872,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!( @@ -891,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, }), @@ -911,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, }); @@ -935,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)), @@ -944,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, }); @@ -1002,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, }) @@ -1135,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 { @@ -1145,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, }); @@ -1171,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, }); @@ -1193,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, }); @@ -1214,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, }); @@ -1289,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, }); @@ -4473,7 +4501,15 @@ fn native_legacy_stream_next_reports_validation_item_and_setup_failures() { .await .expect("allocation failure should terminate the legacy stream"); }); - assert!(callback.done.load(Ordering::Acquire)); + 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); @@ -7468,7 +7504,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(), @@ -7925,7 +7962,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 { @@ -7974,7 +8012,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(); @@ -8710,17 +8749,6 @@ unsafe extern "C" fn llm_execution_error_with_output( NemoRelayStatus::InvalidArg } -unsafe extern "C" fn llm_stream_execution_error( - _user_data: *mut c_void, - _name: *const NemoRelayNativeString, - _request_json: *const NemoRelayNativeString, - _next_fn: NemoRelayNativeLlmStreamNextFn, - _next_ctx: *mut c_void, - _out_stream: *mut NemoRelayNativeLlmStreamV1, -) -> NemoRelayStatus { - NemoRelayStatus::InvalidArg -} - #[test] fn native_callback_helpers_cover_success_error_and_invalid_output() { assert_eq!( @@ -8878,7 +8906,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(), From 6dac757f2ff68cb4294867c2578fdeb80830fd9f Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 3 Aug 2026 10:53:50 -0600 Subject: [PATCH 30/32] test(plugin): gate Unix-only native callbacks Signed-off-by: Bryan Bednarski --- crates/core/tests/unit/native_plugin_tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 22af0077e..9f3574bca 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -8718,6 +8718,7 @@ 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, @@ -8728,6 +8729,7 @@ unsafe extern "C" fn llm_response_error_with_output( NemoRelayStatus::InvalidArg } +#[cfg(unix)] unsafe extern "C" fn llm_response_none( _user_data: *mut c_void, _response_json: *const NemoRelayNativeString, @@ -8737,6 +8739,7 @@ unsafe extern "C" fn llm_response_none( NemoRelayStatus::Ok } +#[cfg(unix)] unsafe extern "C" fn llm_execution_error_with_output( _user_data: *mut c_void, _name: *const NemoRelayNativeString, From 6598e0c06dcb93e861b6812cfc9cea9f28fbd4a7 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 3 Aug 2026 11:19:16 -0600 Subject: [PATCH 31/32] test(plugin): split native layout assertions Signed-off-by: Bryan Bednarski --- crates/plugin/tests/typed_callbacks.rs | 133 +++++++++++++------------ 1 file changed, 67 insertions(+), 66 deletions(-) diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index f6b9df57a..9714bde4b 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -460,73 +460,74 @@ fn native_abi_struct_sizes_are_self_describing() { 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::(), 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]); - } + 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::(), 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]); - } +#[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] { From 3a8f8f0745fc6545a9162bd585da2273fa052785 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 3 Aug 2026 11:22:42 -0600 Subject: [PATCH 32/32] test(plugin): split native failure assertions Signed-off-by: Bryan Bednarski --- crates/core/tests/unit/native_plugin_tests.rs | 134 ++++++++++-------- 1 file changed, 75 insertions(+), 59 deletions(-) diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 9f3574bca..ced07d06a 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -3602,56 +3602,7 @@ fn native_api_v2_split_callbacks_settle_when_native_string_allocation_fails() { let dispatch = test_v2_dispatch_json(); let live_with_dispatch = native_string_live_allocations(); - 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) }; - } + assert_unary_allocation_failures_settle(&runtime, dispatch, live_with_dispatch); let open_next = Arc::new(NativeAsyncNext::new( NativeAsyncNextInner::LlmStream(Arc::new(|_| { @@ -3749,6 +3700,63 @@ fn native_api_v2_split_callbacks_settle_when_native_string_allocation_fails() { 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 = [ @@ -8419,13 +8427,25 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { 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"); - output = 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, ) @@ -8442,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, @@ -8460,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, ) @@ -8469,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(