diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index 3712b4e6a..7ae5c12ea 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -452,7 +452,7 @@ pub fn create_scope_stack_from_propagation( /// # Examples /// /// ```no_run -/// # async fn example() -> nemo_relay::Result<()> { +/// # async fn example() -> nemo_relay::error::Result<()> { /// use nemo_relay::api::runtime::{TASK_SCOPE_STACK, fork_scope_stack}; /// /// let stack = fork_scope_stack()?; diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index d6fde6414..3b0124b5b 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -2218,80 +2218,11 @@ unsafe extern "C" fn native_async_next_invoke_stream( 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 { - match item { - Ok(chunk) => { - if let Some(chunk) = native_string_from_json(&chunk) { - let keep_going = unsafe { - cb( - user_data as *mut c_void, - chunk, - ptr::null(), - false, - ) - }; - unsafe { - native_string_free(chunk); - } - if !keep_going { - callback_guard.finish(); - return; - } - } else { - break; - } - } - Err(error) => { - if let Some(message) = - native_string_from_str(&error.to_string()) - { - unsafe { - let _ = cb( - user_data as *mut c_void, - ptr::null(), - message, - false, - ); - native_string_free(message); - } - callback_guard.finish(); - } - return; - } - } - } - unsafe { - let _ = - cb(user_data as *mut c_void, ptr::null(), ptr::null(), true); - } - callback_guard.finish(); - } - Err(error) => { - if let Some(message) = native_string_from_str(&error.to_string()) { - unsafe { - let _ = - cb(user_data as *mut c_void, ptr::null(), message, false); - native_string_free(message); - } - callback_guard.finish(); - } - } - } - }) - .catch_unwind() - .await; - if let Err(payload) = result { - callback_guard.fail(&format!( - "native async stream continuation panicked: {}", - panic_payload_message(payload.as_ref()) - )); - } - }) + .run(deliver_native_async_next_stream( + next_fn, + request, + callback_guard, + )) .await; output_stream_for_cleanup .downstream_aborts @@ -2305,6 +2236,79 @@ unsafe extern "C" fn native_async_next_invoke_stream( NemoRelayStatus::Ok } +async fn deliver_native_async_next_stream( + next_fn: LlmStreamExecutionNextFn, + request: LlmRequest, + mut callback_guard: NativeAsyncStreamCallbackGuard, +) { + let result = AssertUnwindSafe(async { + match next_fn(request).await { + Ok(stream) => forward_native_async_next_stream(stream, &mut callback_guard).await, + Err(error) => callback_guard.fail(&error.to_string()), + } + }) + .catch_unwind() + .await; + if let Err(payload) = result { + callback_guard.fail(&format!( + "native async stream continuation panicked: {}", + panic_payload_message(payload.as_ref()) + )); + } +} + +async fn forward_native_async_next_stream( + stream: LlmJsonStream, + callback_guard: &mut NativeAsyncStreamCallbackGuard, +) { + forward_native_async_next_stream_with(stream, callback_guard, native_string_from_json).await; +} + +async fn forward_native_async_next_stream_with( + mut stream: LlmJsonStream, + callback_guard: &mut NativeAsyncStreamCallbackGuard, + to_native_string: impl Fn(&Json) -> Option<*mut NemoRelayNativeString>, +) { + while let Some(item) = stream.next().await { + match item { + Ok(chunk) => { + let Some(chunk) = to_native_string(&chunk) else { + callback_guard.fail( + "failed to serialize or allocate native async stream continuation chunk", + ); + return; + }; + let keep_going = unsafe { + (callback_guard.cb)( + callback_guard.user_data as *mut c_void, + chunk, + ptr::null(), + false, + ) + }; + unsafe { native_string_free(chunk) }; + if !keep_going { + callback_guard.finish(); + return; + } + } + Err(error) => { + callback_guard.fail(&error.to_string()); + return; + } + } + } + unsafe { + let _ = (callback_guard.cb)( + callback_guard.user_data as *mut c_void, + ptr::null(), + ptr::null(), + true, + ); + } + callback_guard.finish(); +} + fn 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 e180909bd..f0978fc6d 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -15,6 +15,8 @@ use nemo_relay_plugin::{ NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamNextFn, NemoRelayNativeToolNextFn, }; +#[cfg(unix)] +use nemo_relay_plugin::{NemoRelayNativePluginRegisterFn, NemoRelayNativePluginValidateFn}; use serde_json::json; use crate::api::optimization::{ @@ -185,592 +187,1319 @@ unsafe extern "C" fn stop_after_first_native_stream_item( false } -struct InvokeNativeNextThenReturnState { - callback_state: u32, - invoke_status: AtomicUsize, - started: Mutex>, -} +#[test] +fn native_async_entrypoints_reject_null_handles() { + unsafe { + assert_eq!( + native_async_completion_resolve_json(ptr::null(), ptr::null()), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_async_completion_reject(ptr::null(), ptr::null()), + NemoRelayStatus::NullPointer + ); + assert!(native_async_completion_is_cancelled(ptr::null())); + native_async_completion_release(ptr::null()); + native_async_next_release(ptr::null()); -unsafe extern "C" fn invoke_native_next_then_return_state( - user_data: *mut c_void, - invocation_json: *const NemoRelayNativeString, - next: *const NemoRelayNativeAsyncNext, - completion: *const NemoRelayNativeAsyncCompletion, -) -> u32 { - let state = unsafe { &*user_data.cast::() }; - let status = unsafe { native_async_next_invoke(next, invocation_json, completion) }; - state - .invoke_status - .store(status as usize, Ordering::Release); - if status == NemoRelayStatus::Ok { - let _ = state - .started - .lock() - .unwrap_or_else(|error| error.into_inner()) - .recv_timeout(Duration::from_secs(1)); - } - unsafe { native_async_next_release(next) }; - state.callback_state -} + assert_eq!( + native_async_stream_push_json(ptr::null(), ptr::null()), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_async_stream_finish(ptr::null()), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_async_stream_reject(ptr::null(), ptr::null()), + NemoRelayStatus::NullPointer + ); + assert!(native_async_stream_is_cancelled(ptr::null())); + native_async_stream_release(ptr::null()); -unsafe extern "C" fn invoke_native_stream_next_then_return_state( - user_data: *mut c_void, - invocation_json: *const NemoRelayNativeString, - next: *const NemoRelayNativeAsyncNext, - stream: *const NemoRelayNativeAsyncStream, -) -> u32 { - let state = unsafe { &*user_data.cast::() }; - let request = read_native_string(invocation_json) - .ok() - .and_then(|invocation| serde_json::from_str::(&invocation).ok()) - .and_then(|invocation| invocation.get("request").cloned()) - .and_then(|request| native_string_from_json(&request)); - let status = if let Some(request) = request { - let status = unsafe { + assert_eq!( + native_async_next_invoke(ptr::null(), ptr::null(), ptr::null()), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_async_next_invoke_result( + ptr::null(), + ptr::null(), + complete_native_next_result, + ptr::null_mut(), + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( native_async_next_invoke_stream( - next, - request, - stream, + ptr::null(), + ptr::null(), + ptr::null(), accept_native_stream_item, ptr::null_mut(), - ) - }; - unsafe { native_string_free(request) }; - status - } else { - NemoRelayStatus::InvalidJson - }; - state - .invoke_status - .store(status as usize, Ordering::Release); - if status == NemoRelayStatus::Ok { - let _ = state - .started - .lock() - .unwrap_or_else(|error| error.into_inner()) - .recv_timeout(Duration::from_secs(1)); - } - unsafe { - native_async_next_release(next); - native_async_stream_release(stream); + ), + NemoRelayStatus::NullPointer + ); } - state.callback_state } -unsafe extern "C" fn invoke_detached_next_and_finish_replacement_stream( - user_data: *mut c_void, - invocation_json: *const NemoRelayNativeString, - next: *const NemoRelayNativeAsyncNext, - stream: *const NemoRelayNativeAsyncStream, -) -> u32 { - let invoke_status = unsafe { &*user_data.cast::() }; - let request = read_native_string(invocation_json) - .ok() - .and_then(|invocation| serde_json::from_str::(&invocation).ok()) - .and_then(|invocation| invocation.get("request").cloned()) - .and_then(|request| native_string_from_json(&request)); - let status = if let Some(request) = request { - let status = unsafe { - native_async_next_invoke_stream( - next, - request, - stream, - accept_native_stream_item, - ptr::null_mut(), - ) - }; - unsafe { native_string_free(request) }; - status - } else { - NemoRelayStatus::InvalidJson - }; - invoke_status.store(status as usize, Ordering::Release); - let replacement = native_string_from_json(&json!({"source": "replacement"})).unwrap(); - unsafe { - let _ = native_async_stream_push_json(stream, replacement); - native_string_free(replacement); - let _ = native_async_stream_finish(stream); - native_async_next_release(next); - native_async_stream_release(stream); - } - NemoRelayNativeAsyncCallbackState::Complete as u32 +#[cfg(unix)] +unsafe extern "C" fn native_test_validate_invalid_json( + _user_data: *mut c_void, + _config: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out = native_string("not-json") }; + NemoRelayStatus::Ok } -struct FailingNativeCodec; +#[cfg(unix)] +unsafe extern "C" fn native_test_validate_error( + _user_data: *mut c_void, + _config: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out = native_string("unused") }; + set_native_last_error("validation callback failed"); + NemoRelayStatus::InvalidArg +} -impl LlmCodec for FailingNativeCodec { - fn decode(&self, _request: &LlmRequest) -> FlowResult { - Err(FlowError::Internal("request decode rejected".into())) - } +#[cfg(unix)] +unsafe extern "C" fn native_test_validate_empty( + _user_data: *mut c_void, + _config: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} - fn encode( - &self, - _annotated: &AnnotatedLlmRequest, - _original: &LlmRequest, - ) -> FlowResult { - Err(FlowError::Internal("request encode rejected".into())) - } +unsafe extern "C" fn native_test_register_ok( + _user_data: *mut c_void, + _config: *const NemoRelayNativeString, + _ctx: *mut NemoRelayNativePluginContext, +) -> NemoRelayStatus { + NemoRelayStatus::Ok } -impl LlmResponseCodec for FailingNativeCodec { - fn decode_response(&self, _response: &Json) -> FlowResult { - Err(FlowError::Internal("response decode rejected".into())) +#[cfg(unix)] +unsafe extern "C" fn native_test_register_error( + _user_data: *mut c_void, + _config: *const NemoRelayNativeString, + _ctx: *mut NemoRelayNativePluginContext, +) -> NemoRelayStatus { + set_native_last_error("registration callback failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +fn native_test_adapter( + validate: Option, + register: Option, +) -> NativePluginAdapter { + let plugin = NemoRelayNativePluginV1 { + validate, + register, + ..Default::default() + }; + NativePluginAdapter { + plugin_kind: "test.native.adapter".into(), + allows_multiple_components: false, + instance: Arc::new(NativePluginInstance { + plugin_kind: "test.native.adapter".into(), + relay_compat: "^0.7".into(), + allows_multiple_components: false, + plugin: Mutex::new(plugin), + _library: libloading::os::unix::Library::this().into(), + }), } } -struct PanickingNativeCodec; +#[cfg(unix)] +#[tokio::test] +async fn native_plugin_adapter_covers_validation_and_registration_results() { + let no_validate = native_test_adapter(None, Some(native_test_register_ok)); + assert_eq!(no_validate.plugin_kind(), "test.native.adapter"); + assert!(!no_validate.allows_multiple_components()); + assert!(no_validate.validate(&Map::new()).is_empty()); + + let empty = native_test_adapter( + Some(native_test_validate_empty), + Some(native_test_register_ok), + ); + assert!(empty.validate(&Map::new()).is_empty()); + + let invalid = native_test_adapter( + Some(native_test_validate_invalid_json), + Some(native_test_register_ok), + ); + let diagnostics = invalid.validate(&Map::new()); + assert_eq!(diagnostics.len(), 1); + assert!(diagnostics[0].message.contains("invalid diagnostics JSON")); + + let failing = native_test_adapter( + Some(native_test_validate_error), + Some(native_test_register_error), + ); + let diagnostics = failing.validate(&Map::new()); + assert_eq!(diagnostics.len(), 1); + assert_eq!(diagnostics[0].message, "validation callback failed"); + + let mut context = PluginRegistrationContext::new(); + assert!(empty.register(&Map::new(), &mut context).await.is_ok()); + let error = failing + .register(&Map::new(), &mut context) + .await + .expect_err("registration callback should fail"); + assert!(error.to_string().contains("registration callback failed")); + + let missing = native_test_adapter(Some(native_test_validate_empty), None); + let error = missing + .register(&Map::new(), &mut context) + .await + .expect_err("missing registration callback should fail"); + assert!( + error + .to_string() + .contains("did not return a register callback") + ); +} -impl LlmCodec for PanickingNativeCodec { - fn decode(&self, _request: &LlmRequest) -> FlowResult { - panic!("request decode panic") - } +#[test] +fn native_loader_helpers_cover_compatibility_descriptor_and_digest_edges() { + assert_native_compatibility_edges(); + assert_native_descriptor_edges(); + assert_native_digest_edges(); + assert_native_host_api_versions(); +} - fn encode( - &self, - _annotated: &AnnotatedLlmRequest, - _original: &LlmRequest, - ) -> FlowResult { - panic!("request encode panic") - } +fn assert_native_compatibility_edges() { + assert!(validate_relay_compatibility(None).is_err()); + assert!(validate_relay_compatibility(Some(" ")).is_err()); + assert!(validate_relay_compatibility(Some("not a requirement")).is_err()); + assert!(validate_relay_compatibility(Some(">=999.0.0")).is_err()); + let host_requirement = format!("={}", env!("CARGO_PKG_VERSION")); + assert!(validate_relay_compatibility(Some(&host_requirement)).is_ok()); } -impl LlmResponseCodec for PanickingNativeCodec { - fn decode_response(&self, _response: &Json) -> FlowResult { - panic!("response decode panic") - } +fn assert_native_descriptor_edges() { + let mut descriptor = NemoRelayNativePluginV1 { + struct_size: 0, + ..Default::default() + }; + assert!(validate_plugin_descriptor("test", &descriptor).is_err()); + descriptor.struct_size = std::mem::size_of::(); + assert!(validate_plugin_descriptor("test", &descriptor).is_err()); + descriptor.plugin_kind = native_string("test"); + assert!(validate_plugin_descriptor("test", &descriptor).is_err()); + descriptor.register = Some(native_test_register_ok); + assert!(validate_plugin_descriptor("test", &descriptor).is_ok()); + drop_native_plugin_descriptor(&mut descriptor); } -#[test] -fn native_string_and_json_helpers_cover_abi_boundaries() { - clear_native_last_error(); +fn assert_native_digest_edges() { + let temp = tempfile::tempdir().unwrap(); + let manifest = temp.path().join("plugin.toml"); + let library = temp.path().join("plugin.bin"); + std::fs::write(&library, b"native plugin bytes").unwrap(); assert_eq!( - unsafe { native_string_new(ptr::null(), 0, ptr::null_mut()) }, - NemoRelayStatus::NullPointer + resolve_manifest_relative_path(&manifest, "plugin.bin"), + library ); - assert_last_error_contains("out string pointer is null"); + assert_eq!( + resolve_manifest_relative_path(&manifest, temp.path().to_str().unwrap()), + temp.path() + ); + assert_eq!(hex_digest([0x00, 0xab, 0xff]), "00abff"); + let digest = hex_digest(Sha256::digest(b"native plugin bytes")); + assert!(verify_sha256(&library, &format!("sha256:{}", digest.to_uppercase())).is_ok()); + assert!(verify_sha256(&library, "sha256:00").is_err()); + assert!(verify_sha256(&temp.path().join("missing"), "00").is_err()); +} - let mut out = ptr::null_mut(); +fn assert_native_host_api_versions() { + let current = native_host_api(); + let legacy = native_host_api_legacy(); + assert!(!current.is_null()); + assert!(!legacy.is_null()); assert_eq!( - unsafe { native_string_new(ptr::null(), 1, &mut out) }, - NemoRelayStatus::NullPointer + unsafe { (*legacy).abi_version }, + NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY ); - assert!(out.is_null()); - assert_last_error_contains("string data pointer is null"); +} - let invalid_utf8 = [0xff]; +#[tokio::test] +async fn native_async_wait_and_rejection_cover_dropped_and_aborted_continuations() { + let (sender, receiver) = tokio::sync::oneshot::channel::>(); + drop(sender); + let mut wait = NativeAsyncWait { + completion: Arc::new(NativeAsyncCompletion { + sender: Mutex::new(None), + cancelled: AtomicBool::new(false), + next_invoked: AtomicBool::new(false), + next_abort: Mutex::new(None), + before_settlement_lock: None, + _callback_user_data: None, + }), + receiver, + completed: false, + }; + let error = wait + .receive() + .await + .expect_err("a dropped callback must fail the wait"); + assert!(error.to_string().contains("dropped without settling")); + + let task = tokio::spawn(std::future::pending::<()>()); + let abort = task.abort_handle(); + 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(true), + next_abort: Mutex::new(Some(abort)), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let invalid_message = Box::into_raw(Box::new(NativeHostString(vec![0xff]))).cast(); assert_eq!( - unsafe { native_string_new(invalid_utf8.as_ptr(), invalid_utf8.len(), &mut out) }, - NemoRelayStatus::InvalidUtf8 + unsafe { native_async_completion_reject(completion_ref, invalid_message) }, + NemoRelayStatus::InvalidArg ); - assert!(out.is_null()); - assert_last_error_contains("not valid UTF-8"); + unsafe { native_string_free(invalid_message) }; - let text = native_string("hello"); - assert_eq!(unsafe { native_string_len(text) }, 5); + let message = native_string("continuation rejected"); assert_eq!( - unsafe { std::slice::from_raw_parts(native_string_data(text), 5) }, - b"hello" + unsafe { native_async_completion_reject(completion_ref, message) }, + NemoRelayStatus::Ok ); - assert!(unsafe { native_string_data(ptr::null()) }.is_null()); - assert_eq!(unsafe { native_string_len(ptr::null()) }, 0); - assert_eq!(take_native_string(text).unwrap(), "hello"); - unsafe { native_string_free(ptr::null_mut()) }; - - let empty = native_string(""); - assert_eq!(read_native_string(empty).unwrap(), ""); - unsafe { native_string_free(empty) }; - assert_eq!(read_native_string(ptr::null()).unwrap(), ""); + unsafe { native_string_free(message) }; + let error = receiver.await.unwrap().unwrap_err(); + assert!(error.to_string().contains("continuation rejected")); + assert!(task.await.unwrap_err().is_cancelled()); + unsafe { native_async_completion_release(completion_ref) }; +} - let bad = Box::into_raw(Box::new(NativeHostString(vec![0xff]))) as *mut NemoRelayNativeString; - assert!(read_native_string(bad).is_err()); - assert_eq!( - optional_json_from_native_string(bad, "bad json"), - Err(NemoRelayStatus::InvalidUtf8) - ); - unsafe { native_last_error_set(bad) }; - assert_last_error_contains("not valid UTF-8"); - unsafe { native_string_free(bad) }; +#[test] +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)), + 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 callback_state = NativeStreamCallbackState::default(); + let user_data: *mut c_void = ptr::from_ref(&callback_state).cast_mut().cast(); - let message = native_string("explicit native error"); - unsafe { native_last_error_set(message) }; - assert_eq!( - native_last_error_message().as_deref(), - Some("explicit native error") - ); - unsafe { native_string_free(message) }; - unsafe { native_last_error_clear() }; - assert!(native_last_error_message().is_none()); + let mut inactive = NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: user_data as usize, + stream: Arc::clone(&stream), + _library_guard: None, + active: false, + }; + inactive.fail("ignored"); - set_native_last_error("specific fallback"); + stream.cancelled.store(true, Ordering::Release); + let mut cancelled = NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: user_data as usize, + stream: Arc::clone(&stream), + _library_guard: None, + active: true, + }; + cancelled.fail("cancellation owns settlement"); + drop(cancelled); + + stream.cancelled.store(false, Ordering::Release); + stream.settled.store(true, Ordering::Release); + drop(NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: user_data as usize, + stream: Arc::clone(&stream), + _library_guard: None, + active: true, + }); + + stream.settled.store(false, Ordering::Release); + drop(NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: user_data as usize, + stream, + _library_guard: None, + active: true, + }); + assert_eq!(callback_state.callbacks.load(Ordering::Acquire), 3); + assert!(callback_state.done.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn native_async_stream_forwarding_reports_conversion_and_stream_errors() { + let make_stream_state = || { + let (sender, _receiver) = tokio::sync::mpsc::channel(1); + 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 conversion_state = NativeStreamCallbackState::default(); + let mut conversion_guard = NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: ptr::from_ref(&conversion_state) as usize, + stream: make_stream_state(), + _library_guard: None, + active: true, + }; + forward_native_async_next_stream_with( + LlmJsonStream::new(tokio_stream::iter([Ok(json!({"chunk": true}))])), + &mut conversion_guard, + |_| None, + ) + .await; assert!( - json_from_native_string(ptr::null_mut(), "generic fallback") - .unwrap_err() - .to_string() - .contains("specific fallback") + conversion_state + .error + .lock() + .unwrap() + .as_deref() + .is_some_and(|error| error.contains("failed to serialize or allocate")) ); - clear_native_last_error(); + assert!(!conversion_state.done.load(Ordering::Acquire)); + + let stream_error_state = NativeStreamCallbackState::default(); + let mut stream_error_guard = NativeAsyncStreamCallbackGuard { + cb: record_native_stream_result, + user_data: ptr::from_ref(&stream_error_state) as usize, + stream: make_stream_state(), + _library_guard: None, + active: true, + }; + forward_native_async_next_stream( + LlmJsonStream::new(tokio_stream::iter([Err(FlowError::Internal( + "provider stream failed".into(), + ))])), + &mut stream_error_guard, + ) + .await; assert!( - json_from_native_string(ptr::null_mut(), "generic fallback") - .unwrap_err() - .to_string() - .contains("generic fallback") + stream_error_state + .error + .lock() + .unwrap() + .as_deref() + .is_some_and(|error| error.contains("provider stream failed")) ); +} - let invalid_json = native_string("{"); - assert!( - take_json_from_native_string(invalid_json, "unused") - .unwrap_err() - .to_string() - .contains("invalid JSON") +#[tokio::test] +async fn native_async_result_entrypoint_covers_llm_and_stream_continuations() { + let llm_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Llm(Arc::new(|request| { + Box::pin(async move { Ok(request.content) }) + })), + tokio::runtime::Handle::current(), + None, + )); + let llm_ref = Arc::into_raw(llm_next) as *const NemoRelayNativeAsyncNext; + let invalid = native_string("{}"); + assert_eq!( + unsafe { + native_async_next_invoke_result( + llm_ref, + invalid, + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidJson ); + unsafe { native_string_free(invalid) }; + + let request = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"message": "hello"}), + }) + .unwrap(), + ) + .unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::>(); assert_eq!( - optional_json_from_native_string(ptr::null(), "optional"), - Ok(None) + unsafe { + native_async_next_invoke_result( + llm_ref, + request, + complete_native_next_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, + NemoRelayStatus::Ok ); - let valid_json = native_string(r#"{"value":1}"#); assert_eq!( - optional_json_from_native_string(valid_json, "optional").unwrap(), - Some(json!({"value": 1})) + receiver.await.unwrap().unwrap(), + json!({"message": "hello"}) ); - unsafe { native_string_free(valid_json) }; - let invalid_json = native_string("not-json"); + unsafe { + native_string_free(request); + native_async_next_release(llm_ref); + } + + let stream_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })), + tokio::runtime::Handle::current(), + None, + )); + let stream_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; + let invocation = native_string("null"); assert_eq!( - optional_json_from_native_string(invalid_json, "optional"), - Err(NemoRelayStatus::InvalidJson) + unsafe { + native_async_next_invoke_result( + stream_ref, + invocation, + complete_native_next_result, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg ); - assert_last_error_contains("optional is not valid JSON"); - unsafe { native_string_free(invalid_json) }; + unsafe { + native_string_free(invocation); + native_async_next_release(stream_ref); + } +} +#[test] +fn native_async_stream_entrypoints_cover_closed_full_and_settled_channels() { + let chunk = native_string("null"); + + let no_sender = Arc::new(NativeAsyncStream { + sender: Mutex::new(None), + 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 no_sender_ref = Arc::into_raw(no_sender) as *const NemoRelayNativeAsyncStream; assert_eq!( - parse_json_arg(ptr::null(), "null JSON").unwrap_err(), - NemoRelayStatus::InvalidJson + unsafe { native_async_stream_push_json(no_sender_ref, chunk) }, + NemoRelayStatus::InvalidArg ); - let request = LlmRequest { - headers: Map::new(), - content: json!({"model": "test"}), - }; - let request_json = native_string_from_json(&serde_json::to_value(&request).unwrap()).unwrap(); assert_eq!( - parse_llm_request_arg(request_json, "request").unwrap(), - request + unsafe { native_async_stream_finish(no_sender_ref) }, + NemoRelayStatus::InvalidArg ); - unsafe { native_string_free(request_json) }; - let wrong_shape = native_string(r#"{"headers":[]}"#); assert_eq!( - parse_llm_request_arg(wrong_shape, "request").unwrap_err(), - NemoRelayStatus::InvalidJson + unsafe { native_async_stream_reject(no_sender_ref, chunk) }, + NemoRelayStatus::InvalidArg ); - assert_last_error_contains("was not an LLM request"); - unsafe { native_string_free(wrong_shape) }; + unsafe { native_async_stream_release(no_sender_ref) }; + let (full_sender, _full_receiver) = tokio::sync::mpsc::channel::>(1); + full_sender.try_send(Ok(Json::Null)).unwrap(); + let full = Arc::new(NativeAsyncStream { + sender: Mutex::new(Some(full_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 full_ref = Arc::into_raw(full) as *const NemoRelayNativeAsyncStream; assert_eq!( - write_native_json(&json!({"ok": true}), ptr::null_mut()), - NemoRelayStatus::NullPointer + unsafe { native_async_stream_push_json(full_ref, chunk) }, + NemoRelayStatus::Internal ); - let mut json_out = ptr::null_mut(); assert_eq!( - write_native_json(&json!({"ok": true}), &mut json_out), - NemoRelayStatus::Ok + unsafe { native_async_stream_reject(full_ref, chunk) }, + NemoRelayStatus::Internal ); + unsafe { native_async_stream_release(full_ref) }; + + let (closed_sender, closed_receiver) = tokio::sync::mpsc::channel::>(1); + drop(closed_receiver); + let closed = Arc::new(NativeAsyncStream { + sender: Mutex::new(Some(closed_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 closed_ref = Arc::into_raw(closed) as *const NemoRelayNativeAsyncStream; assert_eq!( - take_json_from_native_string(json_out, "unused").unwrap(), - json!({"ok": true}) + unsafe { native_async_stream_push_json(closed_ref, chunk) }, + NemoRelayStatus::InvalidArg ); + assert_eq!( + unsafe { native_async_stream_reject(closed_ref, chunk) }, + NemoRelayStatus::InvalidArg + ); + unsafe { native_async_stream_release(closed_ref) }; - let host_api = unsafe { &*native_host_api() }; - assert_eq!(host_api.abi_version, NEMO_RELAY_NATIVE_ABI_VERSION); + let (settled_sender, _settled_receiver) = tokio::sync::mpsc::channel(1); + let settled = Arc::new(NativeAsyncStream { + sender: Mutex::new(Some(settled_sender)), + cancelled: AtomicBool::new(false), + settled: AtomicBool::new(true), + downstream_aborts: Mutex::new(HashMap::new()), + settlement: Mutex::new(()), + before_settlement_lock: None, + _callback_user_data: None, + }); + let settled_ref = Arc::into_raw(settled) as *const NemoRelayNativeAsyncStream; assert_eq!( - host_api.struct_size, - std::mem::size_of::() + unsafe { native_async_stream_push_json(settled_ref, chunk) }, + NemoRelayStatus::InvalidArg + ); + assert_eq!( + unsafe { native_async_stream_finish(settled_ref) }, + NemoRelayStatus::InvalidArg + ); + assert_eq!( + unsafe { native_async_stream_reject(settled_ref, chunk) }, + NemoRelayStatus::InvalidArg ); + unsafe { + native_async_stream_release(settled_ref); + native_string_free(chunk); + } } -#[test] -fn native_async_next_abi_runs_tool_llm_and_stream_continuations() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let cases: Vec<(NativeAsyncNextInner, Json, Json)> = vec![ - ( - NativeAsyncNextInner::Tool(Arc::new(|value| Box::pin(async move { Ok(value) }))), - json!({"tool": true}), - json!({"result": {"tool": true}, "pending_marks": []}), - ), - ( - NativeAsyncNextInner::Llm(Arc::new(|request| { - Box::pin(async move { Ok(request.content) }) - })), - serde_json::to_value(LlmRequest { - headers: Map::new(), - content: json!({"llm": true}), - }) - .unwrap(), - json!({"llm": true}), - ), - ]; - - for (inner, invocation, expected) in cases { - let next = Arc::new(NativeAsyncNext::new(inner, runtime.handle().clone(), None)); +#[tokio::test] +async fn native_async_result_entrypoint_reports_provider_errors_and_panics() { + for next_fn in [ + Arc::new(|_value| { + Box::pin(async { Err(FlowError::Internal("provider failed".into())) }) + as Pin> + Send>> + }) as ToolExecutionNextFn, + Arc::new(|_value| { + Box::pin(async { + panic!("provider panicked"); + #[allow(unreachable_code)] + Ok(Json::Null) + }) as Pin> + Send>> + }) as ToolExecutionNextFn, + ] { + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(next_fn), + tokio::runtime::Handle::current(), + None, + )); let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - 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), - before_settlement_lock: None, - _callback_user_data: None, - }); - let completion_ref = - Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; - let invocation = native_string_from_json(&invocation).unwrap(); + let invocation = native_string("null"); + let (sender, receiver) = + tokio::sync::oneshot::channel::>(); assert_eq!( - unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, + unsafe { + native_async_next_invoke_result( + next_ref, + invocation, + complete_native_next_result, + Box::into_raw(Box::new(sender)).cast(), + ) + }, NemoRelayStatus::Ok ); - assert_eq!(runtime.block_on(receiver).unwrap().unwrap(), expected); + let error = receiver.await.unwrap().unwrap_err(); + assert!(error.contains("provider")); unsafe { native_string_free(invocation); native_async_next_release(next_ref); - native_async_completion_release(completion_ref); } } +} - let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::LlmStream(Arc::new(|_request| { - Box::pin(async { - Ok(LlmJsonStream::new(tokio_stream::iter(vec![ - Ok(json!({"chunk": 1})), - Ok(json!({"chunk": 2})), - ]))) - }) - })), - runtime.handle().clone(), - None, - )); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let (sender, _receiver) = tokio::sync::oneshot::channel(); - let completion = Arc::new(NativeAsyncCompletion { +#[tokio::test] +async fn native_async_stream_next_entrypoint_validates_handle_kind_and_request() { + let (sender, _receiver) = tokio::sync::mpsc::channel(1); + let stream = Arc::new(NativeAsyncStream { sender: Mutex::new(Some(sender)), cancelled: AtomicBool::new(false), - next_invoked: AtomicBool::new(false), - next_abort: Mutex::new(None), + settled: AtomicBool::new(false), + downstream_aborts: Mutex::new(HashMap::new()), + settlement: Mutex::new(()), before_settlement_lock: None, _callback_user_data: None, }); - let completion_ref = - Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; - let invocation = native_string_from_json( - &serde_json::to_value(LlmRequest { - headers: Map::new(), - content: json!({"stream": true}), - }) - .unwrap(), - ) - .unwrap(); - assert_eq!( - unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, - NemoRelayStatus::InvalidArg - ); - assert_last_error_contains("async_next_invoke_stream"); - unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); - native_async_completion_release(completion_ref); - } -} + let stream_ref = Arc::into_raw(stream) as *const NemoRelayNativeAsyncStream; + let invocation = native_string("null"); -#[test] -fn native_async_next_reports_a_revoked_continuation_without_calling_the_provider() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let provider_calls = Arc::new(AtomicUsize::new(0)); - let (lease, guard) = MiddlewareContinuationLease::capture(); - let next = Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::Tool({ - let provider_calls = provider_calls.clone(); - Arc::new(move |value| { - let provider_calls = provider_calls.clone(); - let invocation = lease.begin(); - Box::pin(async move { - invocation? - .invoke(|| async move { - provider_calls.fetch_add(1, Ordering::SeqCst); - Ok(value) - }) - .await - }) - }) - }), - runtime.handle().clone(), + let tool_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(Arc::new(|value| Box::pin(async move { Ok(value) }))), + tokio::runtime::Handle::current(), None, )); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - 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), - before_settlement_lock: None, - _callback_user_data: None, - }); - let completion_ref = - Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; - let invocation = native_string_from_json(&json!({"tool": true})).unwrap(); - - drop(guard); + let tool_ref = Arc::into_raw(tool_next) as *const NemoRelayNativeAsyncNext; assert_eq!( - unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, - NemoRelayStatus::Ok - ); - let error = runtime - .block_on(receiver) - .expect("native completion should settle") - .expect_err("revoked continuation should reject"); - assert!( - error - .to_string() - .contains("execution continuation is no longer active") + unsafe { + native_async_next_invoke_stream( + tool_ref, + invocation, + stream_ref, + accept_native_stream_item, + ptr::null_mut(), + ) + }, + NemoRelayStatus::InvalidArg ); - assert_eq!(provider_calls.load(Ordering::SeqCst), 0); - - unsafe { - native_string_free(invocation); - native_async_next_release(next_ref); - native_async_completion_release(completion_ref); - } -} -#[test] -fn native_async_next_result_supports_repeated_concurrent_calls() { - 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::Tool({ - let provider_calls = provider_calls.clone(); - Arc::new(move |value| { - provider_calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async move { - tokio::task::yield_now().await; - Ok(json!({ - "value": value, - "scope": crate::api::runtime::task_scope_top().uuid.to_string(), - })) - }) - }) - }), - runtime.handle().clone(), + let stream_next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })), + tokio::runtime::Handle::current(), None, )); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let first = native_string_from_json(&json!({"branch": "first"})).unwrap(); - let second = native_string_from_json(&json!({"branch": "second"})).unwrap(); - let (first_tx, first_rx) = tokio::sync::oneshot::channel::>(); - let (second_tx, second_rx) = - tokio::sync::oneshot::channel::>(); - let first_stack = create_scope_stack(); - let first_scope = first_stack - .read() - .unwrap_or_else(|error| error.into_inner()) - .top() - .uuid - .to_string(); - let second_stack = create_scope_stack(); - let second_scope = second_stack - .read() - .unwrap_or_else(|error| error.into_inner()) - .top() - .uuid - .to_string(); - + let next_ref = Arc::into_raw(stream_next) as *const NemoRelayNativeAsyncNext; assert_eq!( - with_scope_stack(first_stack, || unsafe { - native_async_next_invoke_result( + unsafe { + native_async_next_invoke_stream( next_ref, - first, - complete_native_next_result, - Box::into_raw(Box::new(first_tx)).cast(), + invocation, + ptr::null(), + accept_native_stream_item, + ptr::null_mut(), ) - }), - NemoRelayStatus::Ok + }, + NemoRelayStatus::NullPointer ); assert_eq!( - with_scope_stack(second_stack, || unsafe { - native_async_next_invoke_result( + unsafe { + native_async_next_invoke_stream( next_ref, - second, - complete_native_next_result, - Box::into_raw(Box::new(second_tx)).cast(), + invocation, + stream_ref, + accept_native_stream_item, + ptr::null_mut(), ) - }), - NemoRelayStatus::Ok - ); - let (first_result, second_result) = - runtime.block_on(async { tokio::join!(first_rx, second_rx) }); - - assert_eq!( - first_result.unwrap().unwrap(), - json!({ - "value": {"branch": "first"}, - "scope": first_scope, - }) - ); - assert_eq!( - second_result.unwrap().unwrap(), - json!({ - "value": {"branch": "second"}, - "scope": second_scope, - }) + }, + NemoRelayStatus::InvalidJson ); - assert_eq!(provider_calls.load(Ordering::SeqCst), 2); unsafe { - native_string_free(first); - native_string_free(second); + native_string_free(invocation); + native_async_next_release(tool_ref); native_async_next_release(next_ref); + native_async_stream_release(stream_ref); } } -#[test] -fn native_async_next_result_uses_captured_scope_on_an_unbound_plugin_thread() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let captured_stack = create_scope_stack(); - let captured_scope = captured_stack - .read() - .unwrap_or_else(|error| error.into_inner()) - .top() - .uuid - .to_string(); - let next = with_scope_stack(captured_stack, || { - Arc::new(NativeAsyncNext::new( - NativeAsyncNextInner::Tool(Arc::new(|value| { - Box::pin(async move { - Ok(json!({ - "value": value, - "scope": crate::api::runtime::task_scope_top().uuid.to_string(), - })) - }) - })), - runtime.handle().clone(), - None, - )) - }); - let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; - let invocation = native_string_from_json(&json!({"thread": "plugin"})).unwrap(); - let (sender, receiver) = tokio::sync::oneshot::channel::>(); - let next_address = next_ref as usize; - let invocation_address = invocation as usize; - let sender_address = Box::into_raw(Box::new(sender)) as usize; +struct InvokeNativeNextThenReturnState { + callback_state: u32, + invoke_status: AtomicUsize, + started: Mutex>, +} + +unsafe extern "C" fn invoke_native_next_then_return_state( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + let state = unsafe { &*user_data.cast::() }; + let status = unsafe { native_async_next_invoke(next, invocation_json, completion) }; + state + .invoke_status + .store(status as usize, Ordering::Release); + if status == NemoRelayStatus::Ok { + let _ = state + .started + .lock() + .unwrap_or_else(|error| error.into_inner()) + .recv_timeout(Duration::from_secs(1)); + } + unsafe { native_async_next_release(next) }; + state.callback_state +} + +unsafe extern "C" fn invoke_native_stream_next_then_return_state( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + stream: *const NemoRelayNativeAsyncStream, +) -> u32 { + let state = unsafe { &*user_data.cast::() }; + let request = read_native_string(invocation_json) + .ok() + .and_then(|invocation| serde_json::from_str::(&invocation).ok()) + .and_then(|invocation| invocation.get("request").cloned()) + .and_then(|request| native_string_from_json(&request)); + let status = if let Some(request) = request { + let status = unsafe { + native_async_next_invoke_stream( + next, + request, + stream, + accept_native_stream_item, + ptr::null_mut(), + ) + }; + unsafe { native_string_free(request) }; + status + } else { + NemoRelayStatus::InvalidJson + }; + state + .invoke_status + .store(status as usize, Ordering::Release); + if status == NemoRelayStatus::Ok { + let _ = state + .started + .lock() + .unwrap_or_else(|error| error.into_inner()) + .recv_timeout(Duration::from_secs(1)); + } + unsafe { + native_async_next_release(next); + native_async_stream_release(stream); + } + state.callback_state +} + +unsafe extern "C" fn invoke_detached_next_and_finish_replacement_stream( + user_data: *mut c_void, + invocation_json: *const NemoRelayNativeString, + next: *const NemoRelayNativeAsyncNext, + stream: *const NemoRelayNativeAsyncStream, +) -> u32 { + let invoke_status = unsafe { &*user_data.cast::() }; + let request = read_native_string(invocation_json) + .ok() + .and_then(|invocation| serde_json::from_str::(&invocation).ok()) + .and_then(|invocation| invocation.get("request").cloned()) + .and_then(|request| native_string_from_json(&request)); + let status = if let Some(request) = request { + let status = unsafe { + native_async_next_invoke_stream( + next, + request, + stream, + accept_native_stream_item, + ptr::null_mut(), + ) + }; + unsafe { native_string_free(request) }; + status + } else { + NemoRelayStatus::InvalidJson + }; + invoke_status.store(status as usize, Ordering::Release); + let replacement = native_string_from_json(&json!({"source": "replacement"})).unwrap(); + unsafe { + let _ = native_async_stream_push_json(stream, replacement); + native_string_free(replacement); + let _ = native_async_stream_finish(stream); + native_async_next_release(next); + native_async_stream_release(stream); + } + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +struct FailingNativeCodec; + +impl LlmCodec for FailingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + Err(FlowError::Internal("request decode rejected".into())) + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + Err(FlowError::Internal("request encode rejected".into())) + } +} + +impl LlmResponseCodec for FailingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + Err(FlowError::Internal("response decode rejected".into())) + } +} + +struct PanickingNativeCodec; + +impl LlmCodec for PanickingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + panic!("request decode panic") + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + panic!("request encode panic") + } +} + +impl LlmResponseCodec for PanickingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + panic!("response decode panic") + } +} + +#[test] +fn native_string_and_json_helpers_cover_abi_boundaries() { + assert_native_string_allocation_boundaries(); + assert_native_error_string_boundaries(); + assert_native_json_parsing_boundaries(); + assert_native_json_output_and_host_api(); +} + +fn assert_native_string_allocation_boundaries() { + clear_native_last_error(); + assert_eq!( + unsafe { native_string_new(ptr::null(), 0, ptr::null_mut()) }, + NemoRelayStatus::NullPointer + ); + assert_last_error_contains("out string pointer is null"); + + let mut out = ptr::null_mut(); + assert_eq!( + unsafe { native_string_new(ptr::null(), 1, &mut out) }, + NemoRelayStatus::NullPointer + ); + assert!(out.is_null()); + assert_last_error_contains("string data pointer is null"); + + let invalid_utf8 = [0xff]; + assert_eq!( + unsafe { native_string_new(invalid_utf8.as_ptr(), invalid_utf8.len(), &mut out) }, + NemoRelayStatus::InvalidUtf8 + ); + assert!(out.is_null()); + assert_last_error_contains("not valid UTF-8"); + + let text = native_string("hello"); + assert_eq!(unsafe { native_string_len(text) }, 5); + assert_eq!( + unsafe { std::slice::from_raw_parts(native_string_data(text), 5) }, + b"hello" + ); + assert!(unsafe { native_string_data(ptr::null()) }.is_null()); + assert_eq!(unsafe { native_string_len(ptr::null()) }, 0); + assert_eq!(take_native_string(text).unwrap(), "hello"); + unsafe { native_string_free(ptr::null_mut()) }; + + let empty = native_string(""); + assert_eq!(read_native_string(empty).unwrap(), ""); + unsafe { native_string_free(empty) }; + assert_eq!(read_native_string(ptr::null()).unwrap(), ""); +} + +fn assert_native_error_string_boundaries() { + let bad = Box::into_raw(Box::new(NativeHostString(vec![0xff]))) as *mut NemoRelayNativeString; + assert!(read_native_string(bad).is_err()); + assert_eq!( + optional_json_from_native_string(bad, "bad json"), + Err(NemoRelayStatus::InvalidUtf8) + ); + unsafe { native_last_error_set(bad) }; + assert_last_error_contains("not valid UTF-8"); + unsafe { native_string_free(bad) }; + + let message = native_string("explicit native error"); + unsafe { native_last_error_set(message) }; + assert_eq!( + native_last_error_message().as_deref(), + Some("explicit native error") + ); + unsafe { native_string_free(message) }; + unsafe { native_last_error_clear() }; + assert!(native_last_error_message().is_none()); + + set_native_last_error("specific fallback"); + assert!( + json_from_native_string(ptr::null_mut(), "generic fallback") + .unwrap_err() + .to_string() + .contains("specific fallback") + ); + clear_native_last_error(); + assert!( + json_from_native_string(ptr::null_mut(), "generic fallback") + .unwrap_err() + .to_string() + .contains("generic fallback") + ); +} + +fn assert_native_json_parsing_boundaries() { + let invalid_json = native_string("{"); + assert!( + take_json_from_native_string(invalid_json, "unused") + .unwrap_err() + .to_string() + .contains("invalid JSON") + ); + assert_eq!( + optional_json_from_native_string(ptr::null(), "optional"), + Ok(None) + ); + let valid_json = native_string(r#"{"value":1}"#); + assert_eq!( + optional_json_from_native_string(valid_json, "optional").unwrap(), + Some(json!({"value": 1})) + ); + unsafe { native_string_free(valid_json) }; + let invalid_json = native_string("not-json"); + assert_eq!( + optional_json_from_native_string(invalid_json, "optional"), + Err(NemoRelayStatus::InvalidJson) + ); + assert_last_error_contains("optional is not valid JSON"); + unsafe { native_string_free(invalid_json) }; + + assert_eq!( + parse_json_arg(ptr::null(), "null JSON").unwrap_err(), + NemoRelayStatus::InvalidJson + ); + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "test"}), + }; + let request_json = native_string_from_json(&serde_json::to_value(&request).unwrap()).unwrap(); + assert_eq!( + parse_llm_request_arg(request_json, "request").unwrap(), + request + ); + unsafe { native_string_free(request_json) }; + let wrong_shape = native_string(r#"{"headers":[]}"#); + assert_eq!( + parse_llm_request_arg(wrong_shape, "request").unwrap_err(), + NemoRelayStatus::InvalidJson + ); + assert_last_error_contains("was not an LLM request"); + unsafe { native_string_free(wrong_shape) }; +} + +fn assert_native_json_output_and_host_api() { + assert_eq!( + write_native_json(&json!({"ok": true}), ptr::null_mut()), + NemoRelayStatus::NullPointer + ); + let mut json_out = ptr::null_mut(); + assert_eq!( + write_native_json(&json!({"ok": true}), &mut json_out), + NemoRelayStatus::Ok + ); + assert_eq!( + take_json_from_native_string(json_out, "unused").unwrap(), + json!({"ok": true}) + ); + + let host_api = unsafe { &*native_host_api() }; + assert_eq!(host_api.abi_version, NEMO_RELAY_NATIVE_ABI_VERSION); + assert_eq!( + host_api.struct_size, + std::mem::size_of::() + ); +} + +#[test] +fn native_async_next_abi_runs_tool_llm_and_stream_continuations() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let cases: Vec<(NativeAsyncNextInner, Json, Json)> = vec![ + ( + NativeAsyncNextInner::Tool(Arc::new(|value| Box::pin(async move { Ok(value) }))), + json!({"tool": true}), + json!({"result": {"tool": true}, "pending_marks": []}), + ), + ( + NativeAsyncNextInner::Llm(Arc::new(|request| { + Box::pin(async move { Ok(request.content) }) + })), + serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"llm": true}), + }) + .unwrap(), + json!({"llm": true}), + ), + ]; + + for (inner, invocation, expected) in cases { + let next = Arc::new(NativeAsyncNext::new(inner, runtime.handle().clone(), None)); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + 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), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let invocation = native_string_from_json(&invocation).unwrap(); + assert_eq!( + unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, + NemoRelayStatus::Ok + ); + assert_eq!(runtime.block_on(receiver).unwrap().unwrap(), expected); + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + native_async_completion_release(completion_ref); + } + } + + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::LlmStream(Arc::new(|_request| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![ + Ok(json!({"chunk": 1})), + Ok(json!({"chunk": 2})), + ]))) + }) + })), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let (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), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let invocation = native_string_from_json( + &serde_json::to_value(LlmRequest { + headers: Map::new(), + content: json!({"stream": true}), + }) + .unwrap(), + ) + .unwrap(); + assert_eq!( + unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("async_next_invoke_stream"); + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + native_async_completion_release(completion_ref); + } +} + +#[test] +fn native_async_next_reports_a_revoked_continuation_without_calling_the_provider() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let provider_calls = Arc::new(AtomicUsize::new(0)); + let (lease, guard) = MiddlewareContinuationLease::capture(); + let next = Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool({ + let provider_calls = provider_calls.clone(); + Arc::new(move |value| { + let provider_calls = provider_calls.clone(); + let invocation = lease.begin(); + Box::pin(async move { + invocation? + .invoke(|| async move { + provider_calls.fetch_add(1, Ordering::SeqCst); + Ok(value) + }) + .await + }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + 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), + before_settlement_lock: None, + _callback_user_data: None, + }); + let completion_ref = + Arc::into_raw(Arc::clone(&completion)) as *const NemoRelayNativeAsyncCompletion; + let invocation = native_string_from_json(&json!({"tool": true})).unwrap(); + + drop(guard); + assert_eq!( + unsafe { native_async_next_invoke(next_ref, invocation, completion_ref) }, + NemoRelayStatus::Ok + ); + let error = runtime + .block_on(receiver) + .expect("native completion should settle") + .expect_err("revoked continuation should reject"); + assert!( + error + .to_string() + .contains("execution continuation is no longer active") + ); + assert_eq!(provider_calls.load(Ordering::SeqCst), 0); + + unsafe { + native_string_free(invocation); + native_async_next_release(next_ref); + native_async_completion_release(completion_ref); + } +} + +#[test] +fn native_async_next_result_supports_repeated_concurrent_calls() { + 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::Tool({ + let provider_calls = provider_calls.clone(); + Arc::new(move |value| { + provider_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + tokio::task::yield_now().await; + Ok(json!({ + "value": value, + "scope": crate::api::runtime::task_scope_top().uuid.to_string(), + })) + }) + }) + }), + runtime.handle().clone(), + None, + )); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let first = native_string_from_json(&json!({"branch": "first"})).unwrap(); + let second = native_string_from_json(&json!({"branch": "second"})).unwrap(); + let (first_tx, first_rx) = tokio::sync::oneshot::channel::>(); + let (second_tx, second_rx) = + tokio::sync::oneshot::channel::>(); + let first_stack = create_scope_stack(); + let first_scope = first_stack + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid + .to_string(); + let second_stack = create_scope_stack(); + let second_scope = second_stack + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid + .to_string(); + + assert_eq!( + with_scope_stack(first_stack, || unsafe { + native_async_next_invoke_result( + next_ref, + first, + complete_native_next_result, + Box::into_raw(Box::new(first_tx)).cast(), + ) + }), + NemoRelayStatus::Ok + ); + assert_eq!( + with_scope_stack(second_stack, || unsafe { + native_async_next_invoke_result( + next_ref, + second, + complete_native_next_result, + Box::into_raw(Box::new(second_tx)).cast(), + ) + }), + NemoRelayStatus::Ok + ); + let (first_result, second_result) = + runtime.block_on(async { tokio::join!(first_rx, second_rx) }); + + assert_eq!( + first_result.unwrap().unwrap(), + json!({ + "value": {"branch": "first"}, + "scope": first_scope, + }) + ); + assert_eq!( + second_result.unwrap().unwrap(), + json!({ + "value": {"branch": "second"}, + "scope": second_scope, + }) + ); + assert_eq!(provider_calls.load(Ordering::SeqCst), 2); + + unsafe { + native_string_free(first); + native_string_free(second); + 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() + .enable_all() + .build() + .unwrap(); + let captured_stack = create_scope_stack(); + let captured_scope = captured_stack + .read() + .unwrap_or_else(|error| error.into_inner()) + .top() + .uuid + .to_string(); + let next = with_scope_stack(captured_stack, || { + Arc::new(NativeAsyncNext::new( + NativeAsyncNextInner::Tool(Arc::new(|value| { + Box::pin(async move { + Ok(json!({ + "value": value, + "scope": crate::api::runtime::task_scope_top().uuid.to_string(), + })) + }) + })), + runtime.handle().clone(), + None, + )) + }); + let next_ref = Arc::into_raw(next) as *const NemoRelayNativeAsyncNext; + let invocation = native_string_from_json(&json!({"thread": "plugin"})).unwrap(); + let (sender, receiver) = tokio::sync::oneshot::channel::>(); + let next_address = next_ref as usize; + let invocation_address = invocation as usize; + let sender_address = Box::into_raw(Box::new(sender)) as usize; assert_eq!( std::thread::spawn(move || unsafe { @@ -2581,6 +3310,10 @@ unsafe extern "C" fn count_scope_callback(user_data: *mut c_void) -> NemoRelaySt NemoRelayStatus::Ok } +unsafe extern "C" fn fail_scope_callback(_user_data: *mut c_void) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + #[test] fn native_scope_stack_abi_covers_lifecycle_and_validation() { let _runtime_guard = crate::shared_runtime::runtime_owner_test_mutex() @@ -2589,6 +3322,18 @@ fn native_scope_stack_abi_covers_lifecycle_and_validation() { crate::shared_runtime::reset_runtime_owner_for_tests(); let _global_context_restore = GlobalContextRestore::replace_with_empty(); let _restore = ThreadScopeStackRestore::capture(); + + assert_native_scope_stack_null_validation(); + let stack = create_active_native_scope_stack(); + let strings = NativeScopeTestStrings::new(); + let scope = assert_native_scope_lifecycle(&strings); + assert_native_scope_push_validation(&strings); + assert_native_scope_pop_and_mark_validation(scope, &strings); + assert_native_scope_stack_binding_lifecycle(); + free_native_scope_test_resources(stack, scope, strings); +} + +fn assert_native_scope_stack_null_validation() { assert_eq!( unsafe { native_scope_stack_create(ptr::null_mut()) }, NemoRelayStatus::NullPointer @@ -2615,7 +3360,25 @@ fn native_scope_stack_abi_covers_lifecycle_and_validation() { unsafe { native_scope_get_current(ptr::null_mut()) }, NemoRelayStatus::NullPointer ); + assert_eq!( + unsafe { + native_scope_push( + ptr::null(), + NemoRelayNativeScopeType::Custom, + ptr::null(), + 0, + ptr::null(), + ptr::null(), + ptr::null(), + ptr::null(), + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); +} +fn create_active_native_scope_stack() -> *mut NemoRelayNativeScopeStack { let mut stack = ptr::null_mut(); assert_eq!( unsafe { native_scope_stack_create(&mut stack) }, @@ -2637,78 +3400,246 @@ fn native_scope_stack_abi_covers_lifecycle_and_validation() { (&calls as *const AtomicUsize).cast_mut().cast(), ) }, - NemoRelayStatus::Ok + NemoRelayStatus::Ok + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + unsafe { native_scope_stack_with_current(stack, fail_scope_callback, ptr::null_mut()) }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("scope-stack callback returned InvalidArg"); + + stack +} + +struct NativeScopeTestStrings { + name: *mut NemoRelayNativeString, + data: *mut NemoRelayNativeString, + metadata: *mut NemoRelayNativeString, + input: *mut NemoRelayNativeString, + mark_name: *mut NemoRelayNativeString, + output: *mut NemoRelayNativeString, + invalid: *mut NemoRelayNativeString, +} + +impl NativeScopeTestStrings { + fn new() -> Self { + Self { + name: native_string("native-scope"), + data: native_string(r#"{"source":"native"}"#), + metadata: native_string(r#"{"test":true}"#), + input: native_string(r#"{"input":1}"#), + mark_name: native_string("native-mark"), + output: native_string(r#"{"output":1}"#), + invalid: native_string("not-json"), + } + } +} + +fn assert_native_scope_lifecycle( + strings: &NativeScopeTestStrings, +) -> *mut NemoRelayNativeScopeHandle { + let timestamp = 0_i64; + let mut scope = ptr::null_mut(); + assert_eq!( + unsafe { + native_scope_push( + strings.name, + NemoRelayNativeScopeType::Custom, + ptr::null(), + 0, + strings.data, + strings.metadata, + strings.input, + ×tamp, + &mut scope, + ) + }, + NemoRelayStatus::Ok + ); + assert!(!scope.is_null()); + assert_eq!(native_scope_ref(scope).unwrap().name, "native-scope"); + + let mut current = ptr::null_mut(); + assert_eq!( + unsafe { native_scope_get_current(&mut current) }, + NemoRelayStatus::Ok + ); + assert_eq!(native_scope_ref(current).unwrap().name, "native-scope"); + unsafe { native_scope_handle_free(current) }; + + assert_eq!( + unsafe { + native_emit_mark( + strings.mark_name, + scope, + strings.data, + strings.metadata, + ×tamp, + ) + }, + NemoRelayStatus::Ok + ); + assert_eq!( + unsafe { native_scope_pop(scope, strings.output, strings.metadata, ×tamp) }, + NemoRelayStatus::Ok + ); + + scope +} + +fn assert_native_scope_push_validation(strings: &NativeScopeTestStrings) { + let mut invalid_scope = ptr::null_mut(); + assert_eq!( + unsafe { + native_scope_push( + strings.name, + NemoRelayNativeScopeType::Custom, + ptr::null(), + 0, + strings.invalid, + ptr::null(), + ptr::null(), + ptr::null(), + &mut invalid_scope, + ) + }, + NemoRelayStatus::InvalidJson + ); + assert!(invalid_scope.is_null()); + let invalid_const = strings.invalid.cast_const(); + for (data_arg, metadata_arg, input_arg) in [ + (ptr::null(), invalid_const, ptr::null()), + (ptr::null(), ptr::null(), invalid_const), + ] { + assert_eq!( + unsafe { + native_scope_push( + strings.name, + NemoRelayNativeScopeType::Custom, + ptr::null(), + 0, + data_arg, + metadata_arg, + input_arg, + ptr::null(), + &mut invalid_scope, + ) + }, + NemoRelayStatus::InvalidJson + ); + assert!(invalid_scope.is_null()); + } + let invalid_timestamp = i64::MAX; + assert_eq!( + unsafe { + native_scope_push( + strings.name, + NemoRelayNativeScopeType::Custom, + ptr::null(), + 0, + ptr::null(), + ptr::null(), + ptr::null(), + &invalid_timestamp, + &mut invalid_scope, + ) + }, + NemoRelayStatus::InvalidArg ); - assert_eq!(calls.load(Ordering::SeqCst), 1); - let name = native_string("native-scope"); - let data = native_string(r#"{"source":"native"}"#); - let metadata = native_string(r#"{"test":true}"#); - let input = native_string(r#"{"input":1}"#); - let timestamp = 0_i64; - let mut scope = ptr::null_mut(); + let invalid_name = Box::into_raw(Box::new(NativeHostString(vec![0xff]))).cast(); assert_eq!( unsafe { native_scope_push( - name, + invalid_name, NemoRelayNativeScopeType::Custom, ptr::null(), 0, - data, - metadata, - input, - ×tamp, - &mut scope, + ptr::null(), + ptr::null(), + ptr::null(), + ptr::null(), + &mut invalid_scope, ) }, - NemoRelayStatus::Ok + NemoRelayStatus::InvalidUtf8 ); - assert!(!scope.is_null()); - assert_eq!(native_scope_ref(scope).unwrap().name, "native-scope"); - - let mut current = ptr::null_mut(); assert_eq!( - unsafe { native_scope_get_current(&mut current) }, - NemoRelayStatus::Ok + unsafe { + native_emit_mark( + invalid_name, + ptr::null(), + ptr::null(), + ptr::null(), + ptr::null(), + ) + }, + NemoRelayStatus::InvalidUtf8 ); - assert_eq!(native_scope_ref(current).unwrap().name, "native-scope"); - unsafe { native_scope_handle_free(current) }; + unsafe { native_string_free(invalid_name) }; +} - let mark_name = native_string("native-mark"); +fn assert_native_scope_pop_and_mark_validation( + scope: *mut NemoRelayNativeScopeHandle, + strings: &NativeScopeTestStrings, +) { + let invalid_timestamp = i64::MAX; assert_eq!( - unsafe { native_emit_mark(mark_name, scope, data, metadata, ×tamp) }, - NemoRelayStatus::Ok + unsafe { native_scope_pop(scope, strings.invalid, ptr::null(), ptr::null()) }, + NemoRelayStatus::InvalidJson ); - let output = native_string(r#"{"output":1}"#); assert_eq!( - unsafe { native_scope_pop(scope, output, metadata, ×tamp) }, - NemoRelayStatus::Ok + unsafe { native_scope_pop(scope, ptr::null(), strings.invalid, ptr::null()) }, + NemoRelayStatus::InvalidJson + ); + assert_eq!( + unsafe { native_scope_pop(scope, ptr::null(), ptr::null(), &invalid_timestamp) }, + NemoRelayStatus::InvalidArg ); - - let invalid = native_string("not-json"); - let mut invalid_scope = ptr::null_mut(); assert_eq!( unsafe { - native_scope_push( - name, - NemoRelayNativeScopeType::Custom, + native_emit_mark( + strings.mark_name, + scope, + strings.invalid, ptr::null(), - 0, - invalid, ptr::null(), + ) + }, + NemoRelayStatus::InvalidJson + ); + assert_eq!( + unsafe { + native_emit_mark( + strings.mark_name, + scope, ptr::null(), + strings.invalid, ptr::null(), - &mut invalid_scope, ) }, NemoRelayStatus::InvalidJson ); - assert!(invalid_scope.is_null()); + assert_eq!( + unsafe { + native_emit_mark( + strings.mark_name, + scope, + ptr::null(), + ptr::null(), + &invalid_timestamp, + ) + }, + NemoRelayStatus::InvalidArg + ); assert_eq!( unsafe { native_scope_pop(ptr::null(), ptr::null(), ptr::null(), ptr::null()) }, NemoRelayStatus::NullPointer ); +} +fn assert_native_scope_stack_binding_lifecycle() { let mut binding = ptr::null_mut(); assert_eq!( unsafe { native_scope_stack_capture_thread(&mut binding) }, @@ -2724,8 +3655,22 @@ fn native_scope_stack_abi_covers_lifecycle_and_validation() { NemoRelayStatus::Ok ); unsafe { native_scope_stack_binding_free(disposable_binding) }; +} - for value in [name, data, metadata, input, mark_name, output, invalid] { +fn free_native_scope_test_resources( + stack: *mut NemoRelayNativeScopeStack, + scope: *mut NemoRelayNativeScopeHandle, + strings: NativeScopeTestStrings, +) { + for value in [ + strings.name, + strings.data, + strings.metadata, + strings.input, + strings.mark_name, + strings.output, + strings.invalid, + ] { unsafe { native_string_free(value) }; } unsafe { @@ -2791,184 +3736,841 @@ unsafe extern "C" fn noop_json( NemoRelayStatus::Ok } -unsafe extern "C" fn noop_llm_conditional( - _user_data: *mut c_void, - _request_json: *const NemoRelayNativeString, - _out_reason: *mut *mut NemoRelayNativeString, -) -> NemoRelayStatus { - NemoRelayStatus::Ok -} +unsafe extern "C" fn noop_llm_conditional( + _user_data: *mut c_void, + _request_json: *const NemoRelayNativeString, + _out_reason: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} + +unsafe extern "C" fn noop_llm_request_intercept( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _request_json: *const NemoRelayNativeString, + _annotated_json: *const NemoRelayNativeString, + _out_outcome_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} + +unsafe extern "C" fn noop_llm_execution( + _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 { + NemoRelayStatus::Ok +} + +unsafe extern "C" fn noop_llm_stream_execution( + _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::Ok +} + +#[test] +fn native_registration_entrypoints_reject_null_contexts() { + unsafe { + assert_eq!( + native_plugin_context_register_subscriber( + ptr::null_mut(), + ptr::null(), + noop_subscriber, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_tool_sanitize_request_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_tool_sanitize_response_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_tool_conditional_execution_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_tool_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_tool_request_intercept( + ptr::null_mut(), + ptr::null(), + 0, + false, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_tool_execution_intercept( + ptr::null_mut(), + ptr::null(), + 0, + noop_tool_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_request_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_llm_request, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_response_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_conditional_execution_guardrail( + ptr::null_mut(), + ptr::null(), + 0, + noop_llm_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_request_intercept( + ptr::null_mut(), + ptr::null(), + 0, + false, + noop_llm_request_intercept, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_execution_intercept( + ptr::null_mut(), + ptr::null(), + 0, + noop_llm_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + native_plugin_context_register_llm_stream_execution_intercept( + ptr::null_mut(), + ptr::null(), + 0, + noop_llm_stream_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + } + assert_last_error_contains("plugin context is null"); +} + +#[cfg(unix)] +#[test] +fn native_registration_entrypoints_reject_invalid_host_contexts_and_names() { + let instance = Arc::new(NativePluginInstance { + plugin_kind: "test.native".into(), + relay_compat: "^0.7".into(), + allows_multiple_components: false, + plugin: Mutex::new(NemoRelayNativePluginV1::default()), + _library: libloading::os::unix::Library::this().into(), + }); + let mut invalid_host = NativeHostPluginContext { + ctx: ptr::null_mut(), + instance: Arc::clone(&instance), + }; + assert_eq!( + unsafe { + native_plugin_context_register_subscriber( + ptr::from_mut(&mut invalid_host).cast(), + ptr::null(), + noop_subscriber, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::NullPointer + ); -unsafe extern "C" fn noop_llm_request_intercept( - _user_data: *mut c_void, - _name: *const NemoRelayNativeString, - _request_json: *const NemoRelayNativeString, - _annotated_json: *const NemoRelayNativeString, - _out_outcome_json: *mut *mut NemoRelayNativeString, -) -> NemoRelayStatus { - NemoRelayStatus::Ok -} + let mut registration = PluginRegistrationContext::new(); + let mut host = NativeHostPluginContext { + ctx: ptr::from_mut(&mut registration), + instance, + }; + let ctx = ptr::from_mut(&mut host).cast(); -unsafe extern "C" fn noop_llm_execution( - _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 { - NemoRelayStatus::Ok + assert_registration_entrypoints_reject_invalid_names(ctx); + assert_registration_entrypoints_accept_valid_names(ctx); + assert_async_registration_entrypoints_validate_contracts(ctx); + assert_async_request_registration_rejects_legacy_relay_contract(); } -unsafe extern "C" fn noop_llm_stream_execution( - _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::Ok +#[cfg(unix)] +fn assert_registration_entrypoints_reject_invalid_names(ctx: *mut NemoRelayNativePluginContext) { + let invalid_name = Box::into_raw(Box::new(NativeHostString(vec![0xff]))).cast(); + unsafe { + assert_eq!( + native_plugin_context_register_subscriber( + ctx, + invalid_name, + noop_subscriber, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_tool_sanitize_request_guardrail( + ctx, + invalid_name, + 0, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_tool_sanitize_response_guardrail( + ctx, + invalid_name, + 0, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_tool_conditional_execution_guardrail( + ctx, + invalid_name, + 0, + noop_tool_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_tool_request_intercept( + ctx, + invalid_name, + 0, + false, + noop_tool_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_tool_execution_intercept( + ctx, + invalid_name, + 0, + noop_tool_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_request_guardrail( + ctx, + invalid_name, + 0, + noop_llm_request, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_response_guardrail( + ctx, + invalid_name, + 0, + noop_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_conditional_execution_guardrail( + ctx, + invalid_name, + 0, + noop_llm_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_request_intercept( + ctx, + invalid_name, + 0, + false, + noop_llm_request_intercept, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_execution_intercept( + ctx, + invalid_name, + 0, + noop_llm_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_llm_stream_execution_intercept( + ctx, + invalid_name, + 0, + noop_llm_stream_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_async_stream_middleware( + ctx, + invalid_name, + 0, + invoke_native_stream_next_then_return_state, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + assert_eq!( + native_plugin_context_register_async_middleware( + ctx, + NemoRelayNativeAsyncMiddlewareKind::ToolSanitizeRequest as u32, + invalid_name, + 0, + false, + resolve_async_static_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidUtf8 + ); + drop(Box::from_raw(invalid_name as *mut NativeHostString)); + } } -#[test] -fn native_registration_entrypoints_reject_null_contexts() { +#[cfg(unix)] +fn assert_registration_entrypoints_accept_valid_names(ctx: *mut NemoRelayNativePluginContext) { unsafe { + let name = native_string("registered"); assert_eq!( native_plugin_context_register_subscriber( - ptr::null_mut(), - ptr::null(), + ctx, + name, noop_subscriber, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok ); assert_eq!( native_plugin_context_register_tool_sanitize_request_guardrail( - ptr::null_mut(), - ptr::null(), + ctx, + name, 0, noop_tool_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok ); assert_eq!( native_plugin_context_register_tool_sanitize_response_guardrail( + ctx, + name, + 0, + noop_tool_json, ptr::null_mut(), - ptr::null(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_tool_conditional_execution_guardrail( + ctx, + name, + 0, + noop_tool_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_tool_request_intercept( + ctx, + name, 0, + false, noop_tool_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_tool_execution_intercept( + ctx, + name, + 0, + noop_tool_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_request_guardrail( + ctx, + name, + 0, + noop_llm_request, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_llm_sanitize_response_guardrail( + ctx, + name, + 0, + noop_json, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_llm_conditional_execution_guardrail( + ctx, + name, + 0, + noop_llm_conditional, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_llm_request_intercept( + ctx, + name, + 0, + false, + noop_llm_request_intercept, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + native_plugin_context_register_llm_execution_intercept( + ctx, + name, + 0, + noop_llm_execution, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok ); assert_eq!( - native_plugin_context_register_tool_conditional_execution_guardrail( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_llm_stream_execution_intercept( + ctx, + name, 0, - noop_tool_conditional, + noop_llm_stream_execution, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok ); + native_string_free(name); + } +} + +#[cfg(unix)] +fn assert_async_registration_entrypoints_validate_contracts( + ctx: *mut NemoRelayNativePluginContext, +) { + unsafe { assert_eq!( - native_plugin_context_register_tool_request_intercept( + native_plugin_context_register_async_stream_middleware( ptr::null_mut(), ptr::null(), 0, - false, - noop_tool_json, + invoke_native_stream_next_then_return_state, ptr::null_mut(), None, ), NemoRelayStatus::NullPointer ); assert_eq!( - native_plugin_context_register_tool_execution_intercept( + native_plugin_context_register_async_middleware( ptr::null_mut(), + 0, ptr::null(), 0, - noop_tool_execution, + false, + resolve_async_static_json, ptr::null_mut(), None, ), NemoRelayStatus::NullPointer ); + + let name = native_string("async-registered"); assert_eq!( - native_plugin_context_register_llm_sanitize_request_guardrail( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_middleware( + ctx, + u32::MAX, + name, 0, - noop_llm_request, + false, + resolve_async_static_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::InvalidArg ); assert_eq!( - native_plugin_context_register_llm_sanitize_response_guardrail( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_middleware( + ctx, + NemoRelayNativeAsyncMiddlewareKind::LlmStreamExecutionIntercept as u32, + name, 0, - noop_json, + false, + resolve_async_static_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::InvalidArg ); assert_eq!( - native_plugin_context_register_llm_conditional_execution_guardrail( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_middleware( + ctx, + NemoRelayNativeAsyncMiddlewareKind::ToolSanitizeRequest as u32, + name, 0, - noop_llm_conditional, + false, + resolve_async_static_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok ); assert_eq!( - native_plugin_context_register_llm_request_intercept( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_middleware( + ctx, + NemoRelayNativeAsyncMiddlewareKind::ToolSanitizeRequest as u32, + name, 0, false, - noop_llm_request_intercept, + resolve_async_static_json, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Internal ); + + let stream_name = native_string("async-stream-registered"); assert_eq!( - native_plugin_context_register_llm_execution_intercept( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_stream_middleware( + ctx, + stream_name, 0, - noop_llm_execution, + invoke_native_stream_next_then_return_state, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Ok ); assert_eq!( - native_plugin_context_register_llm_stream_execution_intercept( - ptr::null_mut(), - ptr::null(), + native_plugin_context_register_async_stream_middleware( + ctx, + stream_name, 0, - noop_llm_stream_execution, + invoke_native_stream_next_then_return_state, ptr::null_mut(), None, ), - NemoRelayStatus::NullPointer + NemoRelayStatus::Internal ); + native_string_free(name); + native_string_free(stream_name); + } +} + +#[cfg(unix)] +fn assert_async_request_registration_rejects_legacy_relay_contract() { + let instance = Arc::new(NativePluginInstance { + plugin_kind: "test.native.legacy".into(), + relay_compat: "^0.5".into(), + allows_multiple_components: false, + plugin: Mutex::new(NemoRelayNativePluginV1::default()), + _library: libloading::os::unix::Library::this().into(), + }); + let mut registration = PluginRegistrationContext::new(); + let mut host = NativeHostPluginContext { + ctx: ptr::from_mut(&mut registration), + instance, + }; + let name = native_string("legacy-request"); + assert_eq!( + unsafe { + native_plugin_context_register_async_middleware( + ptr::from_mut(&mut host).cast(), + NemoRelayNativeAsyncMiddlewareKind::LlmRequestIntercept as u32, + name, + 0, + false, + resolve_async_static_json, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::InvalidArg + ); + assert_last_error_contains("excludes Relay 0.5"); + unsafe { native_string_free(name) }; +} + +#[cfg(unix)] +unsafe extern "C" fn resolve_async_static_json( + user_data: *mut c_void, + _invocation_json: *const NemoRelayNativeString, + _next: *const NemoRelayNativeAsyncNext, + completion: *const NemoRelayNativeAsyncCompletion, +) -> u32 { + assert_eq!( + unsafe { native_async_completion_resolve_json(completion, user_data.cast()) }, + NemoRelayStatus::Ok + ); + NemoRelayNativeAsyncCallbackState::Complete as u32 +} + +#[cfg(unix)] +#[tokio::test] +async fn native_async_wrappers_validate_callback_result_shapes() { + let instance = Arc::new(NativePluginInstance { + plugin_kind: "test.native.async".into(), + relay_compat: "^0.7".into(), + allows_multiple_components: false, + plugin: Mutex::new(NemoRelayNativePluginV1::default()), + _library: libloading::os::unix::Library::this().into(), + }); + let result = native_string("true"); + let user_data = result.cast(); + + let tool_json = wrap_native_async_tool_json( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert_eq!( + tool_json("tool".into(), json!({})).await.unwrap(), + json!(true) + ); + + let tool_conditional = wrap_native_async_tool_conditional( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert!( + tool_conditional("tool".into(), json!({})) + .await + .unwrap_err() + .to_string() + .contains("expected string or null") + ); + + let request = LlmRequest { + headers: Map::new(), + content: json!({"messages": []}), + }; + let llm_conditional = wrap_native_async_llm_conditional( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert!( + llm_conditional(request.clone()) + .await + .unwrap_err() + .to_string() + .contains("expected string or null") + ); + + let sanitize_request = wrap_native_async_llm_sanitize_request( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert!( + sanitize_request( + request.clone(), + LlmSanitizeRequestContext::for_request_codec(None), + ) + .await + .is_err() + ); + + let sanitize_response = wrap_native_async_llm_sanitize_response( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert_eq!( + sanitize_response( + json!({"response": true}), + LlmSanitizeResponseContext::for_response_codec(None), + ) + .await + .unwrap(), + Some(json!(true)) + ); + + let request_intercept = wrap_native_async_llm_request_intercept( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert!( + request_intercept("model".into(), request, None) + .await + .unwrap_err() + .to_string() + .contains("invalid native async LLM intercept outcome") + ); + + let tool_execution = wrap_native_async_tool_execution( + Arc::clone(&instance), + resolve_async_static_json, + user_data, + None, + ); + assert!( + tool_execution("tool", json!({}), tool_next(Ok(Json::Null))) + .await + .unwrap_err() + .to_string() + .contains("invalid native async tool outcome") + ); + + let fields = EventSanitizeFields::default(); + let fields_result = native_string_from_json(&serde_json::to_value(&fields).unwrap()).unwrap(); + let event_sanitize = wrap_native_async_event_sanitize( + Arc::clone(&instance), + resolve_async_static_json, + fields_result.cast(), + None, + ); + let event = Event::Mark(crate::api::event::MarkEvent::new( + crate::api::event::BaseEvent::builder() + .name("native-async-event") + .build(), + None, + None, + )); + assert_eq!( + event_sanitize(Arc::new(event), fields.clone()) + .await + .unwrap(), + fields + ); + + drop(event_sanitize); + drop(tool_execution); + drop(request_intercept); + drop(sanitize_response); + drop(sanitize_request); + drop(llm_conditional); + drop(tool_conditional); + drop(tool_json); + unsafe { + native_string_free(fields_result); + native_string_free(result); } - assert_last_error_contains("plugin context is null"); } #[test] @@ -3100,52 +4702,125 @@ fn native_codec_operations_contain_codec_panics() { &mut output, ) }, - NemoRelayStatus::Internal + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode panicked"); + + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + annotated_json, + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode panicked"); + + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&response_codec).cast(), + response_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode panicked"); + + unsafe { + native_string_free(request_json); + native_string_free(annotated_json); + native_string_free(response_json); + } +} + +#[test] +fn native_codec_operations_clear_output_slots_on_null_arguments() { + let request_codec = NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); + let response_codec = + NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); + let request_json = native_string( + r#"{"headers":{},"content":{"model":"gpt-test","messages":[{"role":"user","content":"hello"}]}}"#, + ); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "hello"}] + }), + }; + let annotated = OpenAIChatCodec.decode(&request).unwrap(); + let annotated_json = native_string(&serde_json::to_string(&annotated).unwrap()); + + assert_eq!( + unsafe { native_llm_request_codec_decode(ptr::null(), request_json, &mut ptr::null_mut()) }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&request_codec).cast(), + request_json, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::null(), + annotated_json, + request_json, + &mut ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + annotated_json, + ptr::null(), + &mut ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer ); - assert!(output.is_null()); - assert_last_error_contains("request codec decode panicked"); - assert_eq!( unsafe { native_llm_request_codec_encode( ptr::from_ref(&request_codec).cast(), annotated_json, request_json, - &mut output, + ptr::null_mut(), ) }, - NemoRelayStatus::Internal + 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 panicked"); - assert_eq!( unsafe { native_llm_response_codec_decode( ptr::from_ref(&response_codec).cast(), - response_json, - &mut output, + request_json, + ptr::null_mut(), ) }, - NemoRelayStatus::Internal - ); - assert!(output.is_null()); - assert_last_error_contains("response codec decode panicked"); - - unsafe { - native_string_free(request_json); - native_string_free(annotated_json); - native_string_free(response_json); - } -} - -#[test] -fn native_codec_operations_clear_output_slots_on_null_arguments() { - let request_codec = NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); - let response_codec = - NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); - let request_json = native_string( - r#"{"headers":{},"content":{"model":"gpt-test","messages":[{"role":"user","content":"hello"}]}}"#, + NemoRelayStatus::NullPointer ); let request_decode_sentinel = native_string("request-decode-sentinel"); @@ -3200,6 +4875,7 @@ fn native_codec_operations_clear_output_slots_on_null_arguments() { 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); } } @@ -3226,6 +4902,127 @@ unsafe extern "C" fn tool_json_error( NemoRelayStatus::InvalidArg } +#[cfg(unix)] +unsafe extern "C" fn tool_conditional_error( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _args_json: *const NemoRelayNativeString, + out_reason: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_reason = native_string("discarded reason") }; + set_native_last_error("tool conditional failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn tool_conditional_reason( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _args_json: *const NemoRelayNativeString, + out_reason: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_reason = native_string("blocked by tool policy") }; + NemoRelayStatus::Ok +} + +#[cfg(unix)] +unsafe extern "C" fn llm_conditional_error( + _user_data: *mut c_void, + _request_json: *const NemoRelayNativeString, + out_reason: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_reason = native_string("discarded reason") }; + set_native_last_error("LLM conditional failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn llm_conditional_reason( + _user_data: *mut c_void, + _request_json: *const NemoRelayNativeString, + out_reason: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_reason = native_string("blocked by LLM policy") }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn llm_request_error( + _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}"#) }; + set_native_last_error("LLM request sanitizer failed"); + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn llm_response_error( + _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}"#) }; + set_native_last_error("LLM response sanitizer failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn tool_execution_error( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _args_json: *const NemoRelayNativeString, + _next_fn: NemoRelayNativeToolNextFn, + _next_ctx: *mut c_void, + out_outcome_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_outcome_json = native_string(r#"{"discarded":true}"#) }; + set_native_last_error("tool execution failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn llm_request_intercept_error( + _user_data: *mut c_void, + _name: *const NemoRelayNativeString, + _request_json: *const NemoRelayNativeString, + _annotated_json: *const NemoRelayNativeString, + out_outcome_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_outcome_json = native_string(r#"{"discarded":true}"#) }; + set_native_last_error("LLM request intercept failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +unsafe extern "C" fn llm_execution_error( + _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}"#) }; + set_native_last_error("LLM execution failed"); + NemoRelayStatus::InvalidArg +} + +#[cfg(unix)] +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 { + unsafe { *out_stream = NemoRelayNativeLlmStreamV1::default() }; + set_native_last_error("LLM stream execution failed"); + NemoRelayStatus::InvalidArg +} + unsafe extern "C" fn llm_request_echo( _user_data: *mut c_void, request_json: *const NemoRelayNativeString, @@ -3321,7 +5118,7 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { LlmSanitizeRequestContext::default(), ) .unwrap(), - Some(request) + Some(request.clone()) ); let request = LlmRequest { @@ -3336,7 +5133,7 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { LlmSanitizeRequestContext::default(), ) .unwrap(), - Some(request) + Some(request.clone()) ); let response = json!({"message": "alias"}); @@ -3348,7 +5145,169 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { LlmSanitizeResponseContext::default(), ) .unwrap(), - Some(response) + Some(response.clone()) + ); + + assert!( + call_llm_sanitize_request_callback( + llm_request_error, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap_err() + .to_string() + .contains("LLM request sanitizer failed") + ); + assert_eq!( + call_llm_sanitize_request_callback( + noop_llm_request, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + None + ); + assert!( + call_llm_sanitize_response_callback( + llm_response_error, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::default(), + ) + .unwrap_err() + .to_string() + .contains("LLM response sanitizer failed") + ); + assert_eq!( + call_llm_sanitize_response_callback( + noop_json, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::default(), + ) + .unwrap(), + None + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn native_callback_wrappers_release_error_outputs_and_preserve_reasons() { + let instance = Arc::new(NativePluginInstance { + plugin_kind: "test.native.callback-errors".into(), + relay_compat: "^0.7".into(), + allows_multiple_components: false, + plugin: Mutex::new(NemoRelayNativePluginV1::default()), + _library: libloading::os::unix::Library::this().into(), + }); + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "test"}), + }; + + let tool_conditional = wrap_tool_conditional_fn( + Arc::clone(&instance), + tool_conditional_error, + ptr::null_mut(), + None, + ); + assert!( + tool_conditional("tool".into(), json!({})) + .await + .unwrap_err() + .to_string() + .contains("tool conditional failed") + ); + let tool_conditional = wrap_tool_conditional_fn( + Arc::clone(&instance), + tool_conditional_reason, + ptr::null_mut(), + None, + ); + assert_eq!( + tool_conditional("tool".into(), json!({})).await.unwrap(), + Some("blocked by tool policy".into()) + ); + + let tool_execution = wrap_tool_execution_fn( + Arc::clone(&instance), + tool_execution_error, + ptr::null_mut(), + None, + ); + assert!( + tool_execution("tool", json!({}), tool_next(Ok(Json::Null))) + .await + .unwrap_err() + .to_string() + .contains("tool execution failed") + ); + + let llm_conditional = wrap_llm_conditional_fn( + Arc::clone(&instance), + llm_conditional_error, + ptr::null_mut(), + None, + ); + assert!( + llm_conditional(request.clone()) + .await + .unwrap_err() + .to_string() + .contains("LLM conditional failed") + ); + let llm_conditional = wrap_llm_conditional_fn( + Arc::clone(&instance), + llm_conditional_reason, + ptr::null_mut(), + None, + ); + assert_eq!( + llm_conditional(request.clone()).await.unwrap(), + Some("blocked by LLM policy".into()) + ); + + let request_intercept = wrap_llm_request_intercept_fn( + Arc::clone(&instance), + llm_request_intercept_error, + ptr::null_mut(), + None, + ); + assert!( + request_intercept("model".into(), request.clone(), None) + .await + .unwrap_err() + .to_string() + .contains("LLM request intercept failed") + ); + + let llm_execution = wrap_llm_execution_fn( + Arc::clone(&instance), + llm_execution_error, + ptr::null_mut(), + None, + ); + assert!( + llm_execution("model", request.clone(), llm_next(Ok(Json::Null))) + .await + .unwrap_err() + .to_string() + .contains("LLM execution failed") + ); + + let stream_next: LlmStreamExecutionNextFn = + Arc::new(|_| Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) })); + let llm_stream_execution = + wrap_llm_stream_execution_fn(instance, llm_stream_execution_error, ptr::null_mut(), None); + assert!( + llm_stream_execution("model", request, stream_next) + .await + .err() + .expect("native stream callback should fail") + .to_string() + .contains("LLM stream execution failed") ); } @@ -3572,6 +5531,16 @@ fn native_non_streaming_continuations_cover_success_and_error_paths() { NemoRelayStatus::NotFound ); 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; + assert_eq!( + unsafe { native_tool_next(args, next, &mut out) }, + NemoRelayStatus::Internal + ); + assert_last_error_contains("native tool next panicked"); + unsafe { drop(Box::from_raw(next as *mut ToolExecutionNextFn)) }; unsafe { native_string_free(args) }; let request = LlmRequest { @@ -3599,6 +5568,16 @@ fn native_non_streaming_continuations_cover_success_and_error_paths() { unsafe { native_llm_next(request_json, next, &mut out) }, NemoRelayStatus::GuardrailRejected ); + 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; + assert_eq!( + unsafe { native_llm_next(request_json, next, &mut out) }, + NemoRelayStatus::Internal + ); + assert_last_error_contains("native LLM next panicked"); unsafe { drop(Box::from_raw(next as *mut LlmExecutionNextFn)); native_string_free(request_json); @@ -3611,7 +5590,9 @@ enum NativeStreamItem { InvalidJson, Null, Error(NemoRelayStatus), + ErrorWithJson(NemoRelayStatus), End, + EndWithJson, } struct TestNativeStream { @@ -3634,7 +5615,15 @@ unsafe extern "C" fn test_native_stream_poll( } NativeStreamItem::Null => NemoRelayStatus::Ok, NativeStreamItem::Error(status) => status, + NativeStreamItem::ErrorWithJson(status) => { + unsafe { *out_json = native_string(r#"{"discarded":true}"#) }; + status + } NativeStreamItem::End => NemoRelayStatus::StreamEnd, + NativeStreamItem::EndWithJson => { + unsafe { *out_json = native_string(r#"{"discarded":true}"#) }; + NemoRelayStatus::StreamEnd + } } } @@ -3700,6 +5689,7 @@ async fn native_stream_adapter_covers_chunks_end_errors_and_cancellation() { NativeStreamItem::InvalidJson, NativeStreamItem::Null, NativeStreamItem::Error(NemoRelayStatus::InvalidArg), + NativeStreamItem::ErrorWithJson(NemoRelayStatus::InvalidArg), ] { let (raw, _, drop_count) = test_native_stream([item]); let mut stream = native_stream_to_relay_stream(raw, None, None).unwrap(); @@ -3717,6 +5707,20 @@ async fn native_stream_adapter_covers_chunks_end_errors_and_cancellation() { raw.next = None; 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(), + finished: false, + _next_ctx: None, + _callback_user_data: None, + }; + assert!(invalid.next().await.unwrap().is_err()); + assert!(invalid.next().await.is_none()); } #[tokio::test] @@ -3751,6 +5755,10 @@ async fn relay_stream_adapter_covers_poll_end_error_and_cancel() { unsafe { poll(raw.user_data, &mut out) }, NemoRelayStatus::StreamEnd ); + assert_eq!( + unsafe { poll(raw.user_data, &mut out) }, + NemoRelayStatus::StreamEnd + ); assert_eq!( unsafe { cancel_relay_llm_stream(raw.user_data) }, NemoRelayStatus::Ok @@ -3775,6 +5783,10 @@ async fn relay_stream_adapter_covers_poll_end_error_and_cancel() { let _guard = mutex.lock().unwrap(); panic!("poison native stream lock"); })); + assert_eq!( + unsafe { raw.next.unwrap()(raw.user_data, &mut out) }, + NemoRelayStatus::Internal + ); assert_eq!( unsafe { cancel_relay_llm_stream(raw.user_data) }, NemoRelayStatus::Internal @@ -3823,6 +5835,16 @@ fn native_stream_continuation_covers_success_and_error() { unsafe { native_llm_stream_next(request_json, next_ctx, &mut raw) }, NemoRelayStatus::NotFound ); + unsafe { drop(Box::from_raw(next_ctx as *mut LlmStreamExecutionNextFn)) }; + + let next: LlmStreamExecutionNextFn = + Arc::new(|_| Box::pin(async { panic!("stream next panic") })); + let next_ctx = Box::into_raw(Box::new(next)) as *mut c_void; + assert_eq!( + unsafe { native_llm_stream_next(request_json, next_ctx, &mut raw) }, + NemoRelayStatus::Internal + ); + assert_last_error_contains("native LLM stream next panicked"); unsafe { drop(Box::from_raw(next_ctx as *mut LlmStreamExecutionNextFn)); native_string_free(request_json); diff --git a/crates/node/adaptive.js b/crates/node/adaptive.js index 836e60ab4..c96082012 100644 --- a/crates/node/adaptive.js +++ b/crates/node/adaptive.js @@ -157,7 +157,7 @@ function acgConfig(config = {}) { * @param {object} [config={}] - Partial response-cache settings to override. * @returns {object} A normalized response-cache config object. * @remarks The default backend is in-memory; pass a `backend` (e.g. - * `redisBackend(url)`) for a shared cache. `bypassRate` defaults to `0.0`, + * `redisBackend(url)`) for a shared cache. `bypassRate` defaults to `0`, * while caching nondeterministic requests is opt-in. Set a non-empty * `namespace` identifying one trusted cache-sharing domain before validation; * the empty helper default is an unconfigured sentinel. @@ -168,7 +168,7 @@ function responseCacheConfig(config = {}) { ttlSeconds: 3600, namespace: '', priority: 50, - bypassRate: 0.0, + bypassRate: 0, cacheNondeterministic: false, keyStrategy: 'exact_request', headerAllowlist: [], diff --git a/python/plugin/build_backend.py b/python/plugin/build_backend.py index 2200d104b..8a7481ccf 100644 --- a/python/plugin/build_backend.py +++ b/python/plugin/build_backend.py @@ -160,7 +160,8 @@ def _sdist_manifest() -> Iterator[None]: if previous is None: _MANIFEST.unlink(missing_ok=True) else: - _MANIFEST.write_bytes(previous) + # `_MANIFEST` is resolved and constrained to a direct child of this backend's directory. + _MANIFEST.write_bytes(previous) # NOSONAR def _setuptools_backend() -> Any: