#![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)]
mod async_sdk;
pub use async_sdk::{LlmJsonAsyncStream, LlmNext, LlmStreamNext, NativeExecutorConfig, ToolNext};
use std::ffi::{c_char, c_void};
use std::marker::{PhantomData, PhantomPinned};
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::ptr;
use std::sync::{Arc, Mutex};
pub use nemo_relay_types::Json;
pub use nemo_relay_types::api::event::{
CategoryProfile, DataSchema, Event, EventCategory, EventSanitizeFields, LogSeverity,
METRIC_DATA_SCHEMA_NAME, METRIC_DATA_SCHEMA_VERSION, MetricEnvelope, MetricKind,
MetricMeasurement, MetricValueType, PendingMarkSpec, ScopeCategory,
};
pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome};
pub use nemo_relay_types::api::registry::{
RuntimeRegistrationIdentity, RuntimeRegistrationKind, RuntimeRegistrationOwner,
RuntimeRegistrationOwnerKind,
};
pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType};
pub use nemo_relay_types::api::tool::{
TOOL_EXECUTION_INTERCEPT_OUTCOME_SCHEMA, TOOL_EXECUTION_RESULT_SCHEMA, ToolAttributes,
ToolExecutionInterceptOutcome, ToolExecutionResult,
};
pub use nemo_relay_types::codec::identity::{BuiltinLlmCodec, LlmCodecIdentity};
pub use nemo_relay_types::codec::optimization::{
LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind,
LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload,
LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact,
LlmOptimizationTokens,
};
pub use nemo_relay_types::codec::request::AnnotatedLlmRequest;
pub use nemo_relay_types::codec::response::AnnotatedLlmResponse;
pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Map;
pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 4;
pub const NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE: u32 = 3;
pub const NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY: u32 = 2;
pub struct LlmSanitizeRequestContext<'a> {
pub codec: LlmCodecIdentity,
resolved: Option<LlmSanitizeRequestCodec<'a>>,
}
unsafe impl Send for LlmSanitizeRequestContext<'_> {}
pub struct LlmSanitizeResponseContext<'a> {
pub codec: LlmCodecIdentity,
resolved: Option<LlmSanitizeResponseCodec<'a>>,
}
unsafe impl Send for LlmSanitizeResponseContext<'_> {}
#[repr(i32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NemoRelayStatus {
Ok = 0,
AlreadyExists = 1,
NotFound = 2,
ScopeStackEmpty = 3,
GuardrailRejected = 4,
Internal = 5,
NullPointer = 6,
InvalidJson = 7,
InvalidUtf8 = 8,
InvalidArg = 9,
StreamEnd = 10,
Backpressured = 11,
}
#[repr(C)]
pub struct NemoRelayNativeString {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeLlmRequestCodec {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeLlmResponseCodec {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NemoRelayNativeLlmCodecKind {
None = 0,
BuiltIn = 1,
Runtime = 2,
Opaque = 3,
}
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct NemoRelayNativeLlmSanitizeRequestContext {
pub codec_kind: NemoRelayNativeLlmCodecKind,
pub codec_id: *const NemoRelayNativeString,
pub codec: *const NemoRelayNativeLlmRequestCodec,
}
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct NemoRelayNativeLlmSanitizeResponseContext {
pub codec_kind: NemoRelayNativeLlmCodecKind,
pub codec_id: *const NemoRelayNativeString,
pub codec: *const NemoRelayNativeLlmResponseCodec,
}
pub struct LlmSanitizeRequestCodec<'a> {
async_host: NemoRelayNativeHostApiV4,
completion: *const NemoRelayNativeAsyncCompletion,
completion_release: unsafe extern "C" fn(*const NemoRelayNativeAsyncCompletion),
_lifetime: PhantomData<&'a NemoRelayNativeLlmRequestCodec>,
}
unsafe impl Send for LlmSanitizeRequestCodec<'_> {}
unsafe impl Sync for LlmSanitizeRequestCodec<'_> {}
impl Drop for LlmSanitizeRequestCodec<'_> {
fn drop(&mut self) {
unsafe { (self.completion_release)(self.completion) };
}
}
impl LlmSanitizeRequestCodec<'_> {
pub fn decode(&self, request: &LlmRequest) -> Result<AnnotatedLlmRequest> {
native_codec_call(&self.async_host.v3.v1, |out| unsafe {
let request = HostString::from_json(&self.async_host.v3.v1, request)
.ok_or_else(|| "failed to serialize LLM request".to_string())?;
let status = (self.async_host.async_completion_llm_request_codec_decode)(
self.completion,
request.as_ptr(),
out,
);
codec_status(&self.async_host.v3.v1, status)
})
}
pub fn encode(
&self,
annotated: &AnnotatedLlmRequest,
original: &LlmRequest,
) -> Result<LlmRequest> {
native_codec_call(&self.async_host.v3.v1, |out| unsafe {
let annotated = HostString::from_json(&self.async_host.v3.v1, annotated)
.ok_or_else(|| "failed to serialize annotated request".to_string())?;
let original = HostString::from_json(&self.async_host.v3.v1, original)
.ok_or_else(|| "failed to serialize original request".to_string())?;
let status = (self.async_host.async_completion_llm_request_codec_encode)(
self.completion,
annotated.as_ptr(),
original.as_ptr(),
out,
);
codec_status(&self.async_host.v3.v1, status)
})
}
}
pub struct LlmSanitizeResponseCodec<'a> {
async_host: NemoRelayNativeHostApiV4,
completion: *const NemoRelayNativeAsyncCompletion,
completion_release: unsafe extern "C" fn(*const NemoRelayNativeAsyncCompletion),
_lifetime: PhantomData<&'a NemoRelayNativeLlmResponseCodec>,
}
unsafe impl Send for LlmSanitizeResponseCodec<'_> {}
unsafe impl Sync for LlmSanitizeResponseCodec<'_> {}
impl Drop for LlmSanitizeResponseCodec<'_> {
fn drop(&mut self) {
unsafe { (self.completion_release)(self.completion) };
}
}
impl LlmSanitizeResponseCodec<'_> {
pub fn decode(&self, response: &Json) -> Result<AnnotatedLlmResponse> {
native_codec_call(&self.async_host.v3.v1, |out| unsafe {
let response = HostString::from_json(&self.async_host.v3.v1, response)
.ok_or_else(|| "failed to serialize LLM response".to_string())?;
let status = (self.async_host.async_completion_llm_response_codec_decode)(
self.completion,
response.as_ptr(),
out,
);
codec_status(&self.async_host.v3.v1, status)
})
}
}
impl<'a> LlmSanitizeRequestContext<'a> {
#[must_use]
pub fn resolve_codec(&self) -> Option<&LlmSanitizeRequestCodec<'a>> {
self.resolved.as_ref()
}
}
impl<'a> LlmSanitizeResponseContext<'a> {
#[must_use]
pub fn resolve_codec(&self) -> Option<&LlmSanitizeResponseCodec<'a>> {
self.resolved.as_ref()
}
}
#[repr(C)]
pub struct NemoRelayNativePluginContext {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativePluginRuntime {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeScopeHandle {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeScopeStack {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeScopeStackBinding {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(i32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NemoRelayNativeScopeType {
Agent = 0,
Function = 1,
Tool = 2,
Llm = 3,
Retriever = 4,
Embedder = 5,
Reranker = 6,
Guardrail = 7,
Evaluator = 8,
Custom = 9,
Unknown = 10,
}
pub type NemoRelayNativeFreeFn = Option<unsafe extern "C" fn(user_data: *mut c_void)>;
pub type NemoRelayNativeWithScopeStackCb =
unsafe extern "C" fn(user_data: *mut c_void) -> NemoRelayStatus;
pub type NemoRelayNativeToolNextFn = unsafe extern "C" fn(
args_json: *const NemoRelayNativeString,
next_ctx: *mut c_void,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmNextFn = unsafe extern "C" fn(
request_json: *const NemoRelayNativeString,
next_ctx: *mut c_void,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmStreamPollFn = unsafe extern "C" fn(
user_data: *mut c_void,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmStreamCancelFn =
Option<unsafe extern "C" fn(user_data: *mut c_void) -> NemoRelayStatus>;
pub type NemoRelayNativeLlmStreamDropFn = Option<unsafe extern "C" fn(user_data: *mut c_void)>;
#[repr(C)]
pub struct NemoRelayNativeLlmStreamV1 {
pub struct_size: usize,
pub user_data: *mut c_void,
pub next: Option<NemoRelayNativeLlmStreamPollFn>,
pub cancel: NemoRelayNativeLlmStreamCancelFn,
pub drop: NemoRelayNativeLlmStreamDropFn,
}
impl Default for NemoRelayNativeLlmStreamV1 {
fn default() -> Self {
Self {
struct_size: std::mem::size_of::<Self>(),
user_data: ptr::null_mut(),
next: None,
cancel: None,
drop: None,
}
}
}
pub type NemoRelayNativeLlmStreamNextFn = unsafe extern "C" fn(
request_json: *const NemoRelayNativeString,
next_ctx: *mut c_void,
out_stream: *mut NemoRelayNativeLlmStreamV1,
) -> NemoRelayStatus;
pub type NemoRelayNativeEventSubscriberCb = unsafe extern "C" fn(
user_data: *mut c_void,
event_json: *const NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeEventSanitizeCb = unsafe extern "C" fn(
user_data: *mut c_void,
event_json: *const NemoRelayNativeString,
fields_json: *const NemoRelayNativeString,
out_fields_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeToolJsonCb = unsafe extern "C" fn(
user_data: *mut c_void,
name: *const NemoRelayNativeString,
payload_json: *const NemoRelayNativeString,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeToolConditionalCb = unsafe extern "C" fn(
user_data: *mut c_void,
name: *const NemoRelayNativeString,
args_json: *const NemoRelayNativeString,
out_reason: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeToolExecutionCb = unsafe extern "C" fn(
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;
pub type NemoRelayNativeLlmSanitizeRequestCb = unsafe extern "C" fn(
user_data: *mut c_void,
request_json: *const NemoRelayNativeString,
context: NemoRelayNativeLlmSanitizeRequestContext,
out_request_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmSanitizeResponseCb = unsafe extern "C" fn(
user_data: *mut c_void,
payload_json: *const NemoRelayNativeString,
context: NemoRelayNativeLlmSanitizeResponseContext,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmConditionalCb = unsafe extern "C" fn(
user_data: *mut c_void,
request_json: *const NemoRelayNativeString,
out_reason: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmRequestInterceptCb = unsafe extern "C" fn(
user_data: *mut c_void,
name: *const NemoRelayNativeString,
request_json: *const NemoRelayNativeString,
annotated_json: *const NemoRelayNativeString,
out_outcome_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativeLlmExecutionCb = unsafe extern "C" fn(
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;
pub type NemoRelayNativeLlmStreamExecutionCb = unsafe extern "C" fn(
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;
pub type NemoRelayNativePluginValidateFn = unsafe extern "C" fn(
user_data: *mut c_void,
plugin_config_json: *const NemoRelayNativeString,
out_diagnostics_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus;
pub type NemoRelayNativePluginRegisterFn = unsafe extern "C" fn(
user_data: *mut c_void,
plugin_config_json: *const NemoRelayNativeString,
ctx: *mut NemoRelayNativePluginContext,
) -> NemoRelayStatus;
pub type NemoRelayNativePluginDropFn = Option<unsafe extern "C" fn(user_data: *mut c_void)>;
#[repr(C)]
#[derive(Clone, Copy)]
pub struct NemoRelayNativeHostApiV1 {
pub abi_version: u32,
pub struct_size: usize,
pub relay_version: *const c_char,
pub string_new: unsafe extern "C" fn(
data: *const u8,
len: usize,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub string_data: unsafe extern "C" fn(value: *const NemoRelayNativeString) -> *const u8,
pub string_len: unsafe extern "C" fn(value: *const NemoRelayNativeString) -> usize,
pub string_free: unsafe extern "C" fn(value: *mut NemoRelayNativeString),
pub last_error_clear: unsafe extern "C" fn(),
pub last_error_set: unsafe extern "C" fn(message: *const NemoRelayNativeString),
pub llm_request_codec_decode: unsafe extern "C" fn(
codec: *const NemoRelayNativeLlmRequestCodec,
request_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub llm_request_codec_encode: unsafe extern "C" fn(
codec: *const NemoRelayNativeLlmRequestCodec,
annotated_json: *const NemoRelayNativeString,
original_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub llm_response_codec_decode: unsafe extern "C" fn(
codec: *const NemoRelayNativeLlmResponseCodec,
response_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub plugin_context_register_subscriber: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
cb: NemoRelayNativeEventSubscriberCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_tool_sanitize_request_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_tool_sanitize_response_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_tool_conditional_execution_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeToolConditionalCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_tool_request_intercept: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
)
-> NemoRelayStatus,
pub plugin_context_register_tool_execution_intercept: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeToolExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
)
-> NemoRelayStatus,
pub plugin_context_register_llm_sanitize_request_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeLlmSanitizeRequestCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_llm_sanitize_response_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeLlmSanitizeResponseCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_llm_conditional_execution_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeLlmConditionalCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_llm_request_intercept: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeLlmRequestInterceptCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_llm_execution_intercept: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeLlmExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
)
-> NemoRelayStatus,
pub plugin_context_register_llm_stream_execution_intercept:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeLlmStreamExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub scope_handle_free: unsafe extern "C" fn(handle: *mut NemoRelayNativeScopeHandle),
pub scope_get_current:
unsafe extern "C" fn(out: *mut *mut NemoRelayNativeScopeHandle) -> NemoRelayStatus,
pub scope_push: unsafe extern "C" fn(
name: *const NemoRelayNativeString,
scope_type: NemoRelayNativeScopeType,
parent: *const NemoRelayNativeScopeHandle,
attributes: u32,
data_json: *const NemoRelayNativeString,
metadata_json: *const NemoRelayNativeString,
input_json: *const NemoRelayNativeString,
timestamp_unix_micros: *const i64,
out: *mut *mut NemoRelayNativeScopeHandle,
) -> NemoRelayStatus,
pub scope_pop: unsafe extern "C" fn(
handle: *const NemoRelayNativeScopeHandle,
output_json: *const NemoRelayNativeString,
metadata_json: *const NemoRelayNativeString,
timestamp_unix_micros: *const i64,
) -> NemoRelayStatus,
pub emit_mark: unsafe extern "C" fn(
name: *const NemoRelayNativeString,
parent: *const NemoRelayNativeScopeHandle,
data_json: *const NemoRelayNativeString,
metadata_json: *const NemoRelayNativeString,
timestamp_unix_micros: *const i64,
) -> NemoRelayStatus,
pub scope_stack_create:
unsafe extern "C" fn(out: *mut *mut NemoRelayNativeScopeStack) -> NemoRelayStatus,
pub scope_stack_free: unsafe extern "C" fn(stack: *mut NemoRelayNativeScopeStack),
pub scope_stack_set_thread:
unsafe extern "C" fn(stack: *const NemoRelayNativeScopeStack) -> NemoRelayStatus,
pub scope_stack_capture_thread:
unsafe extern "C" fn(out: *mut *mut NemoRelayNativeScopeStackBinding) -> NemoRelayStatus,
pub scope_stack_restore_thread:
unsafe extern "C" fn(binding: *mut NemoRelayNativeScopeStackBinding) -> NemoRelayStatus,
pub scope_stack_binding_free:
unsafe extern "C" fn(binding: *mut NemoRelayNativeScopeStackBinding),
pub scope_stack_active: unsafe extern "C" fn() -> bool,
pub scope_stack_with_current: unsafe extern "C" fn(
stack: *const NemoRelayNativeScopeStack,
cb: NemoRelayNativeWithScopeStackCb,
user_data: *mut c_void,
) -> NemoRelayStatus,
pub plugin_context_register_mark_sanitize_guardrail: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
)
-> NemoRelayStatus,
pub plugin_context_register_scope_sanitize_start_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub plugin_context_register_scope_sanitize_end_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
}
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NemoRelayNativeAsyncMiddlewareKind {
ToolSanitizeRequest = 0,
ToolSanitizeResponse = 1,
ToolConditionalExecution = 2,
ToolRequestIntercept = 3,
ToolExecutionIntercept = 4,
LlmSanitizeRequest = 5,
LlmSanitizeResponse = 6,
LlmConditionalExecution = 7,
LlmRequestIntercept = 8,
LlmExecutionIntercept = 9,
LlmStreamExecutionIntercept = 10,
MarkSanitize = 11,
ScopeSanitizeStart = 12,
ScopeSanitizeEnd = 13,
EventMetadataInjector = 14,
}
impl TryFrom<u32> for NemoRelayNativeAsyncMiddlewareKind {
type Error = ();
fn try_from(value: u32) -> std::result::Result<Self, Self::Error> {
match value {
0 => Ok(Self::ToolSanitizeRequest),
1 => Ok(Self::ToolSanitizeResponse),
2 => Ok(Self::ToolConditionalExecution),
3 => Ok(Self::ToolRequestIntercept),
4 => Ok(Self::ToolExecutionIntercept),
5 => Ok(Self::LlmSanitizeRequest),
6 => Ok(Self::LlmSanitizeResponse),
7 => Ok(Self::LlmConditionalExecution),
8 => Ok(Self::LlmRequestIntercept),
9 => Ok(Self::LlmExecutionIntercept),
10 => Ok(Self::LlmStreamExecutionIntercept),
11 => Ok(Self::MarkSanitize),
12 => Ok(Self::ScopeSanitizeStart),
13 => Ok(Self::ScopeSanitizeEnd),
14 => Ok(Self::EventMetadataInjector),
_ => Err(()),
}
}
}
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NemoRelayNativeAsyncCallbackState {
Complete = 0,
Pending = 1,
}
impl TryFrom<u32> for NemoRelayNativeAsyncCallbackState {
type Error = ();
fn try_from(value: u32) -> std::result::Result<Self, Self::Error> {
match value {
0 => Ok(Self::Complete),
1 => Ok(Self::Pending),
_ => Err(()),
}
}
}
#[repr(C)]
pub struct NemoRelayNativeAsyncCompletion {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeAsyncNext {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeAsyncStream {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
#[repr(C)]
pub struct NemoRelayNativeLlmAsyncStream {
_private: [u8; 0],
_marker: PhantomData<(*mut u8, PhantomPinned)>,
}
pub type NemoRelayNativeAsyncLlmStreamOpenCb = unsafe extern "C" fn(
user_data: *mut c_void,
stream: *const NemoRelayNativeLlmAsyncStream,
error: *const NemoRelayNativeString,
);
pub type NemoRelayNativeAsyncLlmStreamPullCb = unsafe extern "C" fn(
user_data: *mut c_void,
chunk_json: *const NemoRelayNativeString,
error: *const NemoRelayNativeString,
done: bool,
);
pub type NemoRelayNativeAsyncNextStreamCb = unsafe extern "C" fn(
user_data: *mut c_void,
chunk_json: *const NemoRelayNativeString,
error: *const NemoRelayNativeString,
done: bool,
) -> bool;
pub type NemoRelayNativeAsyncNextResultCb = unsafe extern "C" fn(
user_data: *mut c_void,
value_json: *const NemoRelayNativeString,
error: *const NemoRelayNativeString,
);
pub type NemoRelayNativeAsyncStreamMiddlewareCb = unsafe extern "C" fn(
user_data: *mut c_void,
invocation_json: *const NemoRelayNativeString,
next: *const NemoRelayNativeAsyncNext,
stream: *const NemoRelayNativeAsyncStream,
) -> u32;
pub type NemoRelayNativeAsyncMiddlewareCb = unsafe extern "C" fn(
user_data: *mut c_void,
invocation_json: *const NemoRelayNativeString,
next: *const NemoRelayNativeAsyncNext,
completion: *const NemoRelayNativeAsyncCompletion,
) -> u32;
#[repr(C)]
#[derive(Clone, Copy)]
pub struct NemoRelayNativeHostApiV3 {
pub v1: NemoRelayNativeHostApiV1,
pub async_completion_resolve_json: unsafe extern "C" fn(
completion: *const NemoRelayNativeAsyncCompletion,
value_json: *const NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_completion_reject: unsafe extern "C" fn(
completion: *const NemoRelayNativeAsyncCompletion,
message: *const NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_completion_is_cancelled:
unsafe extern "C" fn(completion: *const NemoRelayNativeAsyncCompletion) -> bool,
pub async_completion_release:
unsafe extern "C" fn(completion: *const NemoRelayNativeAsyncCompletion),
pub async_next_invoke: unsafe extern "C" fn(
next: *const NemoRelayNativeAsyncNext,
invocation_json: *const NemoRelayNativeString,
completion: *const NemoRelayNativeAsyncCompletion,
) -> NemoRelayStatus,
pub async_next_release: unsafe extern "C" fn(next: *const NemoRelayNativeAsyncNext),
pub plugin_context_register_async_middleware: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
kind: u32,
name: *const NemoRelayNativeString,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeAsyncMiddlewareCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus,
pub async_stream_push_json: unsafe extern "C" fn(
stream: *const NemoRelayNativeAsyncStream,
chunk_json: *const NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_stream_finish:
unsafe extern "C" fn(stream: *const NemoRelayNativeAsyncStream) -> NemoRelayStatus,
pub async_stream_reject: unsafe extern "C" fn(
stream: *const NemoRelayNativeAsyncStream,
message: *const NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_stream_is_cancelled:
unsafe extern "C" fn(stream: *const NemoRelayNativeAsyncStream) -> bool,
pub async_stream_release: unsafe extern "C" fn(stream: *const NemoRelayNativeAsyncStream),
pub async_next_invoke_stream: unsafe extern "C" fn(
next: *const NemoRelayNativeAsyncNext,
invocation_json: *const NemoRelayNativeString,
stream: *const NemoRelayNativeAsyncStream,
cb: NemoRelayNativeAsyncNextStreamCb,
user_data: *mut c_void,
) -> NemoRelayStatus,
pub plugin_context_register_async_stream_middleware: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
priority: i32,
cb: NemoRelayNativeAsyncStreamMiddlewareCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
)
-> NemoRelayStatus,
pub async_next_invoke_result: unsafe extern "C" fn(
next: *const NemoRelayNativeAsyncNext,
invocation_json: *const NemoRelayNativeString,
cb: NemoRelayNativeAsyncNextResultCb,
user_data: *mut c_void,
) -> NemoRelayStatus,
}
pub type NemoRelayNativeEmitMarkV2Fn = unsafe extern "C" fn(
name: *const NemoRelayNativeString,
parent: *const NemoRelayNativeScopeHandle,
data_json: *const NemoRelayNativeString,
metadata_json: *const NemoRelayNativeString,
data_schema_json: *const NemoRelayNativeString,
severity: *const NemoRelayNativeString,
timestamp_unix_micros: *const i64,
) -> NemoRelayStatus;
pub type NemoRelayNativeGetRuntimeDiagnosticsFn =
unsafe extern "C" fn(out_json: *mut *mut NemoRelayNativeString) -> NemoRelayStatus;
#[repr(C)]
#[derive(Clone, Copy)]
pub struct NemoRelayNativeHostApiV4 {
pub v3: NemoRelayNativeHostApiV3,
pub async_completion_llm_request_codec_decode: unsafe extern "C" fn(
completion: *const NemoRelayNativeAsyncCompletion,
request_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_completion_llm_request_codec_encode: unsafe extern "C" fn(
completion: *const NemoRelayNativeAsyncCompletion,
annotated_json: *const NemoRelayNativeString,
original_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_completion_llm_response_codec_decode: unsafe extern "C" fn(
completion: *const NemoRelayNativeAsyncCompletion,
response_json: *const NemoRelayNativeString,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub async_next_open_llm_stream: unsafe extern "C" fn(
next: *const NemoRelayNativeAsyncNext,
request_json: *const NemoRelayNativeString,
cb: NemoRelayNativeAsyncLlmStreamOpenCb,
user_data: *mut c_void,
) -> NemoRelayStatus,
pub async_llm_stream_pull: unsafe extern "C" fn(
stream: *const NemoRelayNativeLlmAsyncStream,
cb: NemoRelayNativeAsyncLlmStreamPullCb,
user_data: *mut c_void,
) -> NemoRelayStatus,
pub async_llm_stream_cancel:
unsafe extern "C" fn(stream: *const NemoRelayNativeLlmAsyncStream) -> NemoRelayStatus,
pub async_llm_stream_release:
unsafe extern "C" fn(stream: *const NemoRelayNativeLlmAsyncStream),
pub async_completion_retain:
unsafe extern "C" fn(completion: *const NemoRelayNativeAsyncCompletion) -> NemoRelayStatus,
pub async_stream_is_backpressured:
unsafe extern "C" fn(stream: *const NemoRelayNativeAsyncStream) -> bool,
pub emit_mark_v2: NemoRelayNativeEmitMarkV2Fn,
pub get_runtime_diagnostics: NemoRelayNativeGetRuntimeDiagnosticsFn,
pub plugin_context_runtime: unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
out: *mut *const NemoRelayNativePluginRuntime,
) -> NemoRelayStatus,
pub plugin_runtime_retain:
unsafe extern "C" fn(runtime: *const NemoRelayNativePluginRuntime) -> NemoRelayStatus,
pub plugin_runtime_release: unsafe extern "C" fn(runtime: *const NemoRelayNativePluginRuntime),
pub plugin_runtime_list_registrations: unsafe extern "C" fn(
runtime: *const NemoRelayNativePluginRuntime,
kinds_json: *const NemoRelayNativeString,
out_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub plugin_runtime_register_conditional_middleware_guardrail:
unsafe extern "C" fn(
runtime: *const NemoRelayNativePluginRuntime,
name: *const NemoRelayNativeString,
kinds_json: *const NemoRelayNativeString,
registration_name: *const NemoRelayNativeString,
reason: *const NemoRelayNativeString,
out_handle: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus,
pub plugin_runtime_deregister_conditional_middleware_guardrail:
unsafe extern "C" fn(
runtime: *const NemoRelayNativePluginRuntime,
handle: *const NemoRelayNativeString,
out_removed: *mut bool,
) -> NemoRelayStatus,
pub plugin_context_register_conditional_middleware_guardrail:
unsafe extern "C" fn(
ctx: *mut NemoRelayNativePluginContext,
name: *const NemoRelayNativeString,
kinds_json: *const NemoRelayNativeString,
registration_name: *const NemoRelayNativeString,
reason: *const NemoRelayNativeString,
) -> NemoRelayStatus,
}
unsafe impl Send for NemoRelayNativeHostApiV3 {}
unsafe impl Sync for NemoRelayNativeHostApiV3 {}
unsafe impl Send for NemoRelayNativeHostApiV4 {}
unsafe impl Sync for NemoRelayNativeHostApiV4 {}
unsafe impl Send for NemoRelayNativeHostApiV1 {}
unsafe impl Sync for NemoRelayNativeHostApiV1 {}
#[repr(C)]
pub struct NemoRelayNativePluginV1 {
pub struct_size: usize,
pub plugin_kind: *mut NemoRelayNativeString,
pub allows_multiple_components: bool,
pub user_data: *mut c_void,
pub validate: Option<NemoRelayNativePluginValidateFn>,
pub register: Option<NemoRelayNativePluginRegisterFn>,
pub drop: NemoRelayNativePluginDropFn,
}
impl Default for NemoRelayNativePluginV1 {
fn default() -> Self {
Self {
struct_size: std::mem::size_of::<Self>(),
plugin_kind: ptr::null_mut(),
allows_multiple_components: true,
user_data: ptr::null_mut(),
validate: None,
register: None,
drop: None,
}
}
}
pub type NemoRelayNativePluginEntry = unsafe extern "C" fn(
host: *const NemoRelayNativeHostApiV1,
out: *mut NemoRelayNativePluginV1,
) -> NemoRelayStatus;
pub type Result<T> = std::result::Result<T, String>;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)]
pub struct RuntimeDiagnostic {
pub code: String,
pub message: String,
pub count: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
pub struct RuntimeDiagnostics {
entries: Vec<RuntimeDiagnostic>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConditionalMiddlewareGuardrailHandle(String);
impl RuntimeDiagnostics {
pub fn entries(&self) -> &[RuntimeDiagnostic] {
&self.entries
}
pub fn get(&self, code: &str) -> Option<&RuntimeDiagnostic> {
self.entries
.iter()
.find(|diagnostic| diagnostic.code == code)
}
}
pub type LlmJsonStream = Box<dyn Iterator<Item = Result<Json>> + Send>;
pub struct PluginRuntime {
host: NemoRelayNativeHostApiV1,
emit_mark_v2: Option<NemoRelayNativeEmitMarkV2Fn>,
get_runtime_diagnostics: Option<NemoRelayNativeGetRuntimeDiagnosticsFn>,
v4: Option<NemoRelayNativeHostApiV4>,
capability: *const NemoRelayNativePluginRuntime,
}
unsafe impl Send for PluginRuntime {}
unsafe impl Sync for PluginRuntime {}
impl Clone for PluginRuntime {
fn clone(&self) -> Self {
let mut capability = self.capability;
if let (Some(v4), false) = (self.v4, self.capability.is_null())
&& unsafe { (v4.plugin_runtime_retain)(self.capability) } != NemoRelayStatus::Ok
{
capability = ptr::null();
}
Self {
host: self.host,
emit_mark_v2: self.emit_mark_v2,
get_runtime_diagnostics: self.get_runtime_diagnostics,
v4: self.v4,
capability,
}
}
}
impl Drop for PluginRuntime {
fn drop(&mut self) {
if let (Some(v4), false) = (self.v4, self.capability.is_null()) {
unsafe { (v4.plugin_runtime_release)(self.capability) };
}
}
}
impl PluginRuntime {
pub fn new(host: &NemoRelayNativeHostApiV1) -> Self {
let v4 = (host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION
&& host.struct_size >= std::mem::size_of::<NemoRelayNativeHostApiV4>())
.then(|| unsafe { *(host as *const _ as *const NemoRelayNativeHostApiV4) });
Self {
host: *host,
emit_mark_v2: v4.map(|host| host.emit_mark_v2),
get_runtime_diagnostics: v4.map(|host| host.get_runtime_diagnostics),
v4,
capability: ptr::null(),
}
}
fn from_context(
host: &NemoRelayNativeHostApiV1,
ctx: *mut NemoRelayNativePluginContext,
) -> Self {
let mut runtime = Self::new(host);
let Some(v4) = runtime.v4 else {
return runtime;
};
let mut capability = ptr::null();
if unsafe { (v4.plugin_context_runtime)(ctx, &mut capability) } == NemoRelayStatus::Ok {
runtime.capability = capability;
}
runtime
}
pub fn list_runtime_registrations(
&self,
kinds: Option<&std::collections::BTreeSet<RuntimeRegistrationKind>>,
) -> Result<Vec<RuntimeRegistrationIdentity>> {
let v4 = self.runtime_v4()?;
let kinds = match kinds {
Some(kinds) => Some(
HostString::from_json(&self.host, kinds)
.ok_or_else(|| "failed to serialize runtime registration kinds".to_string())?,
),
None => None,
};
let mut out = ptr::null_mut();
let status = unsafe {
(v4.plugin_runtime_list_registrations)(
self.capability,
kinds.as_ref().map_or(ptr::null(), HostString::as_ptr),
&mut out,
)
};
status_result(&self.host, status, "list runtime registrations")?;
take_host_json(&self.host, out)
}
pub fn register_conditional_middleware_guardrail(
&self,
name: &str,
kinds: &std::collections::BTreeSet<RuntimeRegistrationKind>,
registration_name: &str,
reason: &str,
) -> Result<ConditionalMiddlewareGuardrailHandle> {
let v4 = self.runtime_v4()?;
let name = HostString::new(&self.host, name)
.ok_or_else(|| "failed to allocate gate name".to_string())?;
let kinds = HostString::from_json(&self.host, kinds)
.ok_or_else(|| "failed to serialize runtime registration kinds".to_string())?;
let registration_name = HostString::new(&self.host, registration_name)
.ok_or_else(|| "failed to allocate target name".to_string())?;
let reason = HostString::new(&self.host, reason)
.ok_or_else(|| "failed to allocate gate reason".to_string())?;
let mut out = ptr::null_mut();
let status = unsafe {
(v4.plugin_runtime_register_conditional_middleware_guardrail)(
self.capability,
name.as_ptr(),
kinds.as_ptr(),
registration_name.as_ptr(),
reason.as_ptr(),
&mut out,
)
};
status_result(
&self.host,
status,
"register conditional middleware guardrail",
)?;
take_host_string(&self.host, out).map(ConditionalMiddlewareGuardrailHandle)
}
pub fn deregister_conditional_middleware_guardrail(
&self,
handle: &ConditionalMiddlewareGuardrailHandle,
) -> Result<bool> {
let v4 = self.runtime_v4()?;
let handle = HostString::new(&self.host, &handle.0)
.ok_or_else(|| "failed to allocate gate handle".to_string())?;
let mut removed = false;
let status = unsafe {
(v4.plugin_runtime_deregister_conditional_middleware_guardrail)(
self.capability,
handle.as_ptr(),
&mut removed,
)
};
status_result(
&self.host,
status,
"deregister conditional middleware guardrail",
)?;
Ok(removed)
}
fn runtime_v4(&self) -> Result<NemoRelayNativeHostApiV4> {
self.v4
.filter(|_| !self.capability.is_null())
.ok_or_else(|| "host does not support activation-owned runtime gate control".into())
}
pub fn host_api(&self) -> &NemoRelayNativeHostApiV1 {
&self.host
}
pub fn current_scope(&self) -> Result<ScopeHandle<'_>> {
current_scope(&self.host)
}
pub fn push_scope(
&self,
name: &str,
scope_type: ScopeType,
data: Option<&Json>,
metadata: Option<&Json>,
input: Option<&Json>,
) -> Result<ScopeHandle<'_>> {
push_scope(&self.host, name, scope_type.into(), data, metadata, input)
}
pub fn pop_scope(
&self,
handle: &ScopeHandle<'_>,
output: Option<&Json>,
metadata: Option<&Json>,
) -> Result<()> {
pop_scope(&self.host, handle, output, metadata)
}
pub fn scope(
&self,
name: &str,
scope_type: ScopeType,
data: Option<&Json>,
metadata: Option<&Json>,
input: Option<&Json>,
) -> Result<ScopeGuard<'_>> {
let handle = self.push_scope(name, scope_type, data, metadata, input)?;
Ok(ScopeGuard {
runtime: self,
handle: Some(handle),
})
}
pub fn emit_mark(
&self,
name: &str,
data: Option<&Json>,
metadata: Option<&Json>,
) -> Result<()> {
emit_mark(&self.host, name, data, metadata)
}
pub fn emit_mark_with_options(
&self,
name: &str,
data: Option<&Json>,
metadata: Option<&Json>,
data_schema: Option<&DataSchema>,
severity: Option<LogSeverity>,
) -> Result<()> {
match self.emit_mark_v2 {
Some(emit_mark_v2) => emit_mark_v2_call(
&self.host,
emit_mark_v2,
name,
data,
metadata,
data_schema,
severity,
),
None if data_schema.is_none() && severity.is_none() => {
emit_mark(&self.host, name, data, metadata)
}
None => Err("mark data_schema and severity require native host ABI v4".into()),
}
}
pub fn emit_metric(
&self,
name: &str,
measurements: Vec<MetricMeasurement>,
metadata: Option<&Json>,
) -> Result<()> {
let envelope = MetricEnvelope { measurements };
envelope.validate().map_err(|err| err.to_string())?;
let data = serde_json::to_value(envelope)
.map_err(|err| format!("failed to serialize metric mark: {err}"))?;
let data_schema = DataSchema::builder()
.name(METRIC_DATA_SCHEMA_NAME)
.version(METRIC_DATA_SCHEMA_VERSION)
.build();
self.emit_mark_with_options(name, Some(&data), metadata, Some(&data_schema), None)
}
pub fn runtime_diagnostics(&self) -> Result<RuntimeDiagnostics> {
let Some(get_runtime_diagnostics) = self.get_runtime_diagnostics else {
return Err(
"runtime diagnostics require the native host ABI v4 diagnostics extension".into(),
);
};
native_json_call(&self.host, "runtime diagnostics", |out| {
let status = unsafe { get_runtime_diagnostics(out) };
codec_status(&self.host, status)
})
}
pub fn create_scope_stack(&self) -> Result<ScopeStack<'_>> {
create_scope_stack(&self.host)
}
pub fn capture_scope_stack_thread(&self) -> Result<ScopeStackBinding<'_>> {
capture_scope_stack_thread(&self.host)
}
pub fn scope_stack_active(&self) -> bool {
unsafe { (self.host.scope_stack_active)() }
}
pub fn bind_scope_stack_thread<'a>(
&'a self,
stack: &'a ScopeStack<'a>,
) -> Result<ThreadScopeStackGuard<'a>> {
let previous = self.capture_scope_stack_thread()?;
let status = stack.set_thread();
if status == NemoRelayStatus::Ok {
Ok(ThreadScopeStackGuard {
previous: Some(previous),
})
} else {
let _ = previous.restore();
Err(format!("scope_stack_set_thread failed: {status:?}"))
}
}
}
impl From<ScopeType> for NemoRelayNativeScopeType {
fn from(value: ScopeType) -> Self {
match value {
ScopeType::Agent => Self::Agent,
ScopeType::Function => Self::Function,
ScopeType::Tool => Self::Tool,
ScopeType::Llm => Self::Llm,
ScopeType::Retriever => Self::Retriever,
ScopeType::Embedder => Self::Embedder,
ScopeType::Reranker => Self::Reranker,
ScopeType::Guardrail => Self::Guardrail,
ScopeType::Evaluator => Self::Evaluator,
ScopeType::Custom => Self::Custom,
ScopeType::Unknown => Self::Unknown,
}
}
}
pub struct ScopeGuard<'a> {
runtime: &'a PluginRuntime,
handle: Option<ScopeHandle<'a>>,
}
unsafe impl Send for ScopeGuard<'_> {}
impl<'a> ScopeGuard<'a> {
pub fn handle(&self) -> Option<&ScopeHandle<'a>> {
self.handle.as_ref()
}
pub fn close(&mut self, output: Option<&Json>, metadata: Option<&Json>) -> Result<()> {
let Some(handle) = self.handle.as_ref() else {
return Ok(());
};
self.runtime.pop_scope(handle, output, metadata)?;
self.handle.take();
Ok(())
}
}
impl Drop for ScopeGuard<'_> {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
let _ = self.runtime.pop_scope(&handle, None, None);
}
}
}
pub struct ThreadScopeStackGuard<'a> {
previous: Option<ScopeStackBinding<'a>>,
}
impl ThreadScopeStackGuard<'_> {
pub fn restore(mut self) -> Result<()> {
let Some(previous) = self.previous.take() else {
return Ok(());
};
let status = previous.restore();
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(format!("scope_stack_restore_thread failed: {status:?}"))
}
}
}
impl Drop for ThreadScopeStackGuard<'_> {
fn drop(&mut self) {
if let Some(previous) = self.previous.take() {
let _ = previous.restore();
}
}
}
pub struct LlmStream {
host: NemoRelayNativeHostApiV1,
raw: NemoRelayNativeLlmStreamV1,
finished: bool,
}
unsafe impl Send for LlmStream {}
impl LlmStream {
pub unsafe fn from_raw(
host: &NemoRelayNativeHostApiV1,
mut raw: NemoRelayNativeLlmStreamV1,
) -> Result<Self> {
let expected_size = std::mem::size_of::<NemoRelayNativeLlmStreamV1>();
if raw.struct_size != expected_size {
if raw.struct_size >= expected_size {
unsafe { drop_raw_llm_stream(&mut raw) };
}
return Err(format!(
"unsupported LLM stream struct size: {}",
raw.struct_size
));
}
if raw.next.is_none() {
unsafe { drop_raw_llm_stream(&mut raw) };
return Err("LLM stream next callback was null".into());
}
Ok(Self {
host: *host,
raw,
finished: false,
})
}
pub fn next_chunk(&mut self) -> Result<Option<Json>> {
if self.finished {
return Ok(None);
}
let next = self
.raw
.next
.expect("LLM stream next callback is validated on construction");
let mut out = ptr::null_mut();
let status = unsafe { next(self.raw.user_data, &mut out) };
match status {
NemoRelayStatus::Ok => {
if out.is_null() {
self.finished = true;
return Err("LLM stream returned null chunk".into());
}
let result = read_json_value(&self.host, out, "LLM stream chunk");
unsafe { (self.host.string_free)(out) };
match result {
Ok(chunk) => Ok(Some(chunk)),
Err(status) => {
self.finished = true;
Err(format!("LLM stream returned invalid JSON: {status:?}"))
}
}
}
NemoRelayStatus::StreamEnd => {
if !out.is_null() {
unsafe { (self.host.string_free)(out) };
}
self.finished = true;
Ok(None)
}
other => {
if !out.is_null() {
unsafe { (self.host.string_free)(out) };
}
self.finished = true;
Err(format!("LLM stream failed: {other:?}"))
}
}
}
pub fn cancel(&mut self) -> Result<()> {
if self.finished {
return Ok(());
}
if let Some(cancel) = self.raw.cancel {
let status = unsafe { cancel(self.raw.user_data) };
if status != NemoRelayStatus::Ok {
return Err(format!("LLM stream cancel failed: {status:?}"));
}
}
self.finished = true;
Ok(())
}
}
impl Iterator for LlmStream {
type Item = Result<Json>;
fn next(&mut self) -> Option<Self::Item> {
match self.next_chunk() {
Ok(Some(chunk)) => Some(Ok(chunk)),
Ok(None) => None,
Err(message) => Some(Err(message)),
}
}
}
unsafe fn drop_raw_llm_stream(raw: &mut NemoRelayNativeLlmStreamV1) {
if let Some(drop_fn) = raw.drop.take() {
unsafe { drop_fn(raw.user_data) };
}
raw.user_data = ptr::null_mut();
}
impl Drop for LlmStream {
fn drop(&mut self) {
if !self.finished {
if let Some(cancel) = self.raw.cancel {
let _ = unsafe { cancel(self.raw.user_data) };
}
self.finished = true;
}
unsafe { drop_raw_llm_stream(&mut self.raw) };
}
}
pub struct ScopeHandle<'a> {
host: &'a NemoRelayNativeHostApiV1,
ptr: *mut NemoRelayNativeScopeHandle,
}
unsafe impl Send for ScopeHandle<'_> {}
impl<'a> ScopeHandle<'a> {
pub fn as_ptr(&self) -> *const NemoRelayNativeScopeHandle {
self.ptr
}
}
impl Drop for ScopeHandle<'_> {
fn drop(&mut self) {
unsafe { (self.host.scope_handle_free)(self.ptr) };
}
}
pub struct ScopeStack<'a> {
host: &'a NemoRelayNativeHostApiV1,
ptr: *mut NemoRelayNativeScopeStack,
}
unsafe impl Send for ScopeStack<'_> {}
impl<'a> ScopeStack<'a> {
pub fn as_ptr(&self) -> *const NemoRelayNativeScopeStack {
self.ptr
}
pub fn set_thread(&self) -> NemoRelayStatus {
unsafe { (self.host.scope_stack_set_thread)(self.ptr) }
}
pub fn with_current<F>(&self, f: F) -> Result<()>
where
F: FnOnce() -> Result<()>,
{
struct State<F> {
f: Option<F>,
error: Option<String>,
}
unsafe extern "C" fn trampoline<F>(user_data: *mut c_void) -> NemoRelayStatus
where
F: FnOnce() -> Result<()>,
{
if user_data.is_null() {
return NemoRelayStatus::NullPointer;
}
let state = unsafe { &mut *(user_data as *mut State<F>) };
let result = catch_unwind(AssertUnwindSafe(|| {
let Some(f) = state.f.take() else {
return Err("scope-stack callback was already consumed".to_string());
};
f()
}));
match result {
Ok(Ok(())) => NemoRelayStatus::Ok,
Ok(Err(message)) => {
state.error = Some(message);
NemoRelayStatus::Internal
}
Err(_) => {
state.error = Some("scope-stack callback panicked".into());
NemoRelayStatus::Internal
}
}
}
let mut state = State {
f: Some(f),
error: None,
};
let status = unsafe {
(self.host.scope_stack_with_current)(
self.ptr,
trampoline::<F>,
(&mut state as *mut State<_>).cast(),
)
};
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(state
.error
.unwrap_or_else(|| format!("scope_stack_with_current failed: {status:?}")))
}
}
}
impl Drop for ScopeStack<'_> {
fn drop(&mut self) {
unsafe { (self.host.scope_stack_free)(self.ptr) };
}
}
pub struct ScopeStackBinding<'a> {
host: &'a NemoRelayNativeHostApiV1,
ptr: *mut NemoRelayNativeScopeStackBinding,
}
unsafe impl Send for ScopeStackBinding<'_> {}
impl<'a> ScopeStackBinding<'a> {
pub fn restore(mut self) -> NemoRelayStatus {
let ptr = std::mem::replace(&mut self.ptr, ptr::null_mut());
unsafe { (self.host.scope_stack_restore_thread)(ptr) }
}
}
impl Drop for ScopeStackBinding<'_> {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe { (self.host.scope_stack_binding_free)(self.ptr) };
}
}
}
pub fn current_scope(host: &NemoRelayNativeHostApiV1) -> Result<ScopeHandle<'_>> {
let mut out = ptr::null_mut();
let status = unsafe { (host.scope_get_current)(&mut out) };
if status == NemoRelayStatus::Ok && !out.is_null() {
Ok(ScopeHandle { host, ptr: out })
} else {
Err(format!("scope_get_current failed: {status:?}"))
}
}
pub fn push_scope<'a>(
host: &'a NemoRelayNativeHostApiV1,
name: &str,
scope_type: NemoRelayNativeScopeType,
data: Option<&Json>,
metadata: Option<&Json>,
input: Option<&Json>,
) -> Result<ScopeHandle<'a>> {
let name =
HostString::new(host, name).ok_or_else(|| "failed to allocate scope name".to_string())?;
let data = OptionalHostJson::new(host, data)?;
let metadata = OptionalHostJson::new(host, metadata)?;
let input = OptionalHostJson::new(host, input)?;
let mut out = ptr::null_mut();
let status = unsafe {
(host.scope_push)(
name.as_ptr(),
scope_type,
ptr::null(),
0,
data.as_ptr(),
metadata.as_ptr(),
input.as_ptr(),
ptr::null(),
&mut out,
)
};
if status == NemoRelayStatus::Ok && !out.is_null() {
Ok(ScopeHandle { host, ptr: out })
} else {
Err(format!("scope_push failed: {status:?}"))
}
}
pub fn pop_scope(
host: &NemoRelayNativeHostApiV1,
handle: &ScopeHandle<'_>,
output: Option<&Json>,
metadata: Option<&Json>,
) -> Result<()> {
let output = OptionalHostJson::new(host, output)?;
let metadata = OptionalHostJson::new(host, metadata)?;
let status = unsafe {
(host.scope_pop)(
handle.as_ptr(),
output.as_ptr(),
metadata.as_ptr(),
ptr::null(),
)
};
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(format!("scope_pop failed: {status:?}"))
}
}
pub fn emit_mark(
host: &NemoRelayNativeHostApiV1,
name: &str,
data: Option<&Json>,
metadata: Option<&Json>,
) -> Result<()> {
let name =
HostString::new(host, name).ok_or_else(|| "failed to allocate mark name".to_string())?;
let data = OptionalHostJson::new(host, data)?;
let metadata = OptionalHostJson::new(host, metadata)?;
let status = unsafe {
(host.emit_mark)(
name.as_ptr(),
ptr::null(),
data.as_ptr(),
metadata.as_ptr(),
ptr::null(),
)
};
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(format!("emit_mark failed: {status:?}"))
}
}
#[allow(clippy::too_many_arguments)] fn emit_mark_v2_call(
host: &NemoRelayNativeHostApiV1,
emit_mark_v2: NemoRelayNativeEmitMarkV2Fn,
name: &str,
data: Option<&Json>,
metadata: Option<&Json>,
data_schema: Option<&DataSchema>,
severity: Option<LogSeverity>,
) -> Result<()> {
let name =
HostString::new(host, name).ok_or_else(|| "failed to allocate mark name".to_string())?;
let data = OptionalHostJson::new(host, data)?;
let metadata = OptionalHostJson::new(host, metadata)?;
let data_schema = data_schema
.map(|value| {
HostString::from_json(host, value)
.ok_or_else(|| "failed to serialize mark data schema".to_string())
})
.transpose()?;
let severity = severity
.map(|value| {
serde_json::to_value(value)
.map_err(|err| format!("failed to serialize mark severity: {err}"))
.and_then(|value| {
value
.as_str()
.ok_or_else(|| "mark severity did not serialize as a string".to_string())
.and_then(|value| {
HostString::new(host, value)
.ok_or_else(|| "failed to allocate mark severity".to_string())
})
})
})
.transpose()?;
let status = unsafe {
emit_mark_v2(
name.as_ptr(),
ptr::null(),
data.as_ptr(),
metadata.as_ptr(),
data_schema
.as_ref()
.map(HostString::as_ptr)
.unwrap_or(ptr::null()),
severity
.as_ref()
.map(HostString::as_ptr)
.unwrap_or(ptr::null()),
ptr::null(),
)
};
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(format!("emit_mark_v2 failed: {status:?}"))
}
}
pub fn create_scope_stack(host: &NemoRelayNativeHostApiV1) -> Result<ScopeStack<'_>> {
let mut out = ptr::null_mut();
let status = unsafe { (host.scope_stack_create)(&mut out) };
if status == NemoRelayStatus::Ok && !out.is_null() {
Ok(ScopeStack { host, ptr: out })
} else {
Err(format!("scope_stack_create failed: {status:?}"))
}
}
pub fn capture_scope_stack_thread(
host: &NemoRelayNativeHostApiV1,
) -> Result<ScopeStackBinding<'_>> {
let mut out = ptr::null_mut();
let status = unsafe { (host.scope_stack_capture_thread)(&mut out) };
if status == NemoRelayStatus::Ok && !out.is_null() {
Ok(ScopeStackBinding { host, ptr: out })
} else {
Err(format!("scope_stack_capture_thread failed: {status:?}"))
}
}
pub trait NativePlugin: Send + 'static {
fn plugin_kind(&self) -> &str;
fn allows_multiple_components(&self) -> bool {
true
}
fn executor_config(&self) -> NativeExecutorConfig {
NativeExecutorConfig::default()
}
fn executor_config_for_component(
&self,
plugin_config: &Map<String, Json>,
) -> Result<NativeExecutorConfig> {
self.executor_config().with_component_config(plugin_config)
}
fn validate(&self, plugin_config: &Map<String, Json>) -> Vec<ConfigDiagnostic> {
self.executor_config_for_component(plugin_config)
.err()
.map(|message| ConfigDiagnostic {
level: DiagnosticLevel::Error,
code: "native_executor_config.invalid".into(),
component: None,
field: Some("executor.worker_threads".into()),
message,
})
.into_iter()
.collect()
}
fn register(
&mut self,
plugin_config: &Map<String, Json>,
ctx: &mut PluginContext<'_>,
) -> Result<()>;
}
pub struct PluginContext<'a> {
host: &'a NemoRelayNativeHostApiV1,
raw: *mut NemoRelayNativePluginContext,
executor: Arc<async_sdk::NativeExecutor>,
}
#[allow(clippy::not_unsafe_ptr_arg_deref)]
impl<'a> PluginContext<'a> {
pub unsafe fn from_raw(
host: &'a NemoRelayNativeHostApiV1,
raw: *mut NemoRelayNativePluginContext,
) -> Self {
Self {
host,
raw,
executor: async_sdk::NativeExecutor::new(NativeExecutorConfig::default(), "standalone"),
}
}
unsafe fn from_raw_with_executor(
host: &'a NemoRelayNativeHostApiV1,
raw: *mut NemoRelayNativePluginContext,
executor: Arc<async_sdk::NativeExecutor>,
) -> Self {
Self {
host,
raw,
executor,
}
}
pub fn host_api(&self) -> &'a NemoRelayNativeHostApiV1 {
self.host
}
pub fn runtime(&self) -> PluginRuntime {
PluginRuntime::from_context(self.host, self.raw)
}
pub fn register_conditional_middleware_guardrail(
&mut self,
name: &str,
kinds: &std::collections::BTreeSet<RuntimeRegistrationKind>,
registration_name: &str,
reason: &str,
) -> Result<()> {
if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION
|| self.host.struct_size < std::mem::size_of::<NemoRelayNativeHostApiV4>()
{
return Err("host does not support conditional middleware guardrails".into());
}
let v4 = unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV4) };
let name = HostString::new(self.host, name)
.ok_or_else(|| "failed to allocate gate name".to_string())?;
let kinds = HostString::from_json(self.host, kinds)
.ok_or_else(|| "failed to serialize runtime registration kinds".to_string())?;
let registration_name = HostString::new(self.host, registration_name)
.ok_or_else(|| "failed to allocate target name".to_string())?;
let reason = HostString::new(self.host, reason)
.ok_or_else(|| "failed to allocate gate reason".to_string())?;
let status = unsafe {
(v4.plugin_context_register_conditional_middleware_guardrail)(
self.raw,
name.as_ptr(),
kinds.as_ptr(),
registration_name.as_ptr(),
reason.as_ptr(),
)
};
status_result(
self.host,
status,
"register conditional middleware guardrail",
)
}
pub fn register_subscriber<F>(&mut self, name: &str, callback: F) -> Result<()>
where
F: Fn(&Event) + Send + Sync + 'static,
{
let user_data = typed_callback_user_data(self.host, callback);
let status = unsafe {
self.register_subscriber_raw(
name,
typed_subscriber_trampoline::<F>,
user_data,
Some(drop_typed_callback::<F>),
)
};
finish_typed_registration(self.host, status, user_data, "subscriber")
}
pub unsafe fn register_subscriber_raw(
&mut self,
name: &str,
cb: NemoRelayNativeEventSubscriberCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_subscriber)(self.raw, name, cb, user_data, free_fn)
})
}
pub unsafe fn register_mark_sanitize_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_mark_sanitize_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_scope_sanitize_start_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_scope_sanitize_start_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_scope_sanitize_end_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeEventSanitizeCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_scope_sanitize_end_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_tool_sanitize_request_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_tool_sanitize_request_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_tool_sanitize_response_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_tool_sanitize_response_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_tool_conditional_execution_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeToolConditionalCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_tool_conditional_execution_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_tool_request_intercept_raw(
&mut self,
name: &str,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeToolJsonCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_tool_request_intercept)(
self.raw,
name,
priority,
break_chain,
cb,
user_data,
free_fn,
)
})
}
pub unsafe fn register_tool_execution_intercept_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeToolExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_tool_execution_intercept)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_llm_sanitize_request_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeLlmSanitizeRequestCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_sanitize_request_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_llm_sanitize_response_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeLlmSanitizeResponseCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_sanitize_response_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_llm_conditional_execution_guardrail_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeLlmConditionalCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_conditional_execution_guardrail)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_llm_request_intercept_raw(
&mut self,
name: &str,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeLlmRequestInterceptCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_request_intercept)(
self.raw,
name,
priority,
break_chain,
cb,
user_data,
free_fn,
)
})
}
pub unsafe fn register_llm_execution_intercept_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeLlmExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_execution_intercept)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
pub unsafe fn register_llm_stream_execution_intercept_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeLlmStreamExecutionCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
self.with_name_and_callback(name, user_data, free_fn, |host, name| unsafe {
(host.plugin_context_register_llm_stream_execution_intercept)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
#[allow(clippy::too_many_arguments)] pub unsafe fn register_async_middleware_raw(
&mut self,
kind: NemoRelayNativeAsyncMiddlewareKind,
name: &str,
priority: i32,
break_chain: bool,
cb: NemoRelayNativeAsyncMiddlewareCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE
|| self.host.struct_size < std::mem::size_of::<NemoRelayNativeHostApiV3>()
{
if let Some(free_fn) = free_fn {
unsafe { free_fn(user_data) };
}
return NemoRelayStatus::InvalidArg;
}
let host = unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV3) };
self.with_name_and_callback(name, user_data, free_fn, |_, name| unsafe {
(host.plugin_context_register_async_middleware)(
self.raw,
kind as u32,
name,
priority,
break_chain,
cb,
user_data,
free_fn,
)
})
}
pub unsafe fn register_async_stream_middleware_raw(
&mut self,
name: &str,
priority: i32,
cb: NemoRelayNativeAsyncStreamMiddlewareCb,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
) -> NemoRelayStatus {
if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE
|| self.host.struct_size < std::mem::size_of::<NemoRelayNativeHostApiV3>()
{
if let Some(free_fn) = free_fn {
unsafe { free_fn(user_data) };
}
return NemoRelayStatus::InvalidArg;
}
let host = unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV3) };
self.with_name_and_callback(name, user_data, free_fn, |_, name| unsafe {
(host.plugin_context_register_async_stream_middleware)(
self.raw, name, priority, cb, user_data, free_fn,
)
})
}
fn with_name_and_callback(
&self,
name: &str,
user_data: *mut c_void,
free_fn: NemoRelayNativeFreeFn,
f: impl FnOnce(&NemoRelayNativeHostApiV1, *const NemoRelayNativeString) -> NemoRelayStatus,
) -> NemoRelayStatus {
let name = match HostString::try_new(self.host, name) {
Ok(name) => name,
Err(status) => {
if let Some(free_fn) = free_fn {
unsafe { free_fn(user_data) };
}
return status;
}
};
f(self.host, name.as_ptr())
}
}
struct TypedCallback<F> {
host: NemoRelayNativeHostApiV1,
callback: F,
}
fn typed_callback_user_data<F>(host: &NemoRelayNativeHostApiV1, callback: F) -> *mut c_void {
Box::into_raw(Box::new(TypedCallback {
host: *host,
callback,
})) as *mut c_void
}
unsafe extern "C" fn drop_typed_callback<F>(user_data: *mut c_void) {
if !user_data.is_null() {
let callback = unsafe { Box::from_raw(user_data as *mut TypedCallback<F>) };
let host = callback.host;
if catch_unwind(AssertUnwindSafe(|| drop(callback))).is_err() {
set_last_error(&host, "native plugin typed callback state drop panicked");
}
}
}
fn finish_typed_registration(
host: &NemoRelayNativeHostApiV1,
status: NemoRelayStatus,
user_data: *mut c_void,
label: &str,
) -> Result<()> {
let _ = user_data;
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(status_error(host, status, label))
}
}
fn status_error(host: &NemoRelayNativeHostApiV1, status: NemoRelayStatus, label: &str) -> String {
debug_assert_ne!(status, NemoRelayStatus::Ok);
set_last_error(host, &format!("{label} failed: {status:?}"));
format!("{label} failed: {status:?}")
}
fn status_result(
host: &NemoRelayNativeHostApiV1,
status: NemoRelayStatus,
label: &str,
) -> Result<()> {
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(status_error(host, status, label))
}
}
fn callback_panic(host: &NemoRelayNativeHostApiV1, label: &str) -> NemoRelayStatus {
set_last_error(host, &format!("{label} panicked"));
NemoRelayStatus::Internal
}
unsafe extern "C" fn typed_subscriber_trampoline<F>(
user_data: *mut c_void,
event_json: *const NemoRelayNativeString,
) -> NemoRelayStatus
where
F: Fn(&Event) + Send + Sync + 'static,
{
if user_data.is_null() {
return NemoRelayStatus::NullPointer;
}
let state = unsafe { &*(user_data as *const TypedCallback<F>) };
let result = catch_unwind(AssertUnwindSafe(|| {
let event: Event = read_json_value(&state.host, event_json, "event")?;
(state.callback)(&event);
Ok::<_, NemoRelayStatus>(())
}));
match result {
Ok(Ok(())) => NemoRelayStatus::Ok,
Ok(Err(status)) => status,
Err(_) => callback_panic(&state.host, "subscriber callback"),
}
}
struct HostString<'a> {
host: &'a NemoRelayNativeHostApiV1,
ptr: *mut NemoRelayNativeString,
}
unsafe impl Send for HostString<'_> {}
impl<'a> HostString<'a> {
fn try_new(
host: &'a NemoRelayNativeHostApiV1,
value: &str,
) -> std::result::Result<Self, NemoRelayStatus> {
let mut out = ptr::null_mut();
let status = unsafe { (host.string_new)(value.as_ptr(), value.len(), &mut out) };
if status != NemoRelayStatus::Ok {
return Err(status);
}
if out.is_null() {
return Err(NemoRelayStatus::Internal);
}
Ok(Self { host, ptr: out })
}
fn new(host: &'a NemoRelayNativeHostApiV1, value: &str) -> Option<Self> {
Self::try_new(host, value).ok()
}
fn from_json<T: Serialize>(host: &'a NemoRelayNativeHostApiV1, value: &T) -> Option<Self> {
serde_json::to_string(value)
.ok()
.and_then(|json| Self::new(host, &json))
}
fn as_ptr(&self) -> *const NemoRelayNativeString {
self.ptr
}
}
impl Drop for HostString<'_> {
fn drop(&mut self) {
unsafe { (self.host.string_free)(self.ptr) };
}
}
fn codec_status(host: &NemoRelayNativeHostApiV1, status: NemoRelayStatus) -> Result<()> {
if status == NemoRelayStatus::Ok {
Ok(())
} else {
Err(status_error(host, status, "LLM codec operation"))
}
}
fn native_codec_call<T: DeserializeOwned>(
host: &NemoRelayNativeHostApiV1,
call: impl FnOnce(*mut *mut NemoRelayNativeString) -> Result<()>,
) -> Result<T> {
native_json_call(host, "LLM codec operation", call)
}
fn native_json_call<T: DeserializeOwned>(
host: &NemoRelayNativeHostApiV1,
operation: &str,
call: impl FnOnce(*mut *mut NemoRelayNativeString) -> Result<()>,
) -> Result<T> {
let mut out = ptr::null_mut();
call(&mut out)?;
if out.is_null() {
return Err(format!("{operation} returned null"));
}
let out = HostString { host, ptr: out };
let text = read_host_string(host, out.as_ptr())
.map_err(|_| format!("{operation} returned invalid UTF-8"))?;
serde_json::from_str(&text).map_err(|error| format!("invalid {operation} result: {error}"))
}
struct OptionalHostJson<'a>(Option<HostString<'a>>);
impl<'a> OptionalHostJson<'a> {
fn new(host: &'a NemoRelayNativeHostApiV1, value: Option<&Json>) -> Result<Self> {
match value {
Some(value) => HostString::from_json(host, value)
.map(|value| Self(Some(value)))
.ok_or_else(|| "failed to allocate JSON host string".into()),
None => Ok(Self(None)),
}
}
fn as_ptr(&self) -> *const NemoRelayNativeString {
self.0
.as_ref()
.map(HostString::as_ptr)
.unwrap_or(ptr::null())
}
}
enum OwnedHostApi {
V1(NemoRelayNativeHostApiV1),
V3(NemoRelayNativeHostApiV3),
V4(NemoRelayNativeHostApiV4),
}
impl OwnedHostApi {
unsafe fn copy_from(host: &NemoRelayNativeHostApiV1) -> Self {
if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION
&& host.struct_size >= std::mem::size_of::<NemoRelayNativeHostApiV4>()
{
Self::V4(unsafe { *(host as *const _ as *const NemoRelayNativeHostApiV4) })
} else if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE
&& host.struct_size >= std::mem::size_of::<NemoRelayNativeHostApiV3>()
{
Self::V3(unsafe { *(host as *const _ as *const NemoRelayNativeHostApiV3) })
} else {
Self::V1(*host)
}
}
fn v1(&self) -> &NemoRelayNativeHostApiV1 {
match self {
Self::V1(host) => host,
Self::V3(host) => &host.v1,
Self::V4(host) => &host.v3.v1,
}
}
}
struct PluginState<P> {
host: OwnedHostApi,
plugin: Mutex<P>,
}
unsafe extern "C" fn drop_plugin_state<P: NativePlugin>(user_data: *mut c_void) {
if !user_data.is_null() {
let state = unsafe { Box::from_raw(user_data as *mut PluginState<P>) };
let host = *state.host.v1();
if catch_unwind(AssertUnwindSafe(|| drop(state))).is_err() {
set_last_error(&host, "native plugin state drop panicked");
}
}
}
unsafe extern "C" fn validate_trampoline<P: NativePlugin>(
user_data: *mut c_void,
plugin_config_json: *const NemoRelayNativeString,
out_diagnostics_json: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus {
if user_data.is_null() || out_diagnostics_json.is_null() {
return NemoRelayStatus::NullPointer;
}
unsafe { *out_diagnostics_json = ptr::null_mut() };
let state = unsafe { &*(user_data as *const PluginState<P>) };
let result = catch_unwind(AssertUnwindSafe(|| {
let host = state.host.v1();
let config = match read_json_object(host, plugin_config_json) {
Ok(config) => config,
Err(status) => return status,
};
let plugin = match state.plugin.lock() {
Ok(plugin) => plugin,
Err(_) => {
set_last_error(host, "native plugin state lock poisoned");
return NemoRelayStatus::Internal;
}
};
let diagnostics = plugin.validate(&config);
write_json(host, &diagnostics, out_diagnostics_json)
}));
result.unwrap_or_else(|_| {
set_last_error(state.host.v1(), "native plugin validate callback panicked");
NemoRelayStatus::Internal
})
}
unsafe extern "C" fn register_trampoline<P: NativePlugin>(
user_data: *mut c_void,
plugin_config_json: *const NemoRelayNativeString,
ctx: *mut NemoRelayNativePluginContext,
) -> NemoRelayStatus {
if user_data.is_null() || ctx.is_null() {
return NemoRelayStatus::NullPointer;
}
let state = unsafe { &*(user_data as *const PluginState<P>) };
let result = catch_unwind(AssertUnwindSafe(|| {
let host = state.host.v1();
let config = match read_json_object(host, plugin_config_json) {
Ok(config) => config,
Err(status) => return status,
};
let mut plugin = match state.plugin.lock() {
Ok(plugin) => plugin,
Err(_) => {
set_last_error(host, "native plugin state lock poisoned");
return NemoRelayStatus::Internal;
}
};
let executor_config = match plugin.executor_config_for_component(&config) {
Ok(config) => config,
Err(error) => {
set_last_error(host, &error);
return NemoRelayStatus::InvalidArg;
}
};
let mut ctx = unsafe {
PluginContext::from_raw_with_executor(
host,
ctx,
async_sdk::NativeExecutor::new(executor_config, plugin.plugin_kind()),
)
};
match plugin.register(&config, &mut ctx) {
Ok(()) => NemoRelayStatus::Ok,
Err(message) => {
set_last_error(host, &message);
NemoRelayStatus::Internal
}
}
}));
result.unwrap_or_else(|_| {
set_last_error(state.host.v1(), "native plugin register callback panicked");
NemoRelayStatus::Internal
})
}
fn read_json_object(
host: &NemoRelayNativeHostApiV1,
value: *const NemoRelayNativeString,
) -> std::result::Result<Map<String, Json>, NemoRelayStatus> {
let value: Json = read_json_value(host, value, "plugin config")?;
match value {
Json::Object(map) => Ok(map),
_ => {
set_last_error(host, "plugin config must be a JSON object");
Err(NemoRelayStatus::InvalidJson)
}
}
}
fn read_json_value<T: DeserializeOwned>(
host: &NemoRelayNativeHostApiV1,
value: *const NemoRelayNativeString,
label: &str,
) -> std::result::Result<T, NemoRelayStatus> {
let text = read_required_host_string(host, value, label)?;
serde_json::from_str::<T>(&text).map_err(|error| {
set_last_error(host, &format!("{label} was invalid JSON: {error}"));
NemoRelayStatus::InvalidJson
})
}
#[derive(Debug)]
enum HostStringReadError {
Null,
InvalidUtf8,
}
fn read_required_host_string(
host: &NemoRelayNativeHostApiV1,
value: *const NemoRelayNativeString,
label: &str,
) -> std::result::Result<String, NemoRelayStatus> {
match read_host_string(host, value) {
Ok(value) => Ok(value),
Err(HostStringReadError::Null) => {
set_last_error(host, &format!("{label} was null"));
Err(NemoRelayStatus::NullPointer)
}
Err(HostStringReadError::InvalidUtf8) => {
set_last_error(host, &format!("{label} contained invalid UTF-8"));
Err(NemoRelayStatus::InvalidUtf8)
}
}
}
fn read_host_string(
host: &NemoRelayNativeHostApiV1,
value: *const NemoRelayNativeString,
) -> std::result::Result<String, HostStringReadError> {
if value.is_null() {
return Err(HostStringReadError::Null);
}
let len = unsafe { (host.string_len)(value) };
let data = unsafe { (host.string_data)(value) };
if data.is_null() && len > 0 {
return Err(HostStringReadError::InvalidUtf8);
}
let bytes = if len == 0 {
&[][..]
} else {
unsafe { std::slice::from_raw_parts(data, len) }
};
std::str::from_utf8(bytes)
.map(str::to_owned)
.map_err(|_| HostStringReadError::InvalidUtf8)
}
fn take_host_string(
host: &NemoRelayNativeHostApiV1,
value: *mut NemoRelayNativeString,
) -> Result<String> {
let result = read_host_string(host, value)
.map_err(|error| format!("host returned an invalid string: {error:?}"));
if !value.is_null() {
unsafe { (host.string_free)(value) };
}
result
}
fn take_host_json<T: DeserializeOwned>(
host: &NemoRelayNativeHostApiV1,
value: *mut NemoRelayNativeString,
) -> Result<T> {
let text = take_host_string(host, value)?;
serde_json::from_str(&text).map_err(|error| format!("host returned invalid JSON: {error}"))
}
fn write_json<T: Serialize>(
host: &NemoRelayNativeHostApiV1,
value: &T,
out: *mut *mut NemoRelayNativeString,
) -> NemoRelayStatus {
if out.is_null() {
return NemoRelayStatus::NullPointer;
}
unsafe { *out = ptr::null_mut() };
let json = serde_json::to_value(value).expect("Relay DTOs and serde_json::Value serialize");
let Some(handle) = HostString::from_json(host, &json) else {
set_last_error(host, "failed to allocate host string");
return NemoRelayStatus::Internal;
};
unsafe { *out = handle.ptr };
std::mem::forget(handle);
NemoRelayStatus::Ok
}
fn set_last_error(host: &NemoRelayNativeHostApiV1, message: &str) {
if let Some(message) = HostString::new(host, message) {
unsafe { (host.last_error_set)(message.as_ptr()) };
}
}
#[doc(hidden)]
pub unsafe fn __set_last_error_from_entry(host: *const NemoRelayNativeHostApiV1, message: &str) {
if !host.is_null() {
set_last_error(unsafe { &*host }, message);
}
}
pub unsafe fn export_plugin<P: NativePlugin>(
host: *const NemoRelayNativeHostApiV1,
out: *mut NemoRelayNativePluginV1,
plugin: P,
) -> NemoRelayStatus {
if host.is_null() || out.is_null() {
return NemoRelayStatus::NullPointer;
}
unsafe { *out = NemoRelayNativePluginV1::default() };
let host_ref = unsafe { &*host };
export_plugin_checked(host_ref, out, || plugin)
}
#[doc(hidden)]
pub unsafe fn __export_plugin_from_constructor<P, F>(
host: *const NemoRelayNativeHostApiV1,
out: *mut NemoRelayNativePluginV1,
constructor: F,
) -> NemoRelayStatus
where
P: NativePlugin,
F: FnOnce() -> P,
{
if host.is_null() || out.is_null() {
return NemoRelayStatus::NullPointer;
}
unsafe { *out = NemoRelayNativePluginV1::default() };
let host_ref = unsafe { &*host };
export_plugin_checked(host_ref, out, constructor)
}
fn export_plugin_checked<P, F>(
host_ref: &NemoRelayNativeHostApiV1,
out: *mut NemoRelayNativePluginV1,
constructor: F,
) -> NemoRelayStatus
where
P: NativePlugin,
F: FnOnce() -> P,
{
let supported_abi = (NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY..=NEMO_RELAY_NATIVE_ABI_VERSION)
.contains(&host_ref.abi_version);
if !supported_abi {
return NemoRelayStatus::InvalidArg;
}
if host_ref.struct_size < std::mem::size_of::<NemoRelayNativeHostApiV1>() {
return NemoRelayStatus::InvalidArg;
}
let plugin = constructor();
let kind = plugin.plugin_kind().to_owned();
let allows_multiple_components = plugin.allows_multiple_components();
let Some(kind_handle) = HostString::new(host_ref, &kind) else {
return NemoRelayStatus::Internal;
};
let state = Box::new(PluginState {
host: unsafe { OwnedHostApi::copy_from(host_ref) },
plugin: Mutex::new(plugin),
});
unsafe {
*out = NemoRelayNativePluginV1 {
struct_size: std::mem::size_of::<NemoRelayNativePluginV1>(),
plugin_kind: kind_handle.ptr,
allows_multiple_components,
user_data: Box::into_raw(state) as *mut c_void,
validate: Some(validate_trampoline::<P>),
register: Some(register_trampoline::<P>),
drop: Some(drop_plugin_state::<P>),
};
}
std::mem::forget(kind_handle);
NemoRelayStatus::Ok
}
#[macro_export]
macro_rules! nemo_relay_plugin {
($symbol:ident, $constructor:expr) => {
#[doc = "Native plugin entry symbol generated by `nemo_relay_plugin!`."]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn $symbol(
host: *const $crate::NemoRelayNativeHostApiV1,
out: *mut $crate::NemoRelayNativePluginV1,
) -> $crate::NemoRelayStatus {
match ::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(|| unsafe {
$crate::__export_plugin_from_constructor(host, out, $constructor)
})) {
Ok(status) => status,
Err(_) => {
unsafe {
$crate::__set_last_error_from_entry(
host,
"native plugin entry callback panicked",
)
};
$crate::NemoRelayStatus::Internal
}
}
}
};
}