#![allow(clippy::not_unsafe_ptr_arg_deref)]
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_void};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use serde::de::DeserializeOwned;
use aimux_core::AbortSignal;
use aimux_core::AiMuxError;
use aimux_core::generate::{
GenerateTextOptions, GenerateTextResult, generate_object, generate_text,
generate_text_as_openai, generate_text_result_to_chat_completion, stream_text,
stream_text_as_openai,
};
use aimux_core::language_model::LanguageModel;
use aimux_core::message::ModelPrompt;
use aimux_core::openai_output::OpenAiStreamOptions;
use aimux_core::parse_tool_call::{
ToolCallRepairReply, apply_tool_call_repair, apply_tool_call_repair_to_result,
tool_call_repair_context, tool_call_repair_inputs,
};
use aimux_core::provider::Provider;
use aimux_core::recording::RecordingError;
use aimux_core::tool::ToolCall;
use aimux_core::trace::{RingTraceStore, TraceFilter, TraceLayer};
use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider};
use aimux_providers::anthropic_aws::{AnthropicAwsProvider, AnthropicAwsProviderConfig};
use aimux_providers::azure::{AzureConfig, AzureProvider};
use aimux_providers::bedrock::{BedrockProvider, BedrockProviderConfig};
use aimux_providers::cohere::{CohereConfig, CohereProvider};
use aimux_providers::google::{GoogleConfig, GoogleProvider};
use aimux_providers::mistral::{MistralConfig, MistralProvider};
use aimux_providers::openai::{OpenAIConfig, OpenAIProvider};
use aimux_providers::tavily::{TavilyConfig, TavilyProvider};
use aimux_providers::vertex::{VertexProvider, VertexProviderConfig};
use aimux_providers::xai::{XAIConfig, XAIProvider};
use aimux_providers::{ProviderOptions, provider, provider_handle};
use futures::StreamExt;
use tokio::runtime::Runtime;
#[derive(Clone)]
enum HandleEntry {
Language(Arc<dyn LanguageModel>),
Provider(Arc<dyn aimux_core::provider::Provider>),
Embedding(Arc<dyn aimux_core::embedding_model::EmbeddingModel>),
Speech(Arc<dyn aimux_core::speech_model::SpeechModel>),
Image(Arc<dyn aimux_core::image_model::ImageModel>),
Transcription(Arc<dyn aimux_core::transcription_model::TranscriptionModel>),
Reranking(Arc<dyn aimux_core::reranking_model::RerankingModel>),
Video(Arc<dyn aimux_core::video_model::VideoModel>),
Search(Arc<dyn aimux_core::search_model::SearchModel>),
Files(Arc<dyn aimux_core::files_model::Files>),
TranscriptionSession(Arc<transcription_session::TranscriptionFfiSession>),
Abort(AbortSignal),
}
type Registry = HashMap<u64, HandleEntry>;
static REGISTRY: OnceLock<Mutex<Registry>> = OnceLock::new();
static NEXT_HANDLE: AtomicU64 = AtomicU64::new(1);
fn registry() -> &'static Mutex<Registry> {
REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}
fn intern_model(model: Arc<dyn LanguageModel>) -> u64 {
intern_handle(HandleEntry::Language(model))
}
fn intern_handle(h: HandleEntry) -> u64 {
let handle = NEXT_HANDLE.fetch_add(1, Ordering::Relaxed);
registry()
.lock()
.expect("aimux-ffi: registry mutex poisoned")
.insert(handle, h);
handle
}
fn get_handle(handle: u64) -> Option<HandleEntry> {
registry()
.lock()
.expect("aimux-ffi: registry mutex poisoned")
.get(&handle)
.cloned()
}
fn get_model(handle: u64) -> Option<Arc<dyn LanguageModel>> {
match get_handle(handle)? {
HandleEntry::Language(m) => Some(m),
_ => None,
}
}
fn model_of(handle: u64) -> Result<Arc<dyn LanguageModel>, FfiError> {
get_model(handle).ok_or(FfiError::InvalidHandle { expected: "model" })
}
fn entry_of(handle: u64, expected: &'static str) -> Result<HandleEntry, FfiError> {
get_handle(handle).ok_or(FfiError::InvalidHandle { expected })
}
fn get_abort_signal(handle: u64) -> Option<AbortSignal> {
match get_handle(handle)? {
HandleEntry::Abort(signal) => Some(signal),
_ => None,
}
}
fn abort_of(handle: u64) -> Result<AbortSignal, FfiError> {
get_abort_signal(handle).ok_or(FfiError::InvalidHandle { expected: "abort" })
}
fn drop_handle(handle: u64) {
let removed = registry()
.lock()
.expect("aimux-ffi: registry mutex poisoned")
.remove(&handle);
if let Some(HandleEntry::TranscriptionSession(s)) = removed {
s.terminate();
}
if let Some(stores) = TRACE_STORES.get() {
stores
.lock()
.expect("aimux-ffi: trace registry mutex poisoned")
.remove(&handle);
}
}
fn drop_abort_signal(handle: u64) {
let mut registry = registry()
.lock()
.expect("aimux-ffi: registry mutex poisoned");
if matches!(registry.get(&handle), Some(HandleEntry::Abort(_))) {
registry.remove(&handle);
}
}
mod transcription_session;
fn runtime() -> &'static Runtime {
static RUNTIME: OnceLock<Runtime> = OnceLock::new();
RUNTIME.get_or_init(|| {
tokio::runtime::Runtime::new().expect("aimux-ffi: failed to build tokio runtime")
})
}
thread_local! {
static IN_FFI_BLOCK_ON: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
fn ffi_block_on<F, T>(f: F) -> Result<T, FfiError>
where
F: std::future::Future<Output = T>,
{
IN_FFI_BLOCK_ON.with(|flag| {
if flag.replace(true) {
return Err(FfiError::ReentrantCall);
}
struct Reset;
impl Drop for Reset {
fn drop(&mut self) {
IN_FFI_BLOCK_ON.with(|flag| flag.set(false));
}
}
let _reset = Reset;
Ok(runtime().block_on(f))
})
}
#[derive(Debug, Clone)]
enum FfiError {
NullPointer { argument: &'static str },
InvalidUtf8 { argument: &'static str },
InvalidWireJson {
argument: &'static str,
message: String,
},
InvalidHandle { expected: &'static str },
ReentrantCall,
ResultSerialization { message: String },
CallbackFailure { message: String },
}
impl std::fmt::Display for FfiError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FfiError::NullPointer { argument } => write!(f, "{argument}: must not be NULL"),
FfiError::InvalidUtf8 { argument } => write!(f, "{argument}: must be valid UTF-8"),
FfiError::InvalidWireJson { argument, message } => {
write!(f, "{argument}: invalid JSON: {message}")
}
FfiError::InvalidHandle { expected } => {
write!(f, "invalid or expired {expected} handle")
}
FfiError::ReentrantCall => {
f.write_str("re-entrant FFI call from within a callback is not allowed")
}
FfiError::ResultSerialization { message } => {
write!(f, "could not serialize result: {message}")
}
FfiError::CallbackFailure { message } => write!(f, "host callback failed: {message}"),
}
}
}
enum AiMuxFfiError {
AiMux(AiMuxError),
Recording(RecordingError),
Ffi(FfiError),
}
impl From<AiMuxError> for AiMuxFfiError {
fn from(inner: AiMuxError) -> Self {
AiMuxFfiError::AiMux(inner)
}
}
impl From<RecordingError> for AiMuxFfiError {
fn from(inner: RecordingError) -> Self {
AiMuxFfiError::Recording(inner)
}
}
impl From<FfiError> for AiMuxFfiError {
fn from(e: FfiError) -> Self {
AiMuxFfiError::Ffi(e)
}
}
impl std::fmt::Display for AiMuxFfiError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AiMuxFfiError::AiMux(e) => e.fmt(f),
AiMuxFfiError::Recording(e) => e.fmt(f),
AiMuxFfiError::Ffi(e) => e.fmt(f),
}
}
}
#[allow(non_camel_case_types)]
pub struct aimux_error_t {
error: AiMuxFfiError,
}
type FfiResult<T> = Result<T, AiMuxFfiError>;
fn finish<T>(r: FfiResult<T>, write: impl FnOnce(T)) -> *mut aimux_error_t {
match r {
Ok(v) => {
write(v);
std::ptr::null_mut()
}
Err(error) => Box::into_raw(Box::new(aimux_error_t { error })),
}
}
fn with_out_handle(out_handle: *mut u64, f: impl FnOnce() -> FfiResult<u64>) -> *mut aimux_error_t {
if out_handle.is_null() {
return finish(
Err(FfiError::NullPointer {
argument: "out_handle",
}
.into()),
|_: u64| {},
);
}
unsafe { *out_handle = 0 };
finish(f(), |h| unsafe { *out_handle = h })
}
fn with_out_string(
out: *mut *mut c_char,
argument: &'static str,
f: impl FnOnce() -> FfiResult<String>,
) -> *mut aimux_error_t {
if out.is_null() {
return finish(
Err(FfiError::NullPointer { argument }.into()),
|_: String| {},
);
}
unsafe { *out = std::ptr::null_mut() };
finish(f(), |s| unsafe { *out = into_cstring_raw(s) })
}
fn no_result(f: impl FnOnce() -> FfiResult<()>) -> *mut aimux_error_t {
finish(f(), |()| {})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_free(err: *mut aimux_error_t) {
if !err.is_null() {
drop(unsafe { Box::from_raw(err) });
}
}
pub const AIMUX_OK: i32 = 0;
pub const AIMUX_E_OTHER: i32 = 1;
pub const AIMUX_E_JSON_PARSE: i32 = 2;
pub const AIMUX_E_INVALID_RESPONSE_DATA: i32 = 3;
pub const AIMUX_E_INVALID_ARGUMENT: i32 = 5;
pub const AIMUX_E_INVALID_PROMPT: i32 = 6;
pub const AIMUX_E_TOKEN_EXPIRED: i32 = 7;
pub const AIMUX_E_UNSUPPORTED_FUNCTIONALITY: i32 = 8;
pub const AIMUX_E_NO_SUCH_MODEL: i32 = 9;
pub const AIMUX_E_NO_SUCH_PROVIDER: i32 = 10;
pub const AIMUX_E_API_CALL: i32 = 11;
pub const AIMUX_E_TIMEOUT: i32 = 12;
pub const AIMUX_E_ABORTED: i32 = 13;
pub const AIMUX_E_RETRY: i32 = 14;
pub const AIMUX_E_NO_SUCH_TOOL: i32 = 15;
pub const AIMUX_E_INVALID_TOOL_INPUT: i32 = 16;
pub const AIMUX_E_TOOL_CALL_REPAIR: i32 = 17;
pub const AIMUX_E_RECORDING_INIT: i32 = 100;
pub const AIMUX_E_RECORDING_OPEN_FILE: i32 = 101;
pub const AIMUX_E_RECORDING_SPAWN: i32 = 102;
pub const AIMUX_E_RECORDING_WRITER_GONE: i32 = 103;
pub const AIMUX_E_RECORDING_FLUSH_TIMEOUT: i32 = 104;
pub const AIMUX_E_RECORDING_WRITE: i32 = 105;
pub const AIMUX_E_FFI_NULL_POINTER: i32 = 200;
pub const AIMUX_E_FFI_INVALID_UTF8: i32 = 201;
pub const AIMUX_E_FFI_INVALID_WIRE_JSON: i32 = 202;
pub const AIMUX_E_FFI_INVALID_HANDLE: i32 = 203;
pub const AIMUX_E_FFI_REENTRANT_CALL: i32 = 204;
pub const AIMUX_E_FFI_RESULT_SERIALIZATION: i32 = 205;
pub const AIMUX_E_FFI_CALLBACK_FAILURE: i32 = 206;
pub const AIMUX_TRANSCRIPTION_NEXT_PART_PART: i32 = 1;
pub const AIMUX_TRANSCRIPTION_NEXT_PART_ENDED: i32 = 2;
pub const AIMUX_TRANSCRIPTION_NEXT_PART_TIMEOUT: i32 = 3;
fn aimux_error_code_of(err: &AiMuxError) -> i32 {
match err {
AiMuxError::ApiCall { .. } => AIMUX_E_API_CALL,
AiMuxError::Retry(_) => AIMUX_E_RETRY,
AiMuxError::JsonParse(_) => AIMUX_E_JSON_PARSE,
AiMuxError::InvalidResponseData(_) => AIMUX_E_INVALID_RESPONSE_DATA,
AiMuxError::NoSuchTool { .. } => AIMUX_E_NO_SUCH_TOOL,
AiMuxError::InvalidToolInput { .. } => AIMUX_E_INVALID_TOOL_INPUT,
AiMuxError::ToolCallRepair { .. } => AIMUX_E_TOOL_CALL_REPAIR,
AiMuxError::InvalidArgument(_) => AIMUX_E_INVALID_ARGUMENT,
AiMuxError::InvalidPrompt(_) => AIMUX_E_INVALID_PROMPT,
AiMuxError::TokenExpired(_) => AIMUX_E_TOKEN_EXPIRED,
AiMuxError::UnsupportedFunctionality(_) => AIMUX_E_UNSUPPORTED_FUNCTIONALITY,
AiMuxError::NoSuchModel { .. } => AIMUX_E_NO_SUCH_MODEL,
AiMuxError::NoSuchProvider { .. } => AIMUX_E_NO_SUCH_PROVIDER,
AiMuxError::Timeout(_) => AIMUX_E_TIMEOUT,
AiMuxError::Aborted(_) => AIMUX_E_ABORTED,
AiMuxError::Other(_) => AIMUX_E_OTHER,
}
}
fn recording_error_code_of(e: &RecordingError) -> i32 {
use RecordingError as R;
match e {
R::Init { .. } => AIMUX_E_RECORDING_INIT,
R::OpenFile { .. } => AIMUX_E_RECORDING_OPEN_FILE,
R::Spawn { .. } => AIMUX_E_RECORDING_SPAWN,
R::WriterGone => AIMUX_E_RECORDING_WRITER_GONE,
R::FlushTimeout => AIMUX_E_RECORDING_FLUSH_TIMEOUT,
R::Write(_) => AIMUX_E_RECORDING_WRITE,
}
}
fn ffi_error_code_of(e: &FfiError) -> i32 {
match e {
FfiError::NullPointer { .. } => AIMUX_E_FFI_NULL_POINTER,
FfiError::InvalidUtf8 { .. } => AIMUX_E_FFI_INVALID_UTF8,
FfiError::InvalidWireJson { .. } => AIMUX_E_FFI_INVALID_WIRE_JSON,
FfiError::InvalidHandle { .. } => AIMUX_E_FFI_INVALID_HANDLE,
FfiError::ReentrantCall => AIMUX_E_FFI_REENTRANT_CALL,
FfiError::ResultSerialization { .. } => AIMUX_E_FFI_RESULT_SERIALIZATION,
FfiError::CallbackFailure { .. } => AIMUX_E_FFI_CALLBACK_FAILURE,
}
}
fn error_code_of(e: &AiMuxFfiError) -> i32 {
match e {
AiMuxFfiError::AiMux(e) => aimux_error_code_of(e),
AiMuxFfiError::Recording(e) => recording_error_code_of(e),
AiMuxFfiError::Ffi(e) => ffi_error_code_of(e),
}
}
fn api_call(e: &AiMuxError) -> Option<&aimux_core::ApiCallError> {
match e {
AiMuxError::ApiCall(d) => Some(d),
_ => None,
}
}
fn retry(e: &AiMuxError) -> Option<&aimux_core::RetryError> {
match e {
AiMuxError::Retry(r) => Some(r),
_ => None,
}
}
fn map_aimux_error<T>(err: *const aimux_error_t, f: impl FnOnce(&AiMuxError) -> T) -> Option<T> {
match unsafe { err.as_ref() } {
Some(aimux_error_t {
error: AiMuxFfiError::AiMux(e),
}) => Some(f(e)),
_ => None,
}
}
fn opt_cstring(v: Option<String>) -> *mut c_char {
v.map_or(std::ptr::null_mut(), into_cstring_raw)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_code(err: *const aimux_error_t) -> i32 {
unsafe { err.as_ref() }.map_or(AIMUX_OK, |e| error_code_of(&e.error))
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_message(err: *const aimux_error_t) -> *mut c_char {
unsafe { err.as_ref() }.map_or(std::ptr::null_mut(), |e| {
into_cstring_raw(e.error.to_string())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_status(err: *const aimux_error_t) -> i32 {
map_aimux_error(err, AiMuxError::status_code)
.flatten()
.map_or(-1, i32::from)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_retry_ms(err: *const aimux_error_t) -> i64 {
map_aimux_error(err, AiMuxError::retry_after_hint)
.flatten()
.unwrap_or(-1)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_retryable(err: *const aimux_error_t) -> i32 {
map_aimux_error(err, AiMuxError::is_retryable).unwrap_or(false) as i32
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_provider_code(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(map_aimux_error(err, |e| api_call(e)?.provider_code.clone()).flatten())
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_provider_message(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| Some(api_call(e)?.message.clone()))
.flatten()
.filter(|m| !m.is_empty()),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_response_body(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(map_aimux_error(err, |e| api_call(e)?.response_body.clone()).flatten())
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_url(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| Some(api_call(e)?.url.clone()))
.flatten()
.filter(|u| !u.is_empty()),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_request_body_values(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| {
let values = &api_call(e)?.request_body_values;
if values.is_null() {
return None;
}
serde_json::to_string(values).ok()
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_response_headers(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| {
serde_json::to_string(api_call(e)?.response_headers.as_ref()?).ok()
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_provider_data(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| {
serde_json::to_string(api_call(e)?.data.as_ref()?).ok()
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_retry_reason(err: *const aimux_error_t) -> *mut c_char {
let reason = map_aimux_error(err, |e| {
Some(match retry(e)?.reason {
aimux_core::RetryErrorReason::MaxRetriesExceeded => "maxRetriesExceeded",
aimux_core::RetryErrorReason::ErrorNotRetryable => "errorNotRetryable",
})
})
.flatten();
opt_cstring(reason.map(str::to_string))
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_retry_count(err: *const aimux_error_t) -> i32 {
map_aimux_error(err, |e| Some(retry(e)?.errors.len() as i32))
.flatten()
.unwrap_or(0)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_retry_error_at(
err: *const aimux_error_t,
index: i32,
) -> *mut aimux_error_t {
if index < 0 {
return std::ptr::null_mut();
}
map_aimux_error(err, |e| {
retry(e)?.errors.get(index as usize).map(|attempt| {
Box::into_raw(Box::new(aimux_error_t {
error: AiMuxFfiError::AiMux(attempt.clone()),
}))
})
})
.flatten()
.unwrap_or_else(std::ptr::null_mut)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_model_id(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::NoSuchModel { model_id, .. } => Some(model_id.clone()),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_model_type(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::NoSuchModel { model_type, .. } => Some(model_type.clone()),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_provider_id(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::NoSuchProvider { provider_id } => Some(provider_id.clone()),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_tool_name(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::NoSuchTool { tool_name, .. }
| AiMuxError::InvalidToolInput { tool_name, .. } => Some(tool_name.clone()),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_available_tools(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::NoSuchTool {
available_tools: Some(tools),
..
} => serde_json::to_string(tools).ok(),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_tool_input(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::InvalidToolInput { tool_input, .. } => Some(tool_input.clone()),
AiMuxError::NoSuchTool { tool_input, .. } => tool_input.clone(),
_ => None,
})
.flatten(),
)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_error_original_error(err: *const aimux_error_t) -> *mut c_char {
opt_cstring(
map_aimux_error(err, |e| match e {
AiMuxError::ToolCallRepair { original_error, .. } => {
serde_json::to_string(original_error).ok()
}
_ => None,
})
.flatten(),
)
}
fn str_arg(ptr: *const c_char, argument: &'static str) -> Result<String, FfiError> {
if ptr.is_null() {
return Err(FfiError::NullPointer { argument });
}
unsafe { CStr::from_ptr(ptr) }
.to_str()
.map(str::to_owned)
.map_err(|_| FfiError::InvalidUtf8 { argument })
}
fn opt_str_arg(ptr: *const c_char, argument: &'static str) -> Result<Option<String>, FfiError> {
if ptr.is_null() {
return Ok(None);
}
str_arg(ptr, argument).map(Some)
}
fn parse_two_args(
a: *const c_char,
an: &'static str,
b: *const c_char,
bn: &'static str,
) -> Result<(String, String), FfiError> {
Ok((str_arg(a, an)?, str_arg(b, bn)?))
}
#[allow(clippy::too_many_arguments)]
fn parse_four_args(
a: *const c_char,
an: &'static str,
b: *const c_char,
bn: &'static str,
c: *const c_char,
cn: &'static str,
d: *const c_char,
dn: &'static str,
) -> Result<(String, String, String, String), FfiError> {
Ok((
str_arg(a, an)?,
str_arg(b, bn)?,
str_arg(c, cn)?,
str_arg(d, dn)?,
))
}
fn parse_base_url(base_url: *const c_char) -> Result<Option<String>, FfiError> {
Ok(opt_str_arg(base_url, "base_url")?.filter(|url| !url.is_empty()))
}
fn wire_failure(argument: &'static str, e: &serde_json::Error, detail: String) -> AiMuxFfiError {
match e.classify() {
serde_json::error::Category::Data => {
AiMuxError::InvalidArgument(format!("{argument}: {detail}")).into()
}
_ => FfiError::InvalidWireJson {
argument,
message: detail,
}
.into(),
}
}
fn wire_err(argument: &'static str, e: serde_json::Error) -> AiMuxFfiError {
let detail = e.to_string();
wire_failure(argument, &e, detail)
}
fn parse_json_arg<T: DeserializeOwned>(json: *const c_char, name: &'static str) -> FfiResult<T> {
let s = str_arg(json, name)?;
serde_json::from_str::<T>(&s).map_err(|e| wire_err(name, e))
}
fn parse_prompt_arg(prompt_json: *const c_char) -> FfiResult<ModelPrompt> {
let s = str_arg(prompt_json, "prompt_json")?;
let parse = |json: &str| -> Result<ModelPrompt, serde_json::Error> {
let value: serde_json::Value = serde_json::from_str(json)?;
let inner = match &value {
serde_json::Value::Object(obj) if obj.len() == 1 && obj.contains_key("prompt") => {
obj.get("prompt").expect("checked by guard")
}
_ => &value,
};
serde_json::from_value(inner.clone())
};
parse(&s).map_err(|e| wire_err("prompt_json", e))
}
fn parse_opts_arg(opts_json: *const c_char) -> FfiResult<GenerateTextOptions> {
match opt_str_arg(opts_json, "opts_json")? {
Some(s) if !s.trim().is_empty() && s.trim() != "null" => {
serde_json::from_str(&s).map_err(|e| wire_err("opts_json", e))
}
_ => Ok(GenerateTextOptions::default()),
}
}
fn parse_provider_options(config_json: *const c_char) -> FfiResult<Option<ProviderOptions>> {
match opt_str_arg(config_json, "config_json")? {
Some(s) if !s.trim().is_empty() && s.trim() != "null" => {
serde_json::from_str::<ProviderOptions>(&s)
.map(Some)
.map_err(|e| wire_err("config_json", e))
}
_ => Ok(None),
}
}
fn normalize_config_json(json: *const c_char, argument: &'static str) -> Result<String, FfiError> {
Ok(match opt_str_arg(json, argument)? {
Some(s) if !s.trim().is_empty() && s.trim() != "null" => s,
_ => String::from("{}"),
})
}
fn into_cstring_raw(s: String) -> *mut c_char {
if s.contains('\0') {
tracing::warn!(
target: "aimux_ffi",
nul_bytes = s.bytes().filter(|&b| b == 0).count(),
"string contains NUL byte(s); replaced with U+FFFD before FFI return"
);
}
let sanitized = s.replace('\0', "\u{FFFD}");
CString::new(sanitized)
.expect("aimux-ffi: impossible: NUL-free string rejected by CString::new")
.into_raw()
}
fn to_json<T: serde::Serialize>(v: &T) -> FfiResult<String> {
serde_json::to_string(v).map_err(|e| {
FfiError::ResultSerialization {
message: format!("serialize: {e}"),
}
.into()
})
}
fn run_json<F, T>(f: F) -> FfiResult<String>
where
F: std::future::Future<Output = Result<T, AiMuxError>>,
T: serde::Serialize,
{
let v = ffi_block_on(f)??;
to_json(&v)
}
fn invoke_stream_callback(callback_name: &str, f: impl FnOnce()) -> Result<(), FfiError> {
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
Ok(()) => Ok(()),
Err(payload) => {
let msg = payload
.downcast_ref::<&'static str>()
.copied()
.or_else(|| payload.downcast_ref::<String>().map(String::as_str))
.unwrap_or("<non-string panic>");
Err(FfiError::CallbackFailure {
message: format!("stream callback '{callback_name}' panicked: {msg}"),
})
}
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let m = OpenAIProvider::new(OpenAIConfig::new(api_key)).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = OpenAIProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_anthropic_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let m = AnthropicProvider::new(AnthropicConfig::new(api_key)).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_anthropic_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = AnthropicConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = AnthropicProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_anthropic_aws_new(
api_key: *const c_char,
region: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let region = str_arg(region, "region")?;
let model_id = str_arg(model_id, "model_id")?;
let m =
AnthropicAwsProvider::new(AnthropicAwsProviderConfig::with_api_key(api_key, region))
.language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_anthropic_aws_new_with_base(
api_key: *const c_char,
region: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let region = str_arg(region, "region")?;
let model_id = str_arg(model_id, "model_id")?;
let mut config = AnthropicAwsProviderConfig::with_api_key(api_key, region);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = AnthropicAwsProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_azure_new(
api_key: *const c_char,
resource_name: *const c_char,
deployment: *const c_char,
api_version: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let resource_name = str_arg(resource_name, "resource_name")?;
let deployment = str_arg(deployment, "deployment")?;
let mut config = AzureConfig::new()
.with_api_key(api_key)
.with_resource_name(resource_name);
if let Some(v) = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()) {
config = config.with_api_version(v);
}
let m = AzureProvider::new(config)?.language_model(&deployment)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_azure_new_with_base(
api_key: *const c_char,
base_url: *const c_char,
deployment: *const c_char,
api_version: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let base_url = str_arg(base_url, "base_url")?;
let deployment = str_arg(deployment, "deployment")?;
let mut config = AzureConfig::new()
.with_api_key(api_key)
.with_base_url(base_url);
if let Some(v) = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()) {
config = config.with_api_version(v);
}
let m = AzureProvider::new(config)?.language_model(&deployment)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_bedrock_new(
access_key_id: *const c_char,
secret_access_key: *const c_char,
region: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (access_key_id, secret_access_key, region, model_id) = parse_four_args(
access_key_id,
"access_key_id",
secret_access_key,
"secret_access_key",
region,
"region",
model_id,
"model_id",
)?;
let m = BedrockProvider::new(BedrockProviderConfig::new(
access_key_id,
secret_access_key,
region,
))
.language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_bedrock_new_with_base(
access_key_id: *const c_char,
secret_access_key: *const c_char,
region: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (access_key_id, secret_access_key, region, model_id) = parse_four_args(
access_key_id,
"access_key_id",
secret_access_key,
"secret_access_key",
region,
"region",
model_id,
"model_id",
)?;
let mut config = BedrockProviderConfig::new(access_key_id, secret_access_key, region);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = BedrockProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_vertex_new(
access_token: *const c_char,
project: *const c_char,
location: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (access_token, project, location, model_id) = parse_four_args(
access_token,
"access_token",
project,
"project",
location,
"location",
model_id,
"model_id",
)?;
let m = VertexProvider::new(VertexProviderConfig::new(access_token, project, location))
.language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_vertex_new_with_base(
access_token: *const c_char,
project: *const c_char,
location: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (access_token, project, location, model_id) = parse_four_args(
access_token,
"access_token",
project,
"project",
location,
"location",
model_id,
"model_id",
)?;
let mut config = VertexProviderConfig::new(access_token, project, location);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = VertexProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let m = CohereProvider::new(CohereConfig::new(api_key)).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = CohereConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = CohereProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_mistral_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let m = MistralProvider::new(MistralConfig::new(api_key)).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_mistral_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = MistralConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = MistralProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_xai_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let m = XAIProvider::new(XAIConfig::new(api_key)).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_xai_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = XAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let m = XAIProvider::new(config).language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_provider_new(
name: *const c_char,
api_key: *const c_char,
model_id: *const c_char,
config_json: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let name = str_arg(name, "name")?;
let model_id = str_arg(model_id, "model_id")?;
let key = opt_str_arg(api_key, "api_key")?;
let opts = parse_provider_options(config_json)?;
let m = provider(&name, key, &model_id, opts)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_provider_from_env(
name: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let name = str_arg(name, "name")?;
let model_id = str_arg(model_id, "model_id")?;
let m = provider(&name, None, &model_id, None)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_provider_handle_new(
name: *const c_char,
api_key: *const c_char,
config_json: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let name = str_arg(name, "name")?;
let key = opt_str_arg(api_key, "api_key")?;
let opts = parse_provider_options(config_json)?;
let p = provider_handle(&name, key, opts)?;
Ok(intern_handle(HandleEntry::Provider(Arc::from(p))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_provider_list_models(
handle: u64,
out_models_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_models_json, "out_models_json", || {
let HandleEntry::Provider(p) = entry_of(handle, "provider")? else {
return Err(FfiError::InvalidHandle {
expected: "provider",
}
.into());
};
run_json(p.list_models())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_provider_model(
handle: u64,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let HandleEntry::Provider(p) = entry_of(handle, "provider")? else {
return Err(FfiError::InvalidHandle {
expected: "provider",
}
.into());
};
let model_id = str_arg(model_id, "model_id")?;
let m = p.language_model(&model_id)?;
Ok(intern_model(Arc::from(m)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_get_model_specs(
source_url: *const c_char,
out_specs_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_specs_json, "out_specs_json", || {
let url = opt_str_arg(source_url, "source_url")?;
run_json(aimux_providers::get_model_specs(url.as_deref()))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_generate_text(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let opts = parse_opts_arg(opts_json)?;
run_json(async move { generate_text(&*model, prompt, opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_generate_object(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let opts = parse_opts_arg(opts_json)?;
run_json(async move { generate_object(&*model, prompt, opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_consume_stream_text(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let opts = parse_opts_arg(opts_json)?;
run_json(async move {
let stream_result = stream_text(&*model, prompt, opts).await?;
stream_result.consume().await
})
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_stream_text(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
) -> *mut aimux_error_t {
no_result(|| {
stream_text_with_signal(
handle,
prompt_json,
opts_json,
on_part,
on_done,
stream_ctx,
None,
)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_abort_signal_new() -> u64 {
intern_handle(HandleEntry::Abort(AbortSignal::new()))
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_abort_signal_abort(handle: u64) {
if let Some(signal) = get_abort_signal(handle) {
signal.abort();
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_abort_signal_drop(handle: u64) {
if handle != 0 {
drop_abort_signal(handle);
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_stream_text_with_abort(
handle: u64,
abort_handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
) -> *mut aimux_error_t {
no_result(|| {
let abort_signal = abort_of(abort_handle)?;
stream_text_with_signal(
handle,
prompt_json,
opts_json,
on_part,
on_done,
stream_ctx,
Some(abort_signal),
)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_tool_call_repair_context(
tool_call_json: *const c_char,
prompt_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let tool_call: ToolCall = parse_json_arg(tool_call_json, "tool_call_json")?;
let prompt = parse_prompt_arg(prompt_json)?;
let opts = parse_opts_arg(opts_json)?;
let (messages, instructions, tools) = tool_call_repair_inputs(prompt, &opts);
to_json(&tool_call_repair_context(
&tool_call,
tools,
&messages,
instructions,
)?)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_apply_tool_call_repair(
tool_call_json: *const c_char,
opts_json: *const c_char,
reply_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let tool_call: ToolCall = parse_json_arg(tool_call_json, "tool_call_json")?;
let opts = parse_opts_arg(opts_json)?;
let reply: ToolCallRepairReply = parse_json_arg(reply_json, "reply_json")?;
to_json(&apply_tool_call_repair(
&tool_call,
opts.tools.as_deref(),
reply,
)?)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_apply_tool_call_repair_to_result(
result_json: *const c_char,
opts_json: *const c_char,
tool_call_id: *const c_char,
reply_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let result: serde_json::Value = parse_json_arg(result_json, "result_json")?;
let opts = parse_opts_arg(opts_json)?;
let tool_call_id = str_arg(tool_call_id, "tool_call_id")?;
let reply: ToolCallRepairReply = parse_json_arg(reply_json, "reply_json")?;
to_json(&apply_tool_call_repair_to_result(
&result,
opts.tools.as_deref(),
&tool_call_id,
reply,
)?)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_generate_text_result_as_openai(
handle: u64,
result_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let model = model_of(handle)?;
let result: GenerateTextResult = parse_json_arg(result_json, "result_json")?;
to_json(&generate_text_result_to_chat_completion(
&result,
model.model_id(),
))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_generate_text_as_openai(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let opts = parse_opts_arg(opts_json)?;
run_json(async move { generate_text_as_openai(&*model, prompt, opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_stream_text_as_openai(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
) -> *mut aimux_error_t {
no_result(|| {
stream_text_as_openai_with_signal(
handle,
prompt_json,
opts_json,
on_part,
on_done,
stream_ctx,
None,
)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_stream_text_as_openai_with_abort(
handle: u64,
abort_handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
) -> *mut aimux_error_t {
no_result(|| {
let abort_signal = abort_of(abort_handle)?;
stream_text_as_openai_with_signal(
handle,
prompt_json,
opts_json,
on_part,
on_done,
stream_ctx,
Some(abort_signal),
)
})
}
fn stream_part_cstring<T: serde::Serialize>(part: &T) -> Result<CString, FfiError> {
let json = serde_json::to_string(part).map_err(|e| FfiError::ResultSerialization {
message: format!("stream part: {e}"),
})?;
CString::new(json).map_err(|e| FfiError::ResultSerialization {
message: format!("stream part contains NUL: {e}"),
})
}
fn stream_text_as_openai_with_signal(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
abort_signal: Option<AbortSignal>,
) -> FfiResult<()> {
let on_part = on_part.ok_or(FfiError::NullPointer {
argument: "on_part",
})?;
let on_done = on_done.ok_or(FfiError::NullPointer {
argument: "on_done",
})?;
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let mut opts = parse_opts_arg(opts_json)?;
let stream_options = opts
.provider_options
.as_ref()
.and_then(|po| po.get("openai"))
.and_then(|o| o.get("stream_options"))
.cloned()
.map(|v| OpenAiStreamOptions {
include_usage: v
.get("include_usage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(true),
include_reasoning: v
.get("include_reasoning")
.and_then(serde_json::Value::as_bool)
.unwrap_or(true),
})
.unwrap_or_default();
opts.abort_signal = abort_signal.clone();
let stream_ctx = stream_ctx as usize;
ffi_block_on(async move {
let sr = match abort_signal.as_ref() {
Some(signal) => {
tokio::select! {
biased;
_ = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)),
result = stream_text_as_openai(&*model, prompt, opts, stream_options) => result,
}
}
None => stream_text_as_openai(&*model, prompt, opts, stream_options).await,
}?;
let mut stream = sr.stream;
loop {
let next = match abort_signal.as_ref() {
Some(signal) => {
tokio::select! {
biased;
_ = signal.cancelled() => {
return Err(AiMuxError::from_abort_signal(signal).into());
}
item = stream.next() => item,
}
}
None => stream.next().await,
};
let Some(item) = next else {
break;
};
let cstr = match item {
Ok(chunk) => stream_part_cstring(&chunk)?,
Err(e) if e.is_recoverable_stream_error() => continue,
Err(e) => return Err(e.into()),
};
invoke_stream_callback("on_part", || {
on_part(cstr.as_ptr(), stream_ctx as *mut c_void);
})?;
}
invoke_stream_callback("on_done", || {
on_done(stream_ctx as *mut c_void);
})?;
Ok::<(), AiMuxFfiError>(())
})?
}
fn stream_text_with_signal(
handle: u64,
prompt_json: *const c_char,
opts_json: *const c_char,
on_part: Option<extern "C-unwind" fn(*const c_char, *mut c_void)>,
on_done: Option<extern "C-unwind" fn(*mut c_void)>,
stream_ctx: *mut c_void,
abort_signal: Option<AbortSignal>,
) -> FfiResult<()> {
let on_part = on_part.ok_or(FfiError::NullPointer {
argument: "on_part",
})?;
let on_done = on_done.ok_or(FfiError::NullPointer {
argument: "on_done",
})?;
let model = model_of(handle)?;
let prompt = parse_prompt_arg(prompt_json)?;
let mut opts = parse_opts_arg(opts_json)?;
opts.abort_signal = abort_signal.clone();
let stream_ctx = stream_ctx as usize;
ffi_block_on(async move {
let sr = match abort_signal.as_ref() {
Some(signal) => {
tokio::select! {
biased;
_ = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)),
result = stream_text(&*model, prompt, opts) => result,
}
}
None => stream_text(&*model, prompt, opts).await,
}?;
let mut stream = sr.stream;
loop {
let next = match abort_signal.as_ref() {
Some(signal) => {
tokio::select! {
biased;
_ = signal.cancelled() => {
return Err(AiMuxError::from_abort_signal(signal).into());
}
item = stream.next() => item,
}
}
None => stream.next().await,
};
let Some(item) = next else {
break;
};
let cstr = match item {
Ok(part) => stream_part_cstring(&part)?,
Err(e) if e.is_recoverable_stream_error() => {
stream_part_cstring(&aimux_core::stream_part::StreamPart::Error { error: e })?
}
Err(e) => return Err(e.into()),
};
invoke_stream_callback("on_part", || {
on_part(cstr.as_ptr(), stream_ctx as *mut c_void);
})?;
}
invoke_stream_callback("on_done", || {
on_done(stream_ctx as *mut c_void);
})?;
Ok::<(), AiMuxFfiError>(())
})?
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_drop_handle(handle: u64) {
if handle != 0 {
drop_handle(handle);
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn aimux_free_string(ptr: *mut c_char) {
if ptr.is_null() {
return;
}
drop(unsafe { CString::from_raw(ptr) });
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_embedding_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_embedding_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = OpenAIProvider::new(config).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_embedding_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = CohereProvider::new(CohereConfig::new(api_key)).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_embedding_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = CohereConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = CohereProvider::new(config).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_embedding_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = GoogleProvider::new(GoogleConfig::new(api_key)).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_embedding_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = GoogleConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = GoogleProvider::new(config).embedding_model(&model_id);
Ok(intern_handle(HandleEntry::Embedding(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_embed(
handle: u64,
values_json: *const c_char,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Embedding(model) = entry_of(handle, "embedding")? else {
return Err(FfiError::InvalidHandle {
expected: "embedding",
}
.into());
};
let values_json = str_arg(values_json, "values_json")?;
let mut opts = aimux_core::embedding_model::EmbeddingCallOptions::new("");
if let Some(s) = opt_str_arg(opts_json, "opts_json")?
&& !s.trim().is_empty()
&& s.trim() != "null"
{
opts = serde_json::from_str(&s).map_err(|e| wire_err("opts_json", e))?;
}
opts.values = serde_json::from_str(&values_json).map_err(|e| wire_err("values_json", e))?;
run_json(async move { aimux_core::embedding_model::embed(model.as_ref(), opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_speech_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).speech(&model_id);
Ok(intern_handle(HandleEntry::Speech(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_speech_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = OpenAIProvider::new(config).speech(&model_id);
Ok(intern_handle(HandleEntry::Speech(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_speech_generate(
handle: u64,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Speech(model) = entry_of(handle, "speech")? else {
return Err(FfiError::InvalidHandle { expected: "speech" }.into());
};
let opts: aimux_core::speech_model::SpeechCallOptions =
parse_json_arg(opts_json, "opts_json")?;
run_json(
async move { aimux_core::speech_model::generate_speech(model.as_ref(), opts).await },
)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_image_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).image(&model_id);
Ok(intern_handle(HandleEntry::Image(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_image_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = OpenAIProvider::new(config).image(&model_id);
Ok(intern_handle(HandleEntry::Image(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_image_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = GoogleProvider::new(GoogleConfig::new(api_key)).image(&model_id);
Ok(intern_handle(HandleEntry::Image(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_image_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = GoogleConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = GoogleProvider::new(config).image(&model_id);
Ok(intern_handle(HandleEntry::Image(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_image_generate(
handle: u64,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Image(model) = entry_of(handle, "image")? else {
return Err(FfiError::InvalidHandle { expected: "image" }.into());
};
let opts: aimux_core::image_model::ImageCallOptions =
parse_json_arg(opts_json, "opts_json")?;
run_json(async move { aimux_core::image_model::generate_image(model.as_ref(), opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_transcription_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).transcription(&model_id);
Ok(intern_handle(HandleEntry::Transcription(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_transcription_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = OpenAIProvider::new(config).transcription(&model_id);
Ok(intern_handle(HandleEntry::Transcription(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_generate(
handle: u64,
audio_base64: *const c_char,
media_type: *const c_char,
_opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Transcription(model) = entry_of(handle, "transcription")? else {
return Err(FfiError::InvalidHandle {
expected: "transcription",
}
.into());
};
let audio_base64 = str_arg(audio_base64, "audio_base64")?;
let media_type = str_arg(media_type, "media_type")?;
let opts = aimux_core::transcription_model::TranscriptionCallOptions::new(
aimux_core::transcription_model::AudioInput::Base64(audio_base64),
media_type,
);
run_json(
async move { aimux_core::transcription_model::transcribe(model.as_ref(), opts).await },
)
})
}
#[derive(serde::Deserialize, Default)]
struct TranscriptionSessionFfiOptions {
input_audio_format: Option<aimux_core::transcription_model::InputAudioFormat>,
provider_options: Option<HashMap<String, serde_json::Value>>,
headers: Option<HashMap<String, String>>,
include_raw_chunks: Option<bool>,
timeout: Option<aimux_core::options::TimeoutConfiguration>,
}
fn session_of(
handle: u64,
) -> Result<Arc<transcription_session::TranscriptionFfiSession>, FfiError> {
match entry_of(handle, "transcription session")? {
HandleEntry::TranscriptionSession(s) => Ok(s),
_ => Err(FfiError::InvalidHandle {
expected: "transcription session",
}),
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_session_new(
model_handle: u64,
abort_handle: u64,
opts_json: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let HandleEntry::Transcription(model) = entry_of(model_handle, "transcription")? else {
return Err(FfiError::InvalidHandle {
expected: "transcription",
}
.into());
};
let abort = if abort_handle == 0 {
None
} else {
Some(abort_of(abort_handle)?)
};
let opts_json = normalize_config_json(opts_json, "opts_json")?;
let ffi_opts: TranscriptionSessionFfiOptions =
serde_json::from_str(&opts_json).map_err(|e| wire_err("opts_json", e))?;
let session = transcription_session::TranscriptionFfiSession::spawn(
model,
transcription_session::SessionOptions {
input_audio_format: ffi_opts.input_audio_format,
provider_options: ffi_opts.provider_options,
headers: ffi_opts.headers,
include_raw_chunks: ffi_opts.include_raw_chunks.unwrap_or(false),
timeout: ffi_opts.timeout,
},
abort,
);
Ok(intern_handle(HandleEntry::TranscriptionSession(session)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_push_audio(
session: u64,
data: *const u8,
len: usize,
) -> *mut aimux_error_t {
no_result(|| {
let session = session_of(session)?;
if data.is_null() {
if len == 0 {
return Ok(()); }
return Err(FfiError::NullPointer { argument: "data" }.into());
}
let bytes = unsafe { std::slice::from_raw_parts(data, len) }.to_vec();
ffi_block_on(session.push_audio(bytes))??;
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_input_done(session: u64) -> *mut aimux_error_t {
no_result(|| {
session_of(session)?.input_done();
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_next_part(
session: u64,
timeout_ms: i64,
out_part: *mut *mut c_char,
out_state: *mut i32,
) -> *mut aimux_error_t {
if out_part.is_null() {
return no_result(|| {
Err(FfiError::NullPointer {
argument: "out_part",
}
.into())
});
}
unsafe { *out_part = std::ptr::null_mut() };
if out_state.is_null() {
return no_result(|| {
Err(FfiError::NullPointer {
argument: "out_state",
}
.into())
});
}
let r: FfiResult<(i32, Option<String>)> = (|| {
let session = session_of(session)?;
let timeout = if timeout_ms < 0 {
None
} else {
Some(std::time::Duration::from_millis(timeout_ms as u64))
};
Ok(match ffi_block_on(session.next_part(timeout))? {
Ok(Some(Ok(part))) => (AIMUX_TRANSCRIPTION_NEXT_PART_PART, Some(to_json(&part)?)),
Ok(None) => (AIMUX_TRANSCRIPTION_NEXT_PART_ENDED, None),
Err(AiMuxError::Timeout(_)) => (AIMUX_TRANSCRIPTION_NEXT_PART_TIMEOUT, None),
Ok(Some(Err(e))) | Err(e) => return Err(e.into()),
})
})();
finish(r, |(state, part)| unsafe {
*out_state = state;
if let Some(p) = part {
*out_part = into_cstring_raw(p);
}
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_transcription_session_drop(session: u64) {
let removed = registry()
.lock()
.expect("aimux-ffi: registry mutex poisoned")
.remove(&session);
if let Some(HandleEntry::TranscriptionSession(s)) = removed {
s.terminate();
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_files_new(
api_key: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let files = OpenAIProvider::new(OpenAIConfig::new(api_key)).files();
Ok(intern_handle(HandleEntry::Files(Arc::new(files))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_openai_files_new_with_base(
api_key: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let files = OpenAIProvider::new(config).files();
Ok(intern_handle(HandleEntry::Files(Arc::new(files))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_file_upload(
handle: u64,
data_base64: *const c_char,
media_type: *const c_char,
_opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Files(model) = entry_of(handle, "files")? else {
return Err(FfiError::InvalidHandle { expected: "files" }.into());
};
let data_base64 = str_arg(data_base64, "data_base64")?;
let media_type = str_arg(media_type, "media_type")?;
let opts = aimux_core::files_model::UploadFileCallOptions::new(
aimux_core::files_model::UploadFileData::Data {
data: aimux_core::shared::FileBytes::Base64(data_base64),
},
media_type,
);
run_json(async move { model.upload_file(&opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_reranking_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = CohereProvider::new(CohereConfig::new(api_key)).reranking_model(&model_id);
Ok(intern_handle(HandleEntry::Reranking(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_cohere_reranking_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = CohereConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = CohereProvider::new(config).reranking_model(&model_id);
Ok(intern_handle(HandleEntry::Reranking(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_rerank(
handle: u64,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Reranking(model) = entry_of(handle, "reranking")? else {
return Err(FfiError::InvalidHandle {
expected: "reranking",
}
.into());
};
let opts: aimux_core::reranking_model::RerankingCallOptions =
parse_json_arg(opts_json, "opts_json")?;
run_json(async move { aimux_core::reranking_model::rerank(model.as_ref(), opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_video_new(
api_key: *const c_char,
model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let model = GoogleProvider::new(GoogleConfig::new(api_key)).video(&model_id);
Ok(intern_handle(HandleEntry::Video(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_google_video_new_with_base(
api_key: *const c_char,
model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?;
let mut config = GoogleConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = GoogleProvider::new(config).video(&model_id);
Ok(intern_handle(HandleEntry::Video(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_video_generate(
handle: u64,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Video(model) = entry_of(handle, "video")? else {
return Err(FfiError::InvalidHandle { expected: "video" }.into());
};
let opts: aimux_core::video_model::VideoCallOptions =
parse_json_arg(opts_json, "opts_json")?;
run_json(async move { aimux_core::video_model::generate_video(model.as_ref(), opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_tavily_search_new(
api_key: *const c_char,
_model_id: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let model = TavilyProvider::new(TavilyConfig::new(api_key)).search_model();
Ok(intern_handle(HandleEntry::Search(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_tavily_search_new_with_base(
api_key: *const c_char,
_model_id: *const c_char,
base_url: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let api_key = str_arg(api_key, "api_key")?;
let mut config = TavilyConfig::new(api_key);
if let Some(url) = parse_base_url(base_url)? {
config = config.with_base_url(url);
}
let model = TavilyProvider::new(config).search_model();
Ok(intern_handle(HandleEntry::Search(Arc::new(model))))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_search(
handle: u64,
opts_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let HandleEntry::Search(model) = entry_of(handle, "search")? else {
return Err(FfiError::InvalidHandle { expected: "search" }.into());
};
let opts: aimux_core::search_model::SearchCallOptions =
parse_json_arg(opts_json, "opts_json")?;
run_json(async move { aimux_core::search_model::search(model.as_ref(), opts).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_codex_refresh(
refresh_token: *const c_char,
client_id: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let (refresh_token, client_id) =
parse_two_args(refresh_token, "refresh_token", client_id, "client_id")?;
run_json(async move { aimux_providers::codex_refresh(&refresh_token, &client_id).await })
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_init_logging(level: *const c_char) -> *mut aimux_error_t {
no_result(|| {
let level = opt_str_arg(level, "level")?;
aimux_providers::init_logging(level.as_deref().unwrap_or("warn"));
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_session_store_init() {
aimux_core::session::init_session_store(std::sync::Arc::new(
aimux_core::session::SessionStore::new(),
));
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_session_infer_init(enabled: i32) {
aimux_core::session::init_session_infer(enabled != 0);
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_session_calls(
session_id: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let id = str_arg(session_id, "session_id")?;
to_json(&aimux_core::session::session_calls(&id))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_list_sessions(out_json: *mut *mut c_char) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
to_json(&aimux_core::session::list_sessions())
})
}
static TRACE_STORES: OnceLock<Mutex<HashMap<u64, Arc<RingTraceStore>>>> = OnceLock::new();
fn trace_stores() -> &'static Mutex<HashMap<u64, Arc<RingTraceStore>>> {
TRACE_STORES.get_or_init(|| Mutex::new(HashMap::new()))
}
fn trace_store_of(handle: u64) -> Result<Arc<RingTraceStore>, FfiError> {
trace_stores()
.lock()
.expect("aimux-ffi: trace registry mutex poisoned")
.get(&handle)
.cloned()
.ok_or(FfiError::InvalidHandle { expected: "trace" })
}
fn trace_wrap(handle: u64, audited: Option<bool>) -> FfiResult<u64> {
let model = model_of(handle)?;
let store = Arc::new(RingTraceStore::new());
let mut layer = TraceLayer::new(model, store.clone());
if let Some(strict) = audited {
layer = layer.with_rules_auditor(strict);
}
let new_handle = intern_model(Arc::new(layer));
trace_stores()
.lock()
.expect("aimux-ffi: trace registry mutex poisoned")
.insert(new_handle, store);
Ok(new_handle)
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_new(handle: u64, out_handle: *mut u64) -> *mut aimux_error_t {
with_out_handle(out_handle, || trace_wrap(handle, None))
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_new_audited(
handle: u64,
strict: i32,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || trace_wrap(handle, Some(strict != 0)))
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_aggregate(
handle: u64,
filter_json: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let store = trace_store_of(handle)?;
let filter = parse_json_arg::<TraceFilter>(filter_json, "filter_json")?;
to_json(&store.aggregate(&filter))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_session_chain(
handle: u64,
session_id: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let store = trace_store_of(handle)?;
let id = str_arg(session_id, "session_id")?;
let view = store
.session_chain(&id)
.ok_or_else(|| AiMuxError::InvalidArgument("unknown session".into()))?;
to_json(&view)
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_session_trajectory(
handle: u64,
session_id: *const c_char,
out_json: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_json, "out_json", || {
let store = trace_store_of(handle)?;
let id = str_arg(session_id, "session_id")?;
to_json(&store.session_cache_trajectory(&id))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_export_jsonl(
handle: u64,
out_jsonl: *mut *mut c_char,
) -> *mut aimux_error_t {
with_out_string(out_jsonl, "out_jsonl", || {
let store = trace_store_of(handle)?;
let mut buf = Vec::new();
store
.export_jsonl(&mut buf)
.map_err(|e| FfiError::ResultSerialization {
message: format!("export: {e}"),
})?;
String::from_utf8(buf).map_err(|e| {
FfiError::ResultSerialization {
message: format!("utf8: {e}"),
}
.into()
})
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_trace_clear(handle: u64) -> *mut aimux_error_t {
no_result(|| {
trace_store_of(handle)?.clear();
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_init_recording(dir: *const c_char) -> *mut aimux_error_t {
no_result(|| {
let dir = str_arg(dir, "dir")?;
let rec = aimux_core::recording::JsonlRecorder::try_new(dir)?;
aimux_core::recording::init_recording(Some(std::sync::Arc::new(rec)));
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_init_recording_ring(cap: u64) -> *mut aimux_error_t {
no_result(|| {
if cap == 0 {
return Err(AiMuxError::InvalidArgument("cap: must be > 0".into()).into());
}
aimux_core::recording::init_recording(Some(std::sync::Arc::new(
aimux_core::recording::RingRecorder::with_capacity(cap as usize),
)));
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_init_recording_ring_default() {
aimux_core::recording::init_recording(Some(std::sync::Arc::new(
aimux_core::recording::RingRecorder::default(),
)));
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_recording_stop() {
aimux_core::recording::init_recording(None);
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_recording_flush() {
if let Some(rec) = aimux_core::recording::recorder() {
rec.flush();
}
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_recording_try_flush() -> *mut aimux_error_t {
no_result(|| {
if let Some(rec) = aimux_core::recording::recorder() {
rec.try_flush()?;
}
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_register_providers(config_json: *const c_char) -> *mut aimux_error_t {
no_result(|| {
let json = str_arg(config_json, "config_json")?;
serde_json::from_str::<serde_json::Value>(&json).map_err(|e| wire_err("config_json", e))?;
aimux_providers::load_providers_from_json(&json).map_err(|e| match e {
AiMuxError::JsonParse(m) => AiMuxError::InvalidArgument(format!("config_json: {m}")),
e => e,
})?;
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_init_proxy(config_json: *const c_char) -> *mut aimux_error_t {
no_result(|| {
let config: aimux_provider_utils::ProxyConfig = parse_json_arg(config_json, "config_json")?;
let _ = aimux_provider_utils::init_proxy(config);
Ok(())
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_mock_replay_new(
recordings_jsonl: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let recordings_jsonl = str_arg(recordings_jsonl, "recordings_jsonl")?;
let mut recordings: Vec<aimux_core::recording::Recording> = Vec::new();
for (idx, line) in recordings_jsonl.lines().enumerate() {
let line = line.trim();
if line.is_empty() {
continue;
}
recordings.push(serde_json::from_str(line).map_err(|e| {
wire_failure("recordings_jsonl", &e, format!("line {}: {e}", idx + 1))
})?);
}
if recordings.is_empty() {
return Err(
AiMuxError::InvalidArgument("recordings_jsonl: no recordings".into()).into(),
);
}
let model = aimux_core::replay::MockReplayModel::new(
recordings[0].provider.provider.clone(),
recordings[0].provider.model_id.clone(),
recordings,
);
Ok(intern_model(Arc::new(model)))
})
}
fn model_handles(
handles: *const u64,
len: usize,
argument: &'static str,
expected: &'static str,
) -> FfiResult<Vec<Arc<dyn aimux_core::LanguageModel>>> {
let slice: &[u64] = if handles.is_null() {
if len != 0 {
return Err(FfiError::NullPointer { argument }.into());
}
&[]
} else {
unsafe { std::slice::from_raw_parts(handles, len) }
};
slice
.iter()
.map(|&h| get_model(h).ok_or_else(|| FfiError::InvalidHandle { expected }.into()))
.collect()
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_router_new(
handles: *const u64,
len: usize,
config_json: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let models = model_handles(handles, len, "handles", "router child")?;
if models.is_empty() {
return Err(FfiError::InvalidHandle {
expected: "router child",
}
.into());
}
let config_json = normalize_config_json(config_json, "config_json")?;
let cfg: RouterFfiConfig =
serde_json::from_str(&config_json).map_err(|e| wire_err("config_json", e))?;
let router: Box<dyn aimux_core::router::Router> = match cfg.router.as_deref() {
Some("weighted") => {
let weights = cfg.weights.unwrap_or_else(|| vec![1.0; models.len()]);
Box::new(aimux_core::router::WeightedRouter::new(weights))
}
_ => Box::new(aimux_core::router::RuleRouter),
};
let fallback = if cfg.fallback.as_deref() == Some("none") {
aimux_core::router::FallbackPolicy::None
} else {
aimux_core::router::FallbackPolicy::OnError
};
let router_cfg = aimux_core::router::RouterConfig {
provider_name: cfg.provider_name.unwrap_or_else(|| "router".into()),
model_id: cfg.model_id.unwrap_or_else(|| "router".into()),
};
let model = aimux_core::router::RouterModel::new(models, router, fallback, router_cfg);
Ok(intern_model(Arc::new(model)))
})
}
#[unsafe(no_mangle)]
pub extern "C" fn aimux_moa_new(
reference_handles: *const u64,
ref_len: usize,
aggregator: u64,
config_json: *const c_char,
out_handle: *mut u64,
) -> *mut aimux_error_t {
with_out_handle(out_handle, || {
let aggregator_model = get_model(aggregator).ok_or(FfiError::InvalidHandle {
expected: "moa aggregator",
})?;
let references = model_handles(
reference_handles,
ref_len,
"reference_handles",
"moa reference",
)?;
let config_json = normalize_config_json(config_json, "config_json")?;
let cfg: aimux_core::moa::MoaConfig =
serde_json::from_str(&config_json).map_err(|e| wire_err("config_json", e))?;
let model = aimux_core::moa::MoaModel::new(references, aggregator_model, cfg);
Ok(intern_model(Arc::new(model)))
})
}
#[derive(Default, serde::Deserialize)]
struct RouterFfiConfig {
router: Option<String>,
weights: Option<Vec<f64>>,
fallback: Option<String>,
provider_name: Option<String>,
model_id: Option<String>,
}
#[cfg(test)]
mod tests {
use aimux_core::{ApiCallError, RetryError, RetryErrorReason};
use std::time::Duration;
use super::*;
fn take(p: *mut c_char) -> Option<String> {
if p.is_null() {
return None;
}
let s = unsafe { CStr::from_ptr(p) }.to_str().unwrap().to_string();
unsafe { aimux_free_string(p) };
Some(s)
}
fn boxed(p: impl Into<AiMuxFfiError>) -> *mut aimux_error_t {
finish::<()>(Err(p.into()), |()| {})
}
fn msg(e: *mut aimux_error_t) -> String {
assert!(!e.is_null(), "expected a returned error");
let m = take(aimux_error_message(e)).unwrap();
aimux_error_free(e);
m
}
fn expect_aimux_error(e: *mut aimux_error_t) -> (i32, String) {
assert!(!e.is_null(), "expected a returned error");
let code = aimux_error_code(e);
if !(AIMUX_E_OTHER..=AIMUX_E_TOOL_CALL_REPAIR).contains(&code) {
panic!("expected an AiMuxError code, got {code}: {}", msg(e));
}
let out = (code, take(aimux_error_message(e)).unwrap());
aimux_error_free(e);
out
}
fn expect_ffi_error(e: *mut aimux_error_t) -> String {
assert!(!e.is_null(), "expected a returned error");
let code = aimux_error_code(e);
assert!(
(AIMUX_E_FFI_NULL_POINTER..=AIMUX_E_FFI_CALLBACK_FAILURE).contains(&code),
"expected a C ABI failure code, got {code}"
);
let m = msg(e);
assert!(!m.is_empty());
m
}
fn zero_err() -> u64 {
0
}
#[test]
fn payload_getters_follow_the_code() {
let owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError {
status_code: Some(429),
provider_code: Some("insufficient_quota".into()),
response_body: Some("{}".into()),
response_headers: Some(std::collections::HashMap::from([(
"retry-after-ms".to_string(),
"1500".to_string(),
)])),
data: Some(serde_json::json!({"code": "insufficient_quota"})),
is_retryable: true,
..ApiCallError::new(
"quota",
"https://api.example.test/v1",
serde_json::json!({"model": "m"}),
)
})));
let h = owner;
assert_eq!(aimux_error_code(h), AIMUX_E_API_CALL);
assert_eq!(aimux_error_status(h), 429);
assert_eq!(aimux_error_retry_ms(h), 1500);
assert_eq!(aimux_error_retryable(h), 1);
assert_eq!(
take(aimux_error_provider_code(h)).as_deref(),
Some("insufficient_quota")
);
assert_eq!(
take(aimux_error_provider_message(h)).as_deref(),
Some("quota")
);
let m = take(aimux_error_message(h)).unwrap();
assert!(m.contains("quota") && m != "quota", "{m}");
assert_eq!(take(aimux_error_response_body(h)).as_deref(), Some("{}"));
assert_eq!(
take(aimux_error_url(h)).as_deref(),
Some("https://api.example.test/v1")
);
assert_eq!(
take(aimux_error_request_body_values(h)).as_deref(),
Some(r#"{"model":"m"}"#)
);
assert_eq!(
take(aimux_error_response_headers(h)).as_deref(),
Some(r#"{"retry-after-ms":"1500"}"#)
);
assert_eq!(
take(aimux_error_provider_data(h)).as_deref(),
Some(r#"{"code":"insufficient_quota"}"#)
);
assert!(aimux_error_model_id(h).is_null());
assert!(aimux_error_model_type(h).is_null());
assert!(aimux_error_provider_id(h).is_null());
assert!(aimux_error_retry_reason(h).is_null());
assert_eq!(aimux_error_retry_count(h), 0);
assert!(aimux_error_retry_error_at(h, 0).is_null());
aimux_error_free(owner);
let owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError::new(
"x",
"https://api.example.test/v1",
serde_json::Value::Null,
))));
let h = owner;
assert!(aimux_error_provider_code(h).is_null());
assert!(aimux_error_response_body(h).is_null());
assert!(aimux_error_response_headers(h).is_null());
assert!(aimux_error_provider_data(h).is_null());
assert!(aimux_error_request_body_values(h).is_null());
assert_eq!(take(aimux_error_provider_message(h)).as_deref(), Some("x"));
aimux_error_free(owner);
let owner = boxed(AiMuxError::Retry(RetryError {
reason: RetryErrorReason::MaxRetriesExceeded,
errors: vec![
AiMuxError::Other("first".into()),
AiMuxError::ApiCall(Box::new(ApiCallError {
status_code: Some(500),
is_retryable: true,
..ApiCallError::new(
"boom",
"https://api.example.test/v1",
serde_json::Value::Null,
)
})),
],
}));
let h = owner;
assert_eq!(aimux_error_code(h), AIMUX_E_RETRY);
assert_eq!(
take(aimux_error_retry_reason(h)).as_deref(),
Some("maxRetriesExceeded")
);
assert_eq!(aimux_error_retry_count(h), 2);
assert_eq!(aimux_error_status(h), -1);
assert_eq!(aimux_error_retryable(h), 0);
assert!(aimux_error_provider_code(h).is_null());
let m = take(aimux_error_message(h)).unwrap();
assert!(m.contains("Failed after 2 attempts"), "{m}");
let first = aimux_error_retry_error_at(h, 0);
assert_eq!(aimux_error_code(first), AIMUX_E_OTHER);
assert_eq!(take(aimux_error_message(first)).as_deref(), Some("first"));
aimux_error_free(first);
let last = aimux_error_retry_error_at(h, 1);
assert_eq!(aimux_error_code(last), AIMUX_E_API_CALL);
assert_eq!(aimux_error_status(last), 500);
assert_eq!(aimux_error_retryable(last), 1);
aimux_error_free(owner);
assert_eq!(
take(aimux_error_provider_message(last)).as_deref(),
Some("boom")
);
aimux_error_free(last);
let owner = boxed(AiMuxError::Retry(RetryError {
reason: RetryErrorReason::ErrorNotRetryable,
errors: vec![AiMuxError::Other("bad".into())],
}));
let h = owner;
assert_eq!(
take(aimux_error_retry_reason(h)).as_deref(),
Some("errorNotRetryable")
);
assert!(aimux_error_retry_error_at(h, 1).is_null(), "out of range");
assert!(aimux_error_retry_error_at(h, -1).is_null());
aimux_error_free(owner);
let owner = boxed(AiMuxError::NoSuchModel {
model_id: "m".into(),
model_type: "language".into(),
});
let h = owner;
assert_eq!(aimux_error_code(h), AIMUX_E_NO_SUCH_MODEL);
assert_eq!(take(aimux_error_model_id(h)).as_deref(), Some("m"));
assert_eq!(take(aimux_error_model_type(h)).as_deref(), Some("language"));
assert!(aimux_error_provider_code(h).is_null());
assert!(aimux_error_provider_id(h).is_null());
assert_eq!(aimux_error_status(h), -1);
assert_eq!(aimux_error_retryable(h), 0);
aimux_error_free(owner);
let owner = boxed(AiMuxError::NoSuchProvider {
provider_id: "p".into(),
});
let h = owner;
assert_eq!(aimux_error_code(h), AIMUX_E_NO_SUCH_PROVIDER);
assert_eq!(take(aimux_error_provider_id(h)).as_deref(), Some("p"));
assert!(aimux_error_model_id(h).is_null());
aimux_error_free(owner);
let n = std::ptr::null();
assert_eq!(aimux_error_code(n), AIMUX_OK);
assert!(aimux_error_message(n).is_null());
assert_eq!(aimux_error_status(n), -1);
assert_eq!(aimux_error_retry_ms(n), -1);
assert_eq!(aimux_error_retryable(n), 0);
assert!(aimux_error_provider_code(n).is_null());
assert!(aimux_error_url(n).is_null());
assert!(aimux_error_request_body_values(n).is_null());
assert!(aimux_error_response_headers(n).is_null());
assert!(aimux_error_provider_data(n).is_null());
assert!(aimux_error_retry_reason(n).is_null());
assert_eq!(aimux_error_retry_count(n), 0);
assert!(aimux_error_retry_error_at(n, 0).is_null());
}
#[test]
fn retryable_crosses_the_abi_and_status_cannot_stand_in() {
let transport_owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError {
is_retryable: true,
..ApiCallError::new(
"connection reset",
"https://api.example.test/v1",
serde_json::Value::Null,
)
})));
let transport = transport_owner;
assert_eq!(aimux_error_status(transport), -1);
assert_eq!(aimux_error_retryable(transport), 1);
aimux_error_free(transport_owner);
let no_key_owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError::new(
"no api key",
"https://api.example.test/v1",
serde_json::Value::Null,
))));
let no_key = no_key_owner;
assert_eq!(
aimux_error_status(no_key),
-1,
"same sentinel as the transport failure"
);
assert_eq!(aimux_error_retryable(no_key), 0, "but the opposite verdict");
aimux_error_free(no_key_owner);
let arg_owner = boxed(AiMuxError::InvalidArgument("bad".into()));
let arg = arg_owner;
assert_eq!(aimux_error_retryable(arg), 0);
aimux_error_free(arg_owner);
}
#[test]
fn unified_error_codes_cover_every_source() {
let e = boxed(AiMuxError::Aborted("request aborted".into()));
assert_eq!(aimux_error_code(e), AIMUX_E_ABORTED);
assert_eq!(msg(e), "request aborted");
let e = boxed(AiMuxError::Retry(RetryError {
reason: RetryErrorReason::ErrorNotRetryable,
errors: vec![AiMuxError::Other("bad".into())],
}));
assert_eq!(aimux_error_code(e), AIMUX_E_RETRY);
assert!(msg(e).contains("non-retryable"));
use RecordingError as R;
let cases = [
(R::WriterGone, AIMUX_E_RECORDING_WRITER_GONE),
(R::FlushTimeout, AIMUX_E_RECORDING_FLUSH_TIMEOUT),
(R::Write("ENOSPC".into()), AIMUX_E_RECORDING_WRITE),
(
R::Spawn {
source: std::io::Error::other("x"),
},
AIMUX_E_RECORDING_SPAWN,
),
];
for (r, code) in cases {
let e = boxed(r);
assert_eq!(aimux_error_code(e), code);
assert!(!msg(e).is_empty());
}
assert_eq!(aimux_error_code(std::ptr::null()), AIMUX_OK);
assert!(aimux_error_message(std::ptr::null()).is_null());
let cases: Vec<(FfiError, i32, &str)> = vec![
(
FfiError::NullPointer {
argument: "api_key",
},
AIMUX_E_FFI_NULL_POINTER,
"api_key: must not be NULL",
),
(
FfiError::InvalidUtf8 {
argument: "model_id",
},
AIMUX_E_FFI_INVALID_UTF8,
"model_id: must be valid UTF-8",
),
(
FfiError::InvalidWireJson {
argument: "opts_json",
message: "expected `,`".into(),
},
AIMUX_E_FFI_INVALID_WIRE_JSON,
"opts_json: invalid JSON: expected `,`",
),
(
FfiError::InvalidHandle { expected: "model" },
AIMUX_E_FFI_INVALID_HANDLE,
"invalid or expired model handle",
),
(
FfiError::ReentrantCall,
AIMUX_E_FFI_REENTRANT_CALL,
"re-entrant FFI call from within a callback is not allowed",
),
(
FfiError::ResultSerialization {
message: "serialize".into(),
},
AIMUX_E_FFI_RESULT_SERIALIZATION,
"could not serialize result: serialize",
),
(
FfiError::CallbackFailure {
message: "on_part panicked".into(),
},
AIMUX_E_FFI_CALLBACK_FAILURE,
"host callback failed: on_part panicked",
),
];
for (f, code, text) in cases {
let e = boxed(f);
assert_eq!(aimux_error_code(e), code);
assert_eq!(expect_ffi_error(e), text);
}
let n: *mut aimux_error_t = std::ptr::null_mut();
assert_eq!(aimux_error_code(n), AIMUX_OK);
assert!(aimux_error_message(n).is_null());
aimux_error_free(n);
}
#[test]
fn recording_init_reports_ffi_codes_for_invalid_input() {
assert_eq!(
expect_ffi_error(aimux_init_recording(std::ptr::null())),
"dir: must not be NULL"
);
let invalid_utf8 = [0xff_u8, 0];
assert_eq!(
expect_ffi_error(aimux_init_recording(invalid_utf8.as_ptr().cast())),
"dir: must be valid UTF-8"
);
}
#[test]
fn error_code_mapping_covers_all_variants() {
let s = |t: &str| t.to_string();
let cases: Vec<(AiMuxError, i32)> = vec![
(
AiMuxError::ApiCall(Box::new(ApiCallError::new(
"x",
"https://example.test",
serde_json::json!({}),
))),
AIMUX_E_API_CALL,
),
(
AiMuxError::Retry(RetryError {
reason: RetryErrorReason::MaxRetriesExceeded,
errors: vec![AiMuxError::Other(s("first")), AiMuxError::Other(s("last"))],
}),
AIMUX_E_RETRY,
),
(AiMuxError::JsonParse(s("x")), AIMUX_E_JSON_PARSE),
(
AiMuxError::InvalidResponseData(s("x")),
AIMUX_E_INVALID_RESPONSE_DATA,
),
(
AiMuxError::NoSuchTool {
tool_name: s("t"),
available_tools: None,
tool_input: None,
},
AIMUX_E_NO_SUCH_TOOL,
),
(
AiMuxError::InvalidToolInput {
tool_name: s("t"),
tool_input: s("{}"),
cause: s("x"),
},
AIMUX_E_INVALID_TOOL_INPUT,
),
(
AiMuxError::ToolCallRepair {
original_error: Box::new(AiMuxError::NoSuchTool {
tool_name: s("t"),
available_tools: None,
tool_input: None,
}),
cause: Box::new(AiMuxError::Other(s("x"))),
},
AIMUX_E_TOOL_CALL_REPAIR,
),
(
AiMuxError::InvalidArgument(s("x")),
AIMUX_E_INVALID_ARGUMENT,
),
(AiMuxError::InvalidPrompt(s("x")), AIMUX_E_INVALID_PROMPT),
(AiMuxError::TokenExpired(s("x")), AIMUX_E_TOKEN_EXPIRED),
(
AiMuxError::UnsupportedFunctionality(s("x")),
AIMUX_E_UNSUPPORTED_FUNCTIONALITY,
),
(
AiMuxError::NoSuchModel {
model_id: s("x"),
model_type: String::new(),
},
AIMUX_E_NO_SUCH_MODEL,
),
(
AiMuxError::NoSuchProvider {
provider_id: s("x"),
},
AIMUX_E_NO_SUCH_PROVIDER,
),
(AiMuxError::Timeout(s("x")), AIMUX_E_TIMEOUT),
(
AiMuxError::Aborted("request aborted".into()),
AIMUX_E_ABORTED,
),
(AiMuxError::Other(s("x")), AIMUX_E_OTHER),
];
for (e, code) in cases {
let expect = format!("variant {e:?}");
let (got, _) = expect_aimux_error(boxed(e));
assert_eq!(got, code, "{expect}");
}
}
#[test]
fn unknown_tool_error_exposes_original_argument_text() {
let raw = r#""hello""#;
let error = boxed(AiMuxError::NoSuchTool {
tool_name: "unknown".into(),
available_tools: None,
tool_input: Some(raw.into()),
});
let text = aimux_error_tool_input(error);
assert!(!text.is_null());
assert_eq!(unsafe { CStr::from_ptr(text) }.to_str().unwrap(), raw);
unsafe { aimux_free_string(text) };
aimux_error_free(error);
}
#[test]
fn error_message_sanitizes_interior_nul() {
let (_, m) = expect_aimux_error(boxed(AiMuxError::Other("a\0b".into())));
assert_eq!(m, "a\u{FFFD}b");
}
#[test]
fn null_out_param_is_an_ffi_error() {
let key = std::ffi::CString::new("k").unwrap();
assert_eq!(
expect_ffi_error(aimux_openai_new(
key.as_ptr(),
key.as_ptr(),
std::ptr::null_mut()
)),
"out_handle: must not be NULL"
);
assert_eq!(
expect_ffi_error(aimux_list_sessions(std::ptr::null_mut())),
"out_json: must not be NULL"
);
}
#[test]
fn stream_callback_ok_passthrough() {
let called = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let called_clone = called.clone();
let result = invoke_stream_callback("on_part", || {
called_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
});
assert!(result.is_ok());
assert_eq!(
called.load(std::sync::atomic::Ordering::SeqCst),
1,
"callback must have been invoked"
);
}
#[test]
fn stream_callback_catches_str_panic() {
let result = invoke_stream_callback("on_part", || {
panic!("callback explosion");
});
match result {
Err(FfiError::CallbackFailure { message: msg }) => {
assert!(
msg.contains("on_part"),
"message should name the callback: {msg}"
);
assert!(
msg.contains("callback explosion"),
"message should include the panic payload: {msg}"
);
}
other => panic!("expected FfiError::CallbackFailure, got {other:?}"),
}
}
#[test]
fn stream_callback_catches_string_panic() {
let result = invoke_stream_callback("on_done", || {
panic!("{}", "dynamic boom".to_string());
});
match result {
Err(FfiError::CallbackFailure { message: msg }) => {
assert!(
msg.contains("on_done"),
"message should name the callback: {msg}"
);
assert!(
msg.contains("dynamic boom"),
"message should include the panic payload: {msg}"
);
}
other => panic!("expected FfiError::CallbackFailure, got {other:?}"),
}
}
#[test]
fn stream_callback_catches_non_string_panic() {
let result = invoke_stream_callback("on_part", || {
std::panic::panic_any(42i32);
});
match result {
Err(FfiError::CallbackFailure { message: msg }) => {
assert!(
msg.contains("on_part"),
"message should name the callback: {msg}"
);
assert!(
msg.contains("<non-string panic>"),
"non-string payload should use placeholder: {msg}"
);
}
other => panic!("expected FfiError::CallbackFailure, got {other:?}"),
}
}
#[test]
fn stream_callback_catches_on_done_panic() {
let result = invoke_stream_callback("on_done", || {
panic!("done callback failed");
});
match result {
Err(FfiError::CallbackFailure { message: msg }) => {
assert!(
msg.contains("on_done"),
"message should name the callback: {msg}"
);
assert!(
msg.contains("done callback failed"),
"message should include the panic payload: {msg}"
);
}
other => panic!("expected FfiError::CallbackFailure, got {other:?}"),
}
}
struct MockText {
provider: &'static str,
model_id: &'static str,
text: &'static str,
}
#[async_trait::async_trait]
impl aimux_core::LanguageModel for MockText {
fn provider(&self) -> &str {
self.provider
}
fn model_id(&self) -> &str {
self.model_id
}
async fn do_generate(
&self,
_options: &aimux_core::options::CallOptions,
) -> Result<aimux_core::result::GenerateResult, AiMuxError> {
Ok(aimux_core::result::GenerateResult {
content: vec![aimux_core::result::GenerateContent::Text {
text: self.text.into(),
provider_metadata: None,
}],
finish_reason: aimux_core::types::FinishReason {
unified: aimux_core::types::FinishReasonUnified::Stop,
raw: None,
},
usage: aimux_core::types::Usage::default(),
warnings: vec![],
provider_metadata: None,
response: aimux_core::types::ResponseMetadata {
model_id: Some(self.model_id.into()),
..Default::default()
},
request_body: None,
response_headers: None,
})
}
async fn do_stream(
&self,
_options: &aimux_core::options::CallOptions,
) -> Result<aimux_core::result::StreamResult, AiMuxError> {
unimplemented!()
}
}
fn mock_handle(provider: &'static str, model_id: &'static str, text: &'static str) -> u64 {
intern_model(Arc::new(MockText {
provider,
model_id,
text,
}))
}
struct RecoverableFrameModel;
#[async_trait::async_trait]
impl aimux_core::LanguageModel for RecoverableFrameModel {
fn provider(&self) -> &str {
"mock"
}
fn model_id(&self) -> &str {
"recoverable"
}
async fn do_generate(
&self,
_options: &aimux_core::options::CallOptions,
) -> Result<aimux_core::result::GenerateResult, AiMuxError> {
unimplemented!()
}
async fn do_stream(
&self,
_options: &aimux_core::options::CallOptions,
) -> Result<aimux_core::result::StreamResult, AiMuxError> {
let delta = |text: &str| aimux_core::stream_part::StreamPart::TextDelta {
id: "1".into(),
delta: text.into(),
provider_metadata: None,
};
Ok(aimux_core::result::StreamResult {
stream: Box::pin(futures::stream::iter([
Ok(delta("a")),
Err(AiMuxError::JsonParse("bad frame".into())),
Ok(delta("b")),
Ok(aimux_core::stream_part::StreamPart::Finish {
finish_reason: aimux_core::types::FinishReason {
unified: aimux_core::types::FinishReasonUnified::Stop,
raw: None,
},
usage: aimux_core::types::Usage::default(),
provider_metadata: None,
}),
])),
request_body: None,
response_headers: None,
})
}
}
#[test]
fn recoverable_frame_error_does_not_end_the_ffi_stream() {
struct Collected {
parts: Vec<String>,
done: bool,
}
extern "C-unwind" fn on_part(json: *const c_char, ctx: *mut c_void) {
let collected = unsafe { &mut *(ctx as *mut Collected) };
let text = unsafe { std::ffi::CStr::from_ptr(json) }
.to_string_lossy()
.into_owned();
collected.parts.push(text);
}
extern "C-unwind" fn on_done(ctx: *mut c_void) {
let collected = unsafe { &mut *(ctx as *mut Collected) };
collected.done = true;
}
let handle = intern_model(Arc::new(RecoverableFrameModel));
let mut collected = Collected {
parts: Vec::new(),
done: false,
};
let prompt = std::ffi::CString::new("\"hi\"").unwrap();
let e = aimux_stream_text(
handle,
prompt.as_ptr(),
std::ptr::null(),
Some(on_part),
Some(on_done),
&mut collected as *mut Collected as *mut c_void,
);
assert!(e.is_null(), "{}", msg(e));
assert!(
collected.done,
"on_done must fire after a recoverable frame"
);
let error_index = collected
.parts
.iter()
.position(|p| p.contains("\"Error\""))
.expect("the recoverable frame arrives as a StreamPart::Error part");
assert!(
collected.parts[error_index + 1..]
.iter()
.any(|p| p.contains("\"b\"")),
"parts after the recoverable frame still arrive: {:?}",
collected.parts
);
}
#[test]
fn recoverable_frame_error_is_skipped_on_the_openai_stream() {
struct Collected {
parts: Vec<String>,
done: bool,
}
extern "C-unwind" fn on_part(json: *const c_char, ctx: *mut c_void) {
let collected = unsafe { &mut *(ctx as *mut Collected) };
let text = unsafe { std::ffi::CStr::from_ptr(json) }
.to_string_lossy()
.into_owned();
collected.parts.push(text);
}
extern "C-unwind" fn on_done(ctx: *mut c_void) {
let collected = unsafe { &mut *(ctx as *mut Collected) };
collected.done = true;
}
let handle = intern_model(Arc::new(RecoverableFrameModel));
let mut collected = Collected {
parts: Vec::new(),
done: false,
};
let prompt = std::ffi::CString::new("\"hi\"").unwrap();
let e = aimux_stream_text_as_openai(
handle,
prompt.as_ptr(),
std::ptr::null(),
Some(on_part),
Some(on_done),
&mut collected as *mut Collected as *mut c_void,
);
assert!(e.is_null(), "{}", msg(e));
assert!(
collected.done,
"on_done must fire after a recoverable frame"
);
assert!(
collected.parts.iter().all(|p| !p.contains("\"error\"")),
"no error item may ride the chunk-typed wire: {:?}",
collected.parts
);
assert!(
collected.parts.iter().any(|p| p.contains("\"b\"")),
"chunks after the recoverable frame still arrive: {:?}",
collected.parts
);
}
#[test]
fn router_new_builds_router_model() {
let handles = [
mock_handle("mock", "primary", "primary-out"),
mock_handle("mock", "backup", "backup-out"),
];
let mut h = zero_err();
let e = aimux_router_new(handles.as_ptr(), handles.len(), std::ptr::null(), &mut h);
assert!(e.is_null(), "{}", msg(e));
assert!(h != 0, "router handle must be non-zero");
let model = get_model(h).expect("router handle resolves");
assert_eq!(model.provider(), "router");
assert_eq!(model.model_id(), "router");
}
#[test]
fn router_new_rejects_empty_handles() {
let mut h = zero_err();
let e = aimux_router_new(std::ptr::null(), 0, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(
expect_ffi_error(e),
"invalid or expired router child handle"
);
let e = aimux_router_new(std::ptr::null(), 2, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(expect_ffi_error(e), "handles: must not be NULL");
}
#[test]
fn router_new_rejects_bad_json() {
let handles = [mock_handle("mock", "m", "x")];
let bad = std::ffi::CString::new("{not json").unwrap();
let mut h = zero_err();
let e = aimux_router_new(handles.as_ptr(), handles.len(), bad.as_ptr(), &mut h);
assert_eq!(h, 0);
assert!(expect_ffi_error(e).starts_with("config_json: invalid JSON:"));
}
#[test]
fn router_new_treats_empty_config_as_defaults() {
let handles = [mock_handle("mock", "m", "x")];
for cfg in ["", " ", "null"] {
let c = std::ffi::CString::new(cfg).unwrap();
let mut h = zero_err();
let e = aimux_router_new(handles.as_ptr(), handles.len(), c.as_ptr(), &mut h);
assert!(e.is_null(), "router with config {cfg:?}: {}", msg(e));
assert!(h != 0);
}
}
#[test]
fn moa_new_builds_moa_model() {
let refs = [
mock_handle("mock", "ref-a", "A"),
mock_handle("mock", "ref-b", "B"),
];
let agg = mock_handle("mock", "aggregator", "agg");
let mut h = zero_err();
let e = aimux_moa_new(refs.as_ptr(), refs.len(), agg, std::ptr::null(), &mut h);
assert!(e.is_null(), "{}", msg(e));
assert!(h != 0, "moa handle must be non-zero");
let model = get_model(h).expect("moa handle resolves");
assert_eq!(model.provider(), "moa");
assert_eq!(model.model_id(), "moa");
}
#[test]
fn moa_new_rejects_bad_aggregator() {
let refs = [mock_handle("mock", "ref-a", "A")];
let mut h = zero_err();
let e = aimux_moa_new(refs.as_ptr(), refs.len(), 999_999, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(
expect_ffi_error(e),
"invalid or expired moa aggregator handle"
);
}
#[test]
fn composites_reject_any_dead_member_handle() {
let live = mock_handle("mock", "live", "L");
let dead = mock_handle("mock", "dead", "D");
aimux_drop_handle(dead);
let handles = [live, dead];
let mut h = zero_err();
let e = aimux_router_new(handles.as_ptr(), handles.len(), std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(
expect_ffi_error(e),
"invalid or expired router child handle"
);
let agg = mock_handle("mock", "aggregator", "agg");
let e = aimux_moa_new(
handles.as_ptr(),
handles.len(),
agg,
std::ptr::null(),
&mut h,
);
assert_eq!(h, 0);
assert_eq!(
expect_ffi_error(e),
"invalid or expired moa reference handle"
);
let e = aimux_moa_new(std::ptr::null(), 2, agg, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(expect_ffi_error(e), "reference_handles: must not be NULL");
}
#[test]
fn moa_new_allows_zero_references() {
let agg = mock_handle("mock", "aggregator", "agg");
let mut h = zero_err();
let e = aimux_moa_new(std::ptr::null(), 0, agg, std::ptr::null(), &mut h);
assert!(e.is_null(), "{}", msg(e));
assert!(h != 0, "moa with 0 references should succeed");
}
struct MockStreamingTranscriber;
#[async_trait::async_trait]
impl aimux_core::transcription_model::TranscriptionModel for MockStreamingTranscriber {
fn provider(&self) -> &str {
"mock"
}
fn model_id(&self) -> &str {
"mock-stream-stt"
}
async fn do_generate(
&self,
_options: &aimux_core::transcription_model::TranscriptionCallOptions,
) -> Result<aimux_core::transcription_model::TranscriptionResult, AiMuxError> {
unimplemented!()
}
async fn do_stream(
&self,
options: aimux_core::transcription_model::TranscriptionStreamOptions,
) -> Result<aimux_core::transcription_model::TranscriptionStreamResult, AiMuxError>
{
use aimux_core::transcription_model::{
AudioChunk, TranscriptionStreamPart, TranscriptionStreamResult,
};
use futures::StreamExt;
let mut audio = options.audio;
let abort = options.abort_signal.clone();
let stream = async_stream::stream! {
yield Ok(TranscriptionStreamPart::StreamStart { warnings: vec![] });
let mut text = String::new();
loop {
let aborted = async {
match &abort {
Some(s) => s.cancelled().await,
None => std::future::pending().await,
}
};
tokio::select! {
_ = aborted => {
yield Err(AiMuxError::Aborted("request aborted".into()));
break;
}
chunk = audio.next() => match chunk {
Some(AudioChunk::Binary(b)) => {
let s = format!("{} ", b.len());
text.push_str(&s);
yield Ok(TranscriptionStreamPart::TranscriptDelta {
id: None,
delta: s,
provider_metadata: None,
});
}
Some(AudioChunk::Base64(s)) => {
text.push_str(&s);
yield Ok(TranscriptionStreamPart::TranscriptDelta {
id: None,
delta: s,
provider_metadata: None,
});
}
None => {
yield Ok(TranscriptionStreamPart::TranscriptFinal {
id: None,
text: text.clone(),
start_second: None,
end_second: None,
channel_index: None,
provider_metadata: None,
});
yield Ok(TranscriptionStreamPart::Finish {
text,
segments: vec![],
language: None,
duration_in_seconds: None,
provider_metadata: None,
});
break;
}
},
}
}
};
Ok(TranscriptionStreamResult {
stream: Box::pin(stream),
request: None,
response: None,
})
}
}
fn mock_transcriber_handle() -> u64 {
intern_handle(HandleEntry::Transcription(Arc::new(
MockStreamingTranscriber,
)))
}
fn session(model: u64, abort: u64) -> u64 {
let mut s = zero_err();
let e = aimux_transcription_session_new(model, abort, std::ptr::null(), &mut s);
assert!(e.is_null(), "{}", msg(e));
assert!(s != 0, "session should start");
s
}
fn pull(session: u64, timeout_ms: i64) -> Result<(i32, Option<String>), *mut aimux_error_t> {
let mut part: *mut c_char = std::ptr::null_mut();
let mut state: i32 = 0;
let e = aimux_transcription_next_part(session, timeout_ms, &mut part, &mut state);
if !e.is_null() {
assert!(part.is_null(), "failure must leave *out_part NULL");
return Err(e);
}
Ok((state, take(part)))
}
fn pull_part(session: u64) -> String {
match pull(session, 2_000) {
Ok((AIMUX_TRANSCRIPTION_NEXT_PART_PART, Some(p))) => p,
Ok(other) => panic!("expected a part, got {other:?}"),
Err(e) => panic!("expected a part, got error: {}", msg(e)),
}
}
#[test]
fn transcription_session_full_lifecycle() {
let model = mock_transcriber_handle();
let session = session(model, 0);
let chunk1 = [1u8, 2, 3];
let chunk2 = [4u8, 5];
assert!(aimux_transcription_push_audio(session, chunk1.as_ptr(), 3).is_null());
assert!(aimux_transcription_push_audio(session, chunk2.as_ptr(), 2).is_null());
assert!(aimux_transcription_input_done(session).is_null());
let p1 = pull_part(session);
assert!(p1.contains("StreamStart"), "part 1: {p1}");
let p2 = pull_part(session);
assert!(p2.contains("3 "), "part 2: {p2}");
let p3 = pull_part(session);
assert!(p3.contains("2 "), "part 3: {p3}");
let p4 = pull_part(session);
assert!(p4.contains("TranscriptFinal"), "part 4: {p4}");
let p5 = pull_part(session);
assert!(p5.contains("Finish"), "part 5: {p5}");
assert!(p5.contains("3 2 "), "finish text: {p5}");
assert_eq!(
pull(session, 2_000).ok(),
Some((AIMUX_TRANSCRIPTION_NEXT_PART_ENDED, None))
);
aimux_transcription_session_drop(session);
}
#[test]
fn transcription_session_timeout_is_a_state_not_an_error() {
let model = mock_transcriber_handle();
let session = session(model, 0);
let _ = pull_part(session);
assert_eq!(
pull(session, 50).ok(),
Some((AIMUX_TRANSCRIPTION_NEXT_PART_TIMEOUT, None))
);
let chunk = [9u8];
assert!(aimux_transcription_push_audio(session, chunk.as_ptr(), 1).is_null());
assert!(aimux_transcription_input_done(session).is_null());
let p = pull_part(session);
assert!(p.contains("1 "), "post-timeout part: {p}");
aimux_transcription_session_drop(session);
}
#[test]
fn transcription_session_abort_via_signal() {
let model = mock_transcriber_handle();
let abort = crate::aimux_abort_signal_new();
let session = session(model, abort);
let _ = pull_part(session);
crate::aimux_abort_signal_abort(abort);
let e = pull(session, 2_000).expect_err("aborted stream must fail");
assert_eq!(expect_aimux_error(e).0, AIMUX_E_ABORTED);
let chunk = [1u8];
let e = aimux_transcription_push_audio(session, chunk.as_ptr(), 1);
aimux_error_free(e);
aimux_transcription_session_drop(session);
}
#[test]
fn transcription_session_drop_is_safe_and_idempotent() {
let model = mock_transcriber_handle();
let session = session(model, 0);
aimux_transcription_session_drop(session);
aimux_transcription_session_drop(session); aimux_transcription_session_drop(0); aimux_transcription_session_drop(999_999);
assert_eq!(
expect_ffi_error(aimux_transcription_push_audio(session, std::ptr::null(), 0)),
"invalid or expired transcription session handle"
);
assert_eq!(
expect_ffi_error(aimux_transcription_input_done(session)),
"invalid or expired transcription session handle"
);
let e = pull(session, 0).expect_err("dead session");
assert_eq!(
expect_ffi_error(e),
"invalid or expired transcription session handle"
);
}
#[test]
fn transcription_session_drop_unblocks_infinite_next_part() {
let model = mock_transcriber_handle();
let session = session(model, 0);
let _ = pull_part(session);
let (done_tx, done_rx) = std::sync::mpsc::channel();
let puller = std::thread::spawn(move || {
if let Err(e) = pull(session, -1) {
aimux_error_free(e);
}
done_tx.send(()).unwrap();
});
std::thread::sleep(Duration::from_millis(50));
assert!(
matches!(
done_rx.try_recv(),
Err(std::sync::mpsc::TryRecvError::Empty)
),
"next_part(-1) should still be blocked before drop"
);
let started = std::time::Instant::now();
aimux_transcription_session_drop(session);
assert!(
started.elapsed() < Duration::from_secs(2),
"session_drop should not wait for next_part"
);
done_rx
.recv_timeout(Duration::from_secs(2))
.expect("session_drop must wake next_part(-1)");
puller.join().unwrap();
drop_handle(model);
}
#[test]
fn transcription_session_drop_unblocks_backpressured_push_audio() {
let model = mock_transcriber_handle();
let session = session(model, 0);
let sent = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let sent_by_pusher = Arc::clone(&sent);
let (done_tx, done_rx) = std::sync::mpsc::channel();
let pusher = std::thread::spawn(move || {
let chunk = [1u8];
for _ in 0..1024 {
let e = aimux_transcription_push_audio(session, chunk.as_ptr(), chunk.len());
if !e.is_null() {
aimux_error_free(e);
break;
}
sent_by_pusher.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
done_tx.send(()).unwrap();
});
let deadline = std::time::Instant::now() + Duration::from_secs(2);
while sent.load(std::sync::atomic::Ordering::Relaxed) < 64
&& std::time::Instant::now() < deadline
{
std::thread::yield_now();
}
assert!(
sent.load(std::sync::atomic::Ordering::Relaxed) >= 64,
"pusher never filled the bounded audio path"
);
std::thread::sleep(Duration::from_millis(50));
assert!(
matches!(
done_rx.try_recv(),
Err(std::sync::mpsc::TryRecvError::Empty)
),
"pusher unexpectedly completed without backpressure"
);
let started = std::time::Instant::now();
aimux_transcription_session_drop(session);
assert!(
started.elapsed() < Duration::from_secs(2),
"session_drop should not wait for push_audio"
);
done_rx
.recv_timeout(Duration::from_secs(2))
.expect("session_drop must wake a backpressured push_audio");
pusher.join().unwrap();
drop_handle(model);
}
#[test]
fn transcription_session_new_rejects_bad_inputs() {
let model = mock_transcriber_handle();
let mut h = zero_err();
let bad = std::ffi::CString::new("{not json").unwrap();
let e = aimux_transcription_session_new(model, 0, bad.as_ptr(), &mut h);
assert_eq!(h, 0);
assert!(expect_ffi_error(e).starts_with("opts_json: invalid JSON:"));
let lang = mock_handle("mock", "m", "x");
let e = aimux_transcription_session_new(lang, 0, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(
expect_ffi_error(e),
"invalid or expired transcription handle"
);
let e = aimux_transcription_session_new(model, 999_999, std::ptr::null(), &mut h);
assert_eq!(h, 0);
assert_eq!(expect_ffi_error(e), "invalid or expired abort handle");
let mut part: *mut c_char = std::ptr::null_mut();
assert_eq!(
expect_ffi_error(aimux_transcription_next_part(
0,
0,
std::ptr::null_mut(),
std::ptr::null_mut()
)),
"out_part: must not be NULL"
);
assert_eq!(
expect_ffi_error(aimux_transcription_next_part(
0,
0,
&mut part,
std::ptr::null_mut()
)),
"out_state: must not be NULL"
);
}
}