diff --git a/aimux-ffi/aimux-ffi.h b/aimux-ffi/aimux-ffi.h index d223c45f..b948fae5 100644 --- a/aimux-ffi/aimux-ffi.h +++ b/aimux-ffi/aimux-ffi.h @@ -625,8 +625,9 @@ uint64_t aimux_mock_replay_new(const char *recordings_jsonl, AimuxError *err); aimux_stream_text); release with aimux_drop_handle. handles: array of `len` existing model handles (e.g. from aimux_openai_new). - Unknown handles are silently dropped; the call only fails if all are - unknown (or len == 0). config_json selects the router + fallback policy: + Lookup is all-or-nothing: any unknown handle fails the call, so the router + never runs with fewer children than requested. len == 0 also fails. + config_json selects the router + fallback policy: { "router": "rule"|"weighted", "weights": [..], "fallback": "on_error"|"none", "provider_name": "router", "model_id": "router" } — all optional; defaults are rule / on_error / "router" / "router". @@ -640,7 +641,7 @@ uint64_t aimux_router_new(const uint64_t *handles, size_t len, reference_handles: array of `ref_len` existing handles (may be NULL/0 — MoaModel then runs just the aggregator). aggregator: a single existing - handle (must be valid). Unknown reference handles are dropped; an unknown + handle (must be valid). Lookup is all-or-nothing: any unknown reference or aggregator handle fails. config_json is a serialized MoaConfig (all fields optional): { "provider_name", "model_id", "aggregator_instructions", "strip_reference_tools", "fail_mode": "best_effort"|"fail_fast" }. diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index 264ea2fa..619a7cb8 100644 --- a/aimux-ffi/src/lib.rs +++ b/aimux-ffi/src/lib.rs @@ -128,6 +128,23 @@ fn get_model(handle: u64) -> Option> { } } +/// Resolve a composite's model handles under one registry lock. This makes +/// the lookup all-or-nothing with respect to concurrent handle drops: either +/// every requested model is cloned, or the constructor fails without silently +/// changing the composite's membership. +fn get_models(handles: &[u64]) -> Option>> { + let registry = registry() + .lock() + .expect("aimux-ffi: registry mutex poisoned"); + handles + .iter() + .map(|handle| match registry.get(handle) { + Some(ModelHandle::Language(model)) => Some(model.clone()), + _ => None, + }) + .collect() +} + /// Look up a provider by handle (RFC-0027 provider handles for list_models). fn get_provider(handle: u64) -> Option> { match get_handle(handle)? { @@ -3017,9 +3034,9 @@ pub extern "C" fn aimux_mock_replay_new( /// optional; defaults are `rule` / `on_error` / `"router"` / `"router"`. /// /// Returns a non-zero handle on success, or 0 with `err` filled: null/invalid -/// pointer, bad JSON, zero-length `handles`, or an unknown child handle -/// (which is dropped from the child list — the call only fails if **all** -/// handles are unknown). +/// pointer, bad JSON, zero-length `handles`, or any unknown child handle. +/// Child lookup is all-or-nothing; unknown children are never silently +/// dropped. #[unsafe(no_mangle)] pub extern "C" fn aimux_router_new( handles: *const u64, @@ -3031,15 +3048,9 @@ pub extern "C" fn aimux_router_new( return unsafe { fail_invalid_args(err) }; } let handle_slice = unsafe { std::slice::from_raw_parts(handles, len) }; - let mut models: Vec> = Vec::with_capacity(len); - for &h in handle_slice { - if let Some(m) = get_model(h) { - models.push(m); - } - } - if models.is_empty() { - return unsafe { fail_other(err, "router: no valid child handles") }; - } + let Some(models) = get_models(handle_slice) else { + return unsafe { fail_other(err, "router: invalid child handle") }; + }; // NULL / empty / "null" all mean "defaults" (matching parse_provider_options // convention). Invalid UTF-8 → invalid args. let Some(config_json) = normalize_config_json(config_json, err) else { @@ -3084,8 +3095,9 @@ pub extern "C" fn aimux_router_new( /// "fail_mode": "best_effort" | "fail_fast" }`. /// /// Returns a non-zero handle on success, or 0 with `err` filled: null/invalid -/// pointer, bad JSON, an unknown aggregator handle, or all-unknown references -/// (references may be empty; aggregator may not). +/// pointer, bad JSON, an unknown aggregator handle, or any unknown reference +/// handle (references may be empty; aggregator may not). Model lookup is +/// all-or-nothing. #[unsafe(no_mangle)] pub extern "C" fn aimux_moa_new( reference_handles: *const u64, @@ -3094,22 +3106,21 @@ pub extern "C" fn aimux_moa_new( config_json: *const c_char, err: *mut CAimuxError, ) -> u64 { - let Some(aggregator_model) = get_model(aggregator) else { - return unsafe { fail_other(err, "moa: invalid aggregator handle") }; - }; - let references: Vec> = - if reference_handles.is_null() || ref_len == 0 { - Vec::new() - } else { - let slice = unsafe { std::slice::from_raw_parts(reference_handles, ref_len) }; - let mut v = Vec::with_capacity(ref_len); - for &h in slice { - if let Some(m) = get_model(h) { - v.push(m); - } - } - v - }; + let reference_slice: &[u64] = if ref_len == 0 { + &[] + } else if reference_handles.is_null() { + return unsafe { fail_invalid_args(err) }; + } else { + unsafe { std::slice::from_raw_parts(reference_handles, ref_len) } + }; + let mut handles = Vec::with_capacity(ref_len + 1); + handles.extend_from_slice(reference_slice); + handles.push(aggregator); + let Some(mut models) = get_models(&handles) else { + return unsafe { fail_other(err, "moa: invalid reference or aggregator handle") }; + }; + let aggregator_model = models.pop().expect("aggregator handle is always present"); + let references = models; // NULL / empty / "null" all mean "defaults" (matching parse_provider_options // convention). Invalid UTF-8 → invalid args. let Some(config_json) = normalize_config_json(config_json, err) else { @@ -3439,6 +3450,15 @@ mod tests { assert_eq!(err.code, AIMUX_E_INVALID_ARGUMENT); } + #[test] + fn router_new_rejects_any_invalid_child() { + let handles = [mock_handle("mock", "valid", "ok"), 999_999]; + let mut err = zero_err(); + let h = aimux_router_new(handles.as_ptr(), handles.len(), std::ptr::null(), &mut err); + assert_eq!(h, 0, "router must not silently drop an invalid child"); + assert_eq!(err.code, AIMUX_E_OTHER); + } + #[test] fn router_new_rejects_bad_json() { let handles = [mock_handle("mock", "m", "x")]; @@ -3505,6 +3525,16 @@ mod tests { assert_eq!(err.code, AIMUX_E_OTHER); } + #[test] + fn moa_new_rejects_any_invalid_reference() { + let refs = [mock_handle("mock", "ref-a", "A"), 999_999]; + let agg = mock_handle("mock", "aggregator", "agg"); + let mut err = zero_err(); + let h = aimux_moa_new(refs.as_ptr(), refs.len(), agg, std::ptr::null(), &mut err); + assert_eq!(h, 0, "MoA must not silently drop an invalid reference"); + assert_eq!(err.code, AIMUX_E_OTHER); + } + #[test] fn moa_new_allows_zero_references() { // 0 references is valid (degrades to aggregator-only). diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index a6e842c8..22289832 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -19,37 +19,35 @@ package aimux #include "aimux-ffi.h" -// Go-side callback trampolines (//export below). stream_ctx is the stream id. -extern void goStreamPart(int64_t id, char* json); -extern void goStreamDone(int64_t id); +// Go-side callback trampoline (//export below). stream_ctx is a cgo.Handle. +extern void goStreamPart(uintptr_t id, char* json); static void trampoline_part(const char* json, void* stream_ctx) { - int64_t id = (int64_t)(intptr_t)stream_ctx; + uintptr_t id = (uintptr_t)stream_ctx; if (id) goStreamPart(id, (char*)json); } static void trampoline_done(void* stream_ctx) { - int64_t id = (int64_t)(intptr_t)stream_ctx; - if (id) goStreamDone(id); + (void)stream_ctx; } // do_stream: blocking cancelable stream; id is passed as stream_ctx. // On failure, *err is filled and return is 0 (no on_done). static int32_t do_stream(uint64_t handle, uint64_t abort_handle, - const char* prompt, const char* opts, int64_t id, + const char* prompt, const char* opts, uintptr_t id, AimuxError* err) { return aimux_stream_text_with_abort( handle, abort_handle, prompt, opts, trampoline_part, trampoline_done, - (void*)(intptr_t)id, err); + (void*)id, err); } static int32_t do_stream_openai(uint64_t handle, uint64_t abort_handle, - const char* prompt, const char* opts, int64_t id, + const char* prompt, const char* opts, uintptr_t id, AimuxError* err) { return aimux_stream_text_as_openai_with_abort( handle, abort_handle, prompt, opts, trampoline_part, trampoline_done, - (void*)(intptr_t)id, err); + (void*)id, err); } */ @@ -61,7 +59,8 @@ import ( "fmt" "io" "runtime" - "sync" + "runtime/cgo" + "sync/atomic" "unsafe" ) @@ -72,43 +71,57 @@ import ( // It implements io.Closer — you MUST call Close (or use defer) to release the // native handle and avoid memory leaks. // -// Concurrency: Model is safe for concurrent use. GenerateText and StreamText -// acquire a read lock; Close acquires a write lock and waits for in-flight -// calls to finish before dropping the native handle. +// Model must not be copied after first use. Its methods and Close are safe to +// call concurrently. Close prevents future calls but never waits for an +// in-flight C call; Rust's registry owns the Arc used by calls that entered +// before Close won the registry race. type Model struct { - mu sync.RWMutex - handle uint64 - closed bool + id atomic.Uint64 // 0 means closed } // Close releases the native handle. Safe to call multiple times. -// It blocks until in-flight GenerateText/StreamText calls finish. +// It never waits for in-flight model calls or streams. func (m *Model) Close() error { - m.mu.Lock() - defer m.mu.Unlock() - if m.closed { + if m == nil { return nil } - m.closed = true - if m.handle != 0 { - C.aimux_drop_handle(C.uint64_t(m.handle)) - m.handle = 0 + if id := m.id.Swap(0); id != 0 { + C.aimux_drop_handle(C.uint64_t(id)) + runtime.SetFinalizer(m, nil) } - runtime.SetFinalizer(m, nil) return nil } -// acquireHandle returns the native handle under a read lock, or an error if -// the model is closed. The caller must call the returned release func when -// done with the handle (deferred after the FFI call returns). -func (m *Model) acquireHandle() (uint64, func(), error) { - m.mu.RLock() - if m.closed { - m.mu.RUnlock() - return 0, nil, newError(CodeInvalidArgument, "aimux: model already closed") +// handle snapshots the native handle. Callers must KeepAlive the receiver +// until after C returns so its finalizer cannot run before registry lookup. +func (m *Model) handle() (uint64, error) { + if m != nil { + if id := m.id.Load(); id != 0 { + return id, nil + } + } + return 0, newError(CodeInvalidArgument, "aimux: model already closed") +} + +func modelHandles(models []*Model) ([]uint64, error) { + handles := make([]uint64, len(models)) + for i, model := range models { + if model == nil { + return nil, newError(CodeInvalidArgument, fmt.Sprintf("aimux: models[%d] is nil", i)) + } + handle, err := model.handle() + if err != nil { + return nil, err + } + handles[i] = handle + } + return handles, nil +} + +func keepModelsAlive(models []*Model) { + for _, model := range models { + runtime.KeepAlive(model) } - h := m.handle - return h, m.mu.RUnlock, nil } // ── Provider constructors ─────────────────────────────────────────────────── @@ -347,15 +360,11 @@ func NewRouter(models []*Model, configJSON string) (*Model, error) { if len(models) == 0 { return nil, newError(CodeInvalidArgument, "router: models must be non-empty") } - handles := make([]uint64, len(models)) - for i, m := range models { - h, unlock, err := m.acquireHandle() - if err != nil { - return nil, err - } - handles[i] = h - unlock() + handles, err := modelHandles(models) + if err != nil { + return nil, err } + defer keepModelsAlive(models) var cerr C.AimuxError C.aimux_error_clear(&cerr) var h C.uint64_t @@ -385,37 +394,24 @@ func NewRouter(models []*Model, configJSON string) (*Model, error) { // // configJSON (optional) is a serialized MoaConfig. func NewMoa(references []*Model, aggregator *Model, configJSON string) (*Model, error) { - aggHandle, unlock, err := aggregator.acquireHandle() + if aggregator == nil { + return nil, newError(CodeInvalidArgument, "aimux: aggregator is nil") + } + models := append(append(make([]*Model, 0, len(references)+1), references...), aggregator) + handles, err := modelHandles(models) if err != nil { return nil, err } - defer unlock() + defer keepModelsAlive(models) + aggHandle := handles[len(references)] var cerr C.AimuxError C.aimux_error_clear(&cerr) - var h C.uint64_t - switch { - case len(references) == 0: - // No references: pass a NULL pointer + 0 length. - h = callMoaNew(nil, 0, aggHandle, configJSON, &cerr) - default: - refHandles := make([]uint64, len(references)) - for i, m := range references { - rh, runlock, rerr := m.acquireHandle() - if rerr != nil { - return nil, rerr - } - refHandles[i] = rh - runlock() - } - h = callMoaNew( - (*C.uint64_t)(unsafe.Pointer(&refHandles[0])), - C.size_t(len(refHandles)), - aggHandle, - configJSON, - &cerr, - ) + var refPtr *C.uint64_t + if len(references) > 0 { + refPtr = (*C.uint64_t)(unsafe.Pointer(&handles[0])) } + h := callMoaNew(refPtr, C.size_t(len(references)), aggHandle, configJSON, &cerr) return wrapHandleU64(h, &cerr) } @@ -458,7 +454,8 @@ func wrapHandleU64(h C.uint64_t, cerr *C.AimuxError) (*Model, error) { if h == 0 { return nil, errorFromC(cerr) } - m := &Model{handle: uint64(h)} + m := &Model{} + m.id.Store(uint64(h)) runtime.SetFinalizer(m, func(m *Model) { m.Close() }) return m, nil } @@ -654,54 +651,58 @@ func ProviderWithConfig(name, apiKey, modelID string, cfg *ProviderConfig) (*Mod // ProviderHandle is a provider handle created by CreateProvider. It supports // ListModels (runtime discovery) and Model (build a model from a discovered id). +// It must not be copied after first use. type ProviderHandle struct { - mu sync.RWMutex - handle uint64 - closed bool + id atomic.Uint64 // 0 means closed } // Close releases the native handle. Safe to call multiple times. func (p *ProviderHandle) Close() error { - p.mu.Lock() - defer p.mu.Unlock() - if p.closed { + if p == nil { return nil } - p.closed = true - if p.handle != 0 { - C.aimux_drop_handle(C.uint64_t(p.handle)) - p.handle = 0 + if id := p.id.Swap(0); id != 0 { + C.aimux_drop_handle(C.uint64_t(id)) + runtime.SetFinalizer(p, nil) } - runtime.SetFinalizer(p, nil) return nil } +func (p *ProviderHandle) handle() (uint64, error) { + if p != nil { + if id := p.id.Load(); id != 0 { + return id, nil + } + } + return 0, newError(CodeInvalidArgument, "aimux: provider handle is closed") +} + // ListModels lists models available on this provider (runtime discovery via // the provider's /models endpoint), enriched with community knowledge (anya2a) // when available. Returns a JSON array of RuntimeModel. func (p *ProviderHandle) ListModels() (string, error) { - p.mu.RLock() - defer p.mu.RUnlock() - if p.closed || p.handle == 0 { - return "", newError(CodeInvalidArgument, "aimux: provider handle is closed") + id, err := p.handle() + if err != nil { + return "", err } + defer runtime.KeepAlive(p) return ffiString(func(cerr *C.AimuxError) *C.char { - return C.aimux_provider_list_models(C.uint64_t(p.handle), cerr) + return C.aimux_provider_list_models(C.uint64_t(id), cerr) }) } // Model builds a language model from a discovered model id. func (p *ProviderHandle) Model(modelID string) (*Model, error) { - p.mu.RLock() - defer p.mu.RUnlock() - if p.closed || p.handle == 0 { - return nil, newError(CodeInvalidArgument, "aimux: provider handle is closed") + id, err := p.handle() + if err != nil { + return nil, err } + defer runtime.KeepAlive(p) cModel := C.CString(modelID) defer C.free(unsafe.Pointer(cModel)) var cerr C.AimuxError C.aimux_error_clear(&cerr) - h := C.aimux_provider_model(C.uint64_t(p.handle), cModel, &cerr) + h := C.aimux_provider_model(C.uint64_t(id), cModel, &cerr) return wrapHandleU64(h, &cerr) } @@ -734,7 +735,8 @@ func CreateProvider(name, apiKey string, cfg *ProviderConfig) (*ProviderHandle, if h == 0 { return nil, errorFromC(&cerr) } - p := &ProviderHandle{handle: uint64(h)} + p := &ProviderHandle{} + p.id.Store(uint64(h)) runtime.SetFinalizer(p, func(p *ProviderHandle) { p.Close() }) return p, nil } @@ -768,11 +770,11 @@ func GetModelSpecs(sourceURL string) (string, error) { // // result, err := model.GenerateText(`"What is Rust?"`, "") func (m *Model) GenerateText(promptJson, optsJson string) (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) @@ -795,11 +797,11 @@ func (m *Model) GenerateText(promptJson, optsJson string) (string, error) { // optsJson for schema control; the function applies JSON repair before // parsing. func (m *Model) GenerateObject(promptJson, optsJson string) (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) @@ -822,11 +824,11 @@ func (m *Model) GenerateObject(promptJson, optsJson string) (string, error) { // Same signature as GenerateText; returns the JSON-serialized // StreamTextResultAggregated. func (m *Model) ConsumeStreamText(promptJson, optsJson string) (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) @@ -849,11 +851,11 @@ func (m *Model) ConsumeStreamText(promptJson, optsJson string) (string, error) { // (OpenAI "chat.completion" object) rather than a GenerateTextResult. Works // with any provider. func (m *Model) GenerateTextAsOpenAI(promptJson, optsJson string) (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) @@ -871,93 +873,15 @@ func (m *Model) GenerateTextAsOpenAI(promptJson, optsJson string) (string, error // ── Streaming generation ───────────────────────────────────────────────────── -// streamEntry holds the channel and error for an active stream. -type streamEntry struct { - parts chan string - mu sync.Mutex - err error - closeOnce sync.Once - terminalOnce sync.Once - terminal chan struct{} - cancelled chan struct{} - abortHandle uint64 -} - -// streamRegistry maps stream IDs to active stream entries. This avoids passing -// Go pointers into C (cgo pointer rules forbid passing Go memory containing Go -// pointers like channels). The ID is a plain int64_t. -var ( - streamRegMu sync.Mutex - streamReg = make(map[int64]*streamEntry) - streamNextID int64 -) - -func registerStream(abortHandle uint64) (*streamEntry, int64) { - streamRegMu.Lock() - defer streamRegMu.Unlock() - streamNextID++ - e := &streamEntry{ - parts: make(chan string, 256), - terminal: make(chan struct{}), - cancelled: make(chan struct{}), - abortHandle: abortHandle, - } - streamReg[streamNextID] = e - return e, streamNextID -} - -func lookupStream(id int64) *streamEntry { - streamRegMu.Lock() - defer streamRegMu.Unlock() - return streamReg[id] -} - -func unregisterStream(id int64) { - streamRegMu.Lock() - defer streamRegMu.Unlock() - delete(streamReg, id) -} - -// closeParts safely closes the stream's parts channel exactly once, guarding -// against the engine firing both on_done and on_error for the same stream -// (which would otherwise panic on double close). -func (e *streamEntry) closeParts() { - e.closeOnce.Do(func() { close(e.parts) }) -} - -// markTerminal records the first terminal result. A cancellation also wakes -// callbacks that are waiting to send to a full parts channel. -func (e *streamEntry) markTerminal(err error, cancelled bool) bool { - marked := false - e.terminalOnce.Do(func() { - marked = true - e.mu.Lock() - e.err = err - e.mu.Unlock() - if cancelled { - close(e.cancelled) - } - close(e.terminal) - }) - return marked -} - -func (e *streamEntry) cancel(err error) { - if err == nil { - err = context.Canceled - } - if e.markTerminal(err, true) && e.abortHandle != 0 { - C.aimux_abort_signal_abort(C.uint64_t(e.abortHandle)) - } -} - -// Stream is a handle to an in-progress or completed stream. -// Consume parts via the Parts() channel; check Err() after the channel closes. -// -// The channel is buffered (256). Call Cancel when the caller stops consuming. +// Stream is an in-progress or completed stream. The producer goroutine is the +// only writer of err and the only goroutine that closes parts; observing the +// channel close publishes err under the Go memory model. type Stream struct { - parts <-chan string - entry *streamEntry + parts chan string + cancelCtx context.Context + cancelFn context.CancelCauseFunc + err error + abortHandle uint64 } // Parts returns a receive-only channel of StreamPart JSON strings. @@ -965,16 +889,28 @@ type Stream struct { func (s *Stream) Parts() <-chan string { return s.parts } // Err returns any error that occurred during streaming. -// Call this after Parts() channel closes. -func (s *Stream) Err() error { - s.entry.mu.Lock() - defer s.entry.mu.Unlock() - return s.entry.err -} +// It must be called only after Parts() has closed. +func (s *Stream) Err() error { return s.err } // Cancel stops this stream. It is safe to call more than once. -func (s *Stream) Cancel() { - s.entry.cancel(context.Canceled) +func (s *Stream) Cancel() { s.cancel(context.Canceled) } + +func (s *Stream) cancel(cause error) { + if cause == nil { + cause = context.Canceled + } + s.cancelFn(cause) + if s.abortHandle != 0 { + C.aimux_abort_signal_abort(C.uint64_t(s.abortHandle)) + } +} + +func (s *Stream) finish(err error) { + if cause := context.Cause(s.cancelCtx); cause != nil { + err = cause + } + s.err = err + close(s.parts) } // StreamText performs streaming text generation. @@ -998,67 +934,7 @@ func (m *Model) StreamText(promptJson, optsJson string) *Stream { // StreamTextContext performs streaming text generation. Context cancellation // stops the native request and makes Err return the context error. func (m *Model) StreamTextContext(ctx context.Context, promptJson, optsJson string) *Stream { - if ctx == nil { - ctx = context.Background() - } - abortHandle := uint64(C.aimux_abort_signal_new()) - entry, id := registerStream(abortHandle) - - if err := ctx.Err(); err != nil { - entry.cancel(err) - } else if done := ctx.Done(); done != nil { - go func() { - select { - case <-done: - entry.cancel(ctx.Err()) - case <-entry.terminal: - } - }() - } - - go func() { - defer unregisterStream(id) - defer C.aimux_abort_signal_drop(C.uint64_t(abortHandle)) - // Safety net: ensure the channel is always closed even if the - // native layer never fires on_done/on_error (defensive against - // future bugs or panic edges in the FFI layer). - defer func() { - entry.markTerminal(nil, false) - entry.closeParts() - }() - - handle, release, err := m.acquireHandle() - if err != nil { - entry.markTerminal(err, false) - return - } - defer release() - - cPrompt := C.CString(promptJson) - defer C.free(unsafe.Pointer(cPrompt)) - - var cOpts *C.char - if optsJson != "" { - cOpts = C.CString(optsJson) - defer C.free(unsafe.Pointer(cOpts)) - } - - var cerr C.AimuxError - C.aimux_error_clear(&cerr) - rc := C.do_stream( - C.uint64_t(handle), - C.uint64_t(abortHandle), - cPrompt, - cOpts, - C.int64_t(id), - &cerr, - ) - if rc == 0 { - entry.markTerminal(errorFromC(&cerr), false) - } - }() - - return &Stream{parts: entry.parts, entry: entry} + return m.startStream(ctx, promptJson, optsJson, false) } // StreamTextAsOpenAI performs streaming text generation with OpenAI Chat @@ -1077,66 +953,87 @@ func (m *Model) StreamTextAsOpenAI(promptJson, optsJson string) *Stream { // StreamTextAsOpenAIContext performs streaming OpenAI-compatible generation. // Context cancellation stops the native request. func (m *Model) StreamTextAsOpenAIContext(ctx context.Context, promptJson, optsJson string) *Stream { + return m.startStream(ctx, promptJson, optsJson, true) +} + +func (m *Model) startStream(ctx context.Context, promptJson, optsJson string, openAI bool) *Stream { if ctx == nil { ctx = context.Background() } abortHandle := uint64(C.aimux_abort_signal_new()) - entry, id := registerStream(abortHandle) + cancelCtx, cancelFn := context.WithCancelCause(context.Background()) + stream := &Stream{ + parts: make(chan string, 256), + cancelCtx: cancelCtx, + cancelFn: cancelFn, + abortHandle: abortHandle, + } + callbackHandle := cgo.NewHandle(stream) - if err := ctx.Err(); err != nil { - entry.cancel(err) - } else if done := ctx.Done(); done != nil { - go func() { - select { - case <-done: - entry.cancel(ctx.Err()) - case <-entry.terminal: - } - }() + if cause := context.Cause(ctx); cause != nil { + stream.cancel(cause) + } + stopContext := func() bool { return true } + if ctx.Done() != nil { + stopContext = context.AfterFunc(ctx, func() { stream.cancel(context.Cause(ctx)) }) } go func() { - defer unregisterStream(id) + defer callbackHandle.Delete() defer C.aimux_abort_signal_drop(C.uint64_t(abortHandle)) - // Safety net: ensure the channel is always closed even if the - // native layer never fires on_done/on_error. - defer func() { - entry.markTerminal(nil, false) - entry.closeParts() - }() - - handle, release, err := m.acquireHandle() - if err != nil { - entry.markTerminal(err, false) - return + err := m.runStream(stream, callbackHandle, promptJson, optsJson, openAI) + // Stop the watcher before the final cause check. If cancellation won + // concurrently, Cause is already visible even when the AfterFunc callback + // itself has not run yet, so terminal publication never has to wait for it. + stopContext() + if cause := context.Cause(ctx); cause != nil { + stream.cancel(cause) } - defer release() + stream.finish(err) + }() - cPrompt := C.CString(promptJson) - defer C.free(unsafe.Pointer(cPrompt)) + return stream +} - var cOpts *C.char - if optsJson != "" { - cOpts = C.CString(optsJson) - defer C.free(unsafe.Pointer(cOpts)) - } +func (m *Model) runStream(stream *Stream, callbackHandle cgo.Handle, promptJson, optsJson string, openAI bool) error { + select { + case <-stream.cancelCtx.Done(): + return context.Cause(stream.cancelCtx) + default: + } - var cerr C.AimuxError - C.aimux_error_clear(&cerr) - rc := C.do_stream_openai( - C.uint64_t(handle), - C.uint64_t(abortHandle), - cPrompt, - cOpts, - C.int64_t(id), - &cerr, - ) - if rc == 0 { - entry.markTerminal(errorFromC(&cerr), false) - } - }() + handle, err := m.handle() + if err != nil { + return err + } + defer runtime.KeepAlive(m) + + cPrompt := C.CString(promptJson) + defer C.free(unsafe.Pointer(cPrompt)) + var cOpts *C.char + if optsJson != "" { + cOpts = C.CString(optsJson) + defer C.free(unsafe.Pointer(cOpts)) + } - return &Stream{parts: entry.parts, entry: entry} + var cerr C.AimuxError + C.aimux_error_clear(&cerr) + var rc C.int32_t + if openAI { + rc = C.do_stream_openai( + C.uint64_t(handle), C.uint64_t(stream.abortHandle), + cPrompt, cOpts, C.uintptr_t(callbackHandle), &cerr, + ) + } else { + rc = C.do_stream( + C.uint64_t(handle), C.uint64_t(stream.abortHandle), + cPrompt, cOpts, C.uintptr_t(callbackHandle), &cerr, + ) + } + if rc == 0 { + return errorFromC(&cerr) + } + return nil } // errorFromC maps aimux-ffi AimuxError into *Error (openai-go style). @@ -1188,31 +1085,18 @@ func ffiString(call func(cerr *C.AimuxError) *C.char) (string, error) { return C.GoString(ptr), nil } -// ── C→Go callback trampolines (called by trampoline_part/done) ───────── +// ── C→Go callback trampoline ──────────────────────────────────────────── //export goStreamPart -func goStreamPart(id C.int64_t, json *C.char) { - e := lookupStream(int64(id)) - if e == nil { - return - } +func goStreamPart(id C.uintptr_t, json *C.char) { if json == nil { return } + stream := cgo.Handle(id).Value().(*Stream) select { - case e.parts <- C.GoString(json): - case <-e.cancelled: - } -} - -//export goStreamDone -func goStreamDone(id C.int64_t) { - e := lookupStream(int64(id)) - if e == nil { - return + case stream.parts <- C.GoString(json): + case <-stream.cancelCtx.Done(): } - e.markTerminal(nil, false) - e.closeParts() } // Ensure io.Closer interface is satisfied. diff --git a/bindings/go/aimux_test.go b/bindings/go/aimux_test.go index 54df5e90..7a6202e3 100644 --- a/bindings/go/aimux_test.go +++ b/bindings/go/aimux_test.go @@ -10,18 +10,52 @@ package aimux import ( "context" "errors" + "fmt" + "net/http" + "net/http/httptest" "strings" + "sync" + "sync/atomic" "testing" "time" ) +// cancelOnDoneContext makes cancellation happen while StreamTextContext is +// installing its context watcher. It deterministically exercises the window +// between the initial context check and context.AfterFunc registration. +type cancelOnDoneContext struct { + done chan struct{} + cancelled atomic.Bool + once sync.Once +} + +func newCancelOnDoneContext() *cancelOnDoneContext { + return &cancelOnDoneContext{done: make(chan struct{})} +} + +func (c *cancelOnDoneContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (c *cancelOnDoneContext) Done() <-chan struct{} { + c.once.Do(func() { + c.cancelled.Store(true) + close(c.done) + }) + return c.done +} +func (c *cancelOnDoneContext) Err() error { + if c.cancelled.Load() { + return context.Canceled + } + return nil +} +func (*cancelOnDoneContext) Value(any) any { return nil } + func TestOpenAI(t *testing.T) { m := OpenAI("sk-test-fake-key", "gpt-4o-mini") if m == nil { t.Fatal("expected non-nil model") } defer m.Close() - if m.handle == 0 { + if m.id.Load() == 0 { t.Fatal("expected non-zero handle") } } @@ -32,7 +66,7 @@ func TestAnthropic(t *testing.T) { t.Fatal("expected non-nil model") } defer m.Close() - if m.handle == 0 { + if m.id.Load() == 0 { t.Fatal("expected non-zero handle") } } @@ -43,7 +77,7 @@ func TestOpenAIWithBase(t *testing.T) { t.Fatal("expected non-nil model") } defer m.Close() - if m.handle == 0 { + if m.id.Load() == 0 { t.Fatal("expected non-zero handle") } } @@ -138,6 +172,56 @@ func TestStreamTextContextAlreadyCanceled(t *testing.T) { stream.Cancel() } +func TestStreamTextContextPreservesCancelCause(t *testing.T) { + m := OpenAI("sk-test-fake-key", "gpt-4o-mini") + defer m.Close() + + cause := errors.New("caller stopped consuming") + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(cause) + stream := m.StreamTextContext(ctx, `"hello"`, "") + for range stream.Parts() { + } + if !errors.Is(stream.Err(), cause) { + t.Fatalf("expected custom cancellation cause, got %v", stream.Err()) + } +} + +func TestStreamTextContextDoesNotLoseRacingCancellation(t *testing.T) { + requestStarted := make(chan struct{}, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestStarted <- struct{}{} + <-r.Context().Done() + })) + defer server.Close() + + m := OpenAIWithBase("sk-test-fake-key", "gpt-4o-mini", server.URL) + defer m.Close() + + stream := m.StreamTextContext(newCancelOnDoneContext(), `"hello"`, "") + done := make(chan struct{}) + go func() { + for range stream.Parts() { + } + close(done) + }() + + select { + case <-done: + case <-requestStarted: + stream.Cancel() + <-done + t.Fatal("request started after the context was cancelled") + case <-time.After(2 * time.Second): + stream.Cancel() + <-done + t.Fatal("racing cancellation did not stop the stream") + } + if !errors.Is(stream.Err(), context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", stream.Err()) + } +} + func TestProviderWithConfigFullOptions(t *testing.T) { retries := uint32(0) m, err := ProviderWithConfig("groq", "sk-test-fake-key", "llama-3.3-70b", &ProviderConfig{ @@ -152,7 +236,7 @@ func TestProviderWithConfigFullOptions(t *testing.T) { t.Fatalf("expected success, got %v", err) } defer m.Close() - if m.handle == 0 { + if m.id.Load() == 0 { t.Fatal("expected non-zero handle") } } @@ -175,6 +259,63 @@ func TestProviderWithBaseQuotedURLDoesNotInjectJSON(t *testing.T) { } } +func TestProviderCloseDoesNotWaitForInFlightListModels(t *testing.T) { + server, requestStarted, releaseRequest := newBlockedHTTPServer( + t, "application/json", `{"data":[]}`, + ) + defer releaseRequest() + + provider, err := CreateProvider("groq", "sk-test-fake-key", &ProviderConfig{ + BaseURL: server.URL, + }) + if err != nil { + t.Fatalf("CreateProvider failed: %v", err) + } + defer provider.Close() + + listDone := make(chan error, 1) + go func() { + _, err := provider.ListModels() + listDone <- err + }() + select { + case <-requestStarted: + case <-time.After(2 * time.Second): + t.Fatal("ListModels request did not start") + } + + closes := make([]func(), 32) + for i := range closes { + closes[i] = func() { + if err := provider.Close(); err != nil { + t.Errorf("Close: %v", err) + } + } + } + runConcurrent(t, "ProviderHandle.Close during ListModels", closes...) + if got := provider.id.Load(); got != 0 { + t.Fatalf("provider handle after Close = %d, want 0", got) + } + + // The call already entered the native registry, so dropping the Go owner + // must not invalidate its Arc. Let the provider response complete now. + releaseRequest() + select { + case err := <-listDone: + if err != nil { + t.Fatalf("in-flight ListModels failed after Close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("in-flight ListModels did not finish after releasing the response") + } + + _, err = provider.ListModels() + var aimuxErr *Error + if !errors.As(err, &aimuxErr) || aimuxErr.Code != CodeInvalidArgument { + t.Fatalf("ListModels after Close: expected flat InvalidArgument, got %v", err) + } +} + // TestInitRecordingRingRejectsZeroCap verifies D7: cap == 0 is rejected with // an error instead of being silently rewritten to 2048. The check happens in // Go before any FFI call, so this does not touch the global recorder state. @@ -207,3 +348,219 @@ func TestInitRecordingRingDefaultNoArg(t *testing.T) { // Reset global recorder state so this doesn't leak into other tests. RecordingStop() } + +// runConcurrent starts every task at one barrier and requires every task, +// including Close, to return. Atomic handle snapshots create no Go lock order +// or wait edge between these operations. +func runConcurrent(t *testing.T, what string, tasks ...func()) { + t.Helper() + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(len(tasks)) + for _, task := range tasks { + go func() { + defer wg.Done() + <-start + task() + }() + } + close(start) + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("%s did not return within 10s; possible deadlock", what) + } +} + +// newBlockedHTTPServer starts handling a request, then withholds its response +// until release is called. It makes an in-flight native call observable +// without adding a production test seam. release is idempotent. +func newBlockedHTTPServer(t *testing.T, contentType, body string) (*httptest.Server, <-chan struct{}, func()) { + t.Helper() + requestStarted := make(chan struct{}) + releaseRequest := make(chan struct{}) + var startedOnce sync.Once + var releaseOnce sync.Once + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + startedOnce.Do(func() { close(requestStarted) }) + <-releaseRequest + if contentType != "" { + w.Header().Set("Content-Type", contentType) + } + _, _ = fmt.Fprint(w, body) + })) + release := func() { releaseOnce.Do(func() { close(releaseRequest) }) } + t.Cleanup(func() { + release() + server.Close() + }) + return server, requestStarted, release +} + +func duplicateModel(m *Model, n int) []*Model { + models := make([]*Model, n) + for i := range models { + models[i] = m + } + return models +} + +func TestNewRouterRepeatedModelRacingClose(t *testing.T) { + m := OpenAI("sk-test-fake-key", "gpt-4o-mini") + if router, err := NewRouter(duplicateModel(m, 8), ""); err != nil { + t.Fatalf("duplicate router input: %v", err) + } else { + router.Close() + } + m.Close() + + for i := 0; i < 25; i++ { + m = OpenAI("sk-test-fake-key", "gpt-4o-mini") + tasks := []func(){func() { m.Close() }} + for j := 0; j < 16; j++ { + tasks = append(tasks, func() { + if router, err := NewRouter(duplicateModel(m, 8), ""); err == nil { + router.Close() + } + }) + } + runConcurrent(t, "NewRouter over a repeated model", tasks...) + } +} + +func TestNewMoaAggregatorAlsoReferenceRacingClose(t *testing.T) { + m := OpenAI("sk-test-fake-key", "gpt-4o-mini") + if moa, err := NewMoa(duplicateModel(m, 7), m, ""); err != nil { + t.Fatalf("duplicate MoA input: %v", err) + } else { + moa.Close() + } + m.Close() + + for i := 0; i < 25; i++ { + m = OpenAI("sk-test-fake-key", "gpt-4o-mini") + tasks := []func(){func() { m.Close() }} + for j := 0; j < 16; j++ { + tasks = append(tasks, func() { + if moa, err := NewMoa(duplicateModel(m, 7), m, ""); err == nil { + moa.Close() + } + }) + } + runConcurrent(t, "NewMoa with the aggregator also in references", tasks...) + } +} + +func TestCompositeConstructorsOppositeOrdersRacingClose(t *testing.T) { + const ( + iterations = 30 + fanout = 8 + pairs = 8 + ) + for i := 0; i < iterations; i++ { + models := make([]*Model, fanout) + for j := range models { + models[j] = OpenAI("sk-test-fake-key", "gpt-4o-mini") + } + reversed := make([]*Model, fanout) + for j, model := range models { + reversed[fanout-1-j] = model + } + + tasks := make([]func(), 0, pairs*2+fanout) + build := func(order []*Model) func() { + return func() { + if router, err := NewRouter(order, ""); err == nil { + router.Close() + } + } + } + for j := 0; j < pairs; j++ { + tasks = append(tasks, build(models), build(reversed)) + } + for _, model := range models { + model := model + tasks = append(tasks, func() { model.Close() }) + } + runConcurrent(t, fmt.Sprintf("opposite-order iteration %d", i), tasks...) + } +} + +func TestZeroValueModelIsClosed(t *testing.T) { + var model Model + _, err := model.GenerateText(`"hello"`, "") + var aimuxErr *Error + if !errors.As(err, &aimuxErr) || aimuxErr.Code != CodeInvalidArgument { + t.Fatalf("zero-value Model: expected InvalidArgument, got %v", err) + } + if err := model.Close(); err != nil { + t.Fatalf("zero-value Close: %v", err) + } +} + +func TestModelConcurrentCloseIsIdempotent(t *testing.T) { + model := OpenAI("sk-test-fake-key", "gpt-4o-mini") + tasks := make([]func(), 64) + for i := range tasks { + tasks[i] = func() { + if err := model.Close(); err != nil { + t.Errorf("Close: %v", err) + } + } + } + runConcurrent(t, "concurrent Model.Close", tasks...) + if got := model.id.Load(); got != 0 { + t.Fatalf("handle after Close = %d, want 0", got) + } + if err := model.Close(); err != nil { + t.Fatalf("repeated Close: %v", err) + } +} + +func TestCompositeConstructorsRejectNilModels(t *testing.T) { + model := OpenAI("sk-test-fake-key", "gpt-4o-mini") + defer model.Close() + + cases := []struct { + name string + want string + call func() (*Model, error) + }{ + {"router nil child", "aimux: models[1] is nil", + func() (*Model, error) { return NewRouter([]*Model{model, nil, model}, "") }}, + {"router only child nil", "aimux: models[0] is nil", + func() (*Model, error) { return NewRouter([]*Model{nil}, "") }}, + {"moa nil aggregator", "aimux: aggregator is nil", + func() (*Model, error) { return NewMoa([]*Model{model}, nil, "") }}, + {"moa nil reference", "aimux: models[1] is nil", + func() (*Model, error) { return NewMoa([]*Model{model, nil}, model, "") }}, + {"moa nil aggregator, no references", "aimux: aggregator is nil", + func() (*Model, error) { return NewMoa(nil, nil, "") }}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + defer func() { + if recovered := recover(); recovered != nil { + t.Fatalf("panicked instead of returning an error: %v", recovered) + } + }() + got, err := test.call() + if got != nil { + got.Close() + t.Fatal("expected a nil model alongside the error") + } + if err == nil { + t.Fatal("expected an error for a nil model") + } + if err.Error() != test.want { + t.Fatalf("error = %q, want %q", err, test.want) + } + }) + } +} diff --git a/bindings/go/e2e_test.go b/bindings/go/e2e_test.go index 268104c8..40ef89b5 100644 --- a/bindings/go/e2e_test.go +++ b/bindings/go/e2e_test.go @@ -511,15 +511,13 @@ func TestStreamCancelUnblocksFullPartsChannel(t *testing.T) { case <-time.After(2 * time.Second): t.Fatal("provider did not send the SSE burst") } - waitForChannelFull(t, stream.entry.parts) - select { - case <-stream.entry.terminal: - t.Fatal("stream ended before cancellation") - default: - } + waitForChannelFull(t, stream.parts) - stream.Cancel() - stream.Cancel() + tasks := []func(){func() { m.Close() }} + for i := 0; i < 32; i++ { + tasks = append(tasks, stream.Cancel) + } + runConcurrent(t, "full stream buffer Cancel + Model.Close", tasks...) drained := make(chan struct{}) go func() { for range stream.Parts() { diff --git a/bindings/go/multimodal.go b/bindings/go/multimodal.go index 132dc24d..b0a82d1d 100644 --- a/bindings/go/multimodal.go +++ b/bindings/go/multimodal.go @@ -25,63 +25,68 @@ import ( "errors" "fmt" "runtime" - "sync" + "sync/atomic" "unsafe" ) // ── Shared helpers ─────────────────────────────────────────────────────────── -// multimodalHandle is the common structure for all multimodal model types. -// It mirrors Model but is kept separate because the C ABI uses distinct -// handle types (Embedding/Speech/Image/... are not interchangeable). +// multimodalHandle is the common state for all multimodal model types. The +// atomic swap makes Close non-blocking: a call that already loaded the handle +// may finish (the Rust registry clones its Arc at entry), while later calls see +// zero and return the existing closed-handle error. No Go lock is held across +// a blocking C call. +// +// A multimodalHandle must not be copied after first use. Public model wrappers +// embed it so atomic.noCopy also lets go vet flag accidental value copies. type multimodalHandle struct { - mu sync.RWMutex - handle uint64 - closed bool + handle atomic.Uint64 } -func (h *multimodalHandle) acquire() (uint64, func(), error) { - h.mu.RLock() - if h.closed { - h.mu.RUnlock() - return 0, nil, newError(CodeInvalidArgument, "aimux: model already closed") +func (h *multimodalHandle) load() (uint64, error) { + handle := h.handle.Load() + if handle == 0 { + return 0, newError(CodeInvalidArgument, "aimux: model already closed") } - return h.handle, h.mu.RUnlock, nil + return handle, nil } func (h *multimodalHandle) close() { - h.mu.Lock() - defer h.mu.Unlock() - if h.closed { - return - } - h.closed = true - if h.handle != 0 { - C.aimux_drop_handle(C.uint64_t(h.handle)) - h.handle = 0 + if handle := h.handle.Swap(0); handle != 0 { + C.aimux_drop_handle(C.uint64_t(handle)) } } -// callFFIString is a helper for the common pattern: acquire handle, call a -// C function that returns char* and fills AimuxError on failure. -func callFFIString(h *multimodalHandle, fn func(handle C.uint64_t, err *C.AimuxError) *C.char) (string, error) { - handle, release, err := h.acquire() +func closeMultimodal(owner any, h *multimodalHandle) { + h.close() + runtime.SetFinalizer(owner, nil) + runtime.KeepAlive(owner) +} + +// callFFIString is the common pattern for a multimodal call returning an +// aimux-allocated char*. owner is kept alive until after C returns so its +// finalizer cannot close the handle while the call is entering the Rust +// registry. Explicit concurrent Close is allowed and may make the call fail; +// it cannot cause use-after-free. +func callFFIString(owner any, h *multimodalHandle, fn func(handle C.uint64_t, err *C.AimuxError) *C.char) (string, error) { + handle, err := h.load() if err != nil { return "", err } - defer release() - return ffiString(func(cerr *C.AimuxError) *C.char { + result, callErr := ffiString(func(cerr *C.AimuxError) *C.char { return fn(C.uint64_t(handle), cerr) }) + runtime.KeepAlive(owner) + return result, callErr } // newMultimodalHandleU64 wraps a constructor handle + AimuxError. -func newMultimodalHandleU64(h C.uint64_t, cerr *C.AimuxError) (*multimodalHandle, error) { +func newMultimodalHandleU64(h C.uint64_t, cerr *C.AimuxError) (uint64, error) { if h == 0 { - return nil, errorFromC(cerr) + return 0, errorFromC(cerr) } - return &multimodalHandle{handle: uint64(h)}, nil + return uint64(h), nil } // cstringPair creates two C strings and returns them with a cleanup func. @@ -113,7 +118,7 @@ func newMultimodalModelWithBase( apiKey, modelID, baseURL string, plain func(ca, cb *C.char, err *C.AimuxError) C.uint64_t, withBase func(ca, cb, cbase *C.char, err *C.AimuxError) C.uint64_t, -) (*multimodalHandle, error) { +) (uint64, error) { var cerr C.AimuxError C.aimux_error_clear(&cerr) if baseURL == "" { @@ -132,7 +137,7 @@ func newMultimodalModelFiles( apiKey, baseURL string, plain func(ca *C.char, err *C.AimuxError) C.uint64_t, withBase func(ca, cbase *C.char, err *C.AimuxError) C.uint64_t, -) (*multimodalHandle, error) { +) (uint64, error) { var cerr C.AimuxError C.aimux_error_clear(&cerr) ca := C.CString(apiKey) @@ -148,13 +153,14 @@ func newMultimodalModelFiles( // ── EmbeddingModel ────────────────────────────────────────────────────────── // EmbeddingModel generates vector embeddings for text. +// An EmbeddingModel must not be copied after first use. type EmbeddingModel struct { - h *multimodalHandle + h multimodalHandle } // Close releases the native handle. func (m *EmbeddingModel) Close() error { - m.h.close() + closeMultimodal(m, &m.h) return nil } @@ -174,11 +180,10 @@ func (m *EmbeddingModel) Embed(values []string, opts *EmbeddingCallOptions) (str optsJSON = string(b) } - handle, release, err := m.h.acquire() + handle, err := m.h.load() if err != nil { return "", err } - defer release() cVals := C.CString(string(valuesJSON)) defer C.free(unsafe.Pointer(cVals)) @@ -188,9 +193,11 @@ func (m *EmbeddingModel) Embed(values []string, opts *EmbeddingCallOptions) (str defer C.free(unsafe.Pointer(cOpts)) } - return ffiString(func(cerr *C.AimuxError) *C.char { + result, callErr := ffiString(func(cerr *C.AimuxError) *C.char { return C.aimux_embed(C.uint64_t(handle), cVals, cOpts, cerr) }) + runtime.KeepAlive(m) + return result, callErr } // ParseEmbeddingResult parses the JSON string returned by Embed. @@ -220,7 +227,8 @@ func NewOpenAIEmbeddingWithBase(apiKey, modelID, baseURL string) (*EmbeddingMode if err != nil { return nil, err } - m := &EmbeddingModel{h: mh} + m := &EmbeddingModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *EmbeddingModel) { m.Close() }) return m, nil } @@ -241,7 +249,8 @@ func NewCohereEmbeddingWithBase(apiKey, modelID, baseURL string) (*EmbeddingMode if err != nil { return nil, err } - m := &EmbeddingModel{h: mh} + m := &EmbeddingModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *EmbeddingModel) { m.Close() }) return m, nil } @@ -262,7 +271,8 @@ func NewGoogleEmbeddingWithBase(apiKey, modelID, baseURL string) (*EmbeddingMode if err != nil { return nil, err } - m := &EmbeddingModel{h: mh} + m := &EmbeddingModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *EmbeddingModel) { m.Close() }) return m, nil } @@ -270,11 +280,12 @@ func NewGoogleEmbeddingWithBase(apiKey, modelID, baseURL string) (*EmbeddingMode // ── SpeechModel (TTS) ──────────────────────────────────────────────────────── // SpeechModel converts text to speech audio. +// A SpeechModel must not be copied after first use. type SpeechModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *SpeechModel) Close() error { m.h.close(); return nil } +func (m *SpeechModel) Close() error { closeMultimodal(m, &m.h); return nil } // Generate generates speech audio from the given options. func (m *SpeechModel) Generate(opts *SpeechCallOptions) (string, error) { @@ -286,7 +297,7 @@ func (m *SpeechModel) Generate(opts *SpeechCallOptions) (string, error) { } optsJSON = string(b) } - return callFFIString(m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { + return callFFIString(m, &m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { cOpts := C.CString(optsJSON) defer C.free(unsafe.Pointer(cOpts)) return C.aimux_speech_generate(handle, cOpts, err) @@ -317,7 +328,8 @@ func NewOpenAISpeechWithBase(apiKey, modelID, baseURL string) (*SpeechModel, err if err != nil { return nil, err } - m := &SpeechModel{h: mh} + m := &SpeechModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *SpeechModel) { m.Close() }) return m, nil } @@ -325,11 +337,12 @@ func NewOpenAISpeechWithBase(apiKey, modelID, baseURL string) (*SpeechModel, err // ── ImageModel ────────────────────────────────────────────────────────────── // ImageModel generates images from prompts. +// An ImageModel must not be copied after first use. type ImageModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *ImageModel) Close() error { m.h.close(); return nil } +func (m *ImageModel) Close() error { closeMultimodal(m, &m.h); return nil } func (m *ImageModel) Generate(opts *ImageCallOptions) (string, error) { optsJSON := "" @@ -340,7 +353,7 @@ func (m *ImageModel) Generate(opts *ImageCallOptions) (string, error) { } optsJSON = string(b) } - return callFFIString(m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { + return callFFIString(m, &m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { cOpts := C.CString(optsJSON) defer C.free(unsafe.Pointer(cOpts)) return C.aimux_image_generate(handle, cOpts, err) @@ -371,7 +384,8 @@ func NewOpenAIImageWithBase(apiKey, modelID, baseURL string) (*ImageModel, error if err != nil { return nil, err } - m := &ImageModel{h: mh} + m := &ImageModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *ImageModel) { m.Close() }) return m, nil } @@ -392,7 +406,8 @@ func NewGoogleImageWithBase(apiKey, modelID, baseURL string) (*ImageModel, error if err != nil { return nil, err } - m := &ImageModel{h: mh} + m := &ImageModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *ImageModel) { m.Close() }) return m, nil } @@ -400,11 +415,12 @@ func NewGoogleImageWithBase(apiKey, modelID, baseURL string) (*ImageModel, error // ── TranscriptionModel (STT) ──────────────────────────────────────────────── // TranscriptionModel converts audio to text. +// A TranscriptionModel must not be copied after first use. type TranscriptionModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *TranscriptionModel) Close() error { m.h.close(); return nil } +func (m *TranscriptionModel) Close() error { closeMultimodal(m, &m.h); return nil } // Generate transcribes audio (base64-encoded) to text. func (m *TranscriptionModel) Generate(audioBase64, mediaType string, opts *TranscriptionCallOptions) (string, error) { @@ -417,11 +433,10 @@ func (m *TranscriptionModel) Generate(audioBase64, mediaType string, opts *Trans optsJSON = string(b) } - handle, release, err := m.h.acquire() + handle, err := m.h.load() if err != nil { return "", err } - defer release() ca, cb, cc, cleanup := cstringTriple(audioBase64, mediaType, optsJSON) defer cleanup() @@ -430,9 +445,11 @@ func (m *TranscriptionModel) Generate(audioBase64, mediaType string, opts *Trans cOpts = cc } - return ffiString(func(cerr *C.AimuxError) *C.char { + result, callErr := ffiString(func(cerr *C.AimuxError) *C.char { return C.aimux_transcription_generate(C.uint64_t(handle), ca, cb, cOpts, cerr) }) + runtime.KeepAlive(m) + return result, callErr } func ParseTranscriptionResult(jsonStr string) (*TranscriptionResult, error) { @@ -461,7 +478,8 @@ func NewOpenAITranscriptionWithBase(apiKey, modelID, baseURL string) (*Transcrip if err != nil { return nil, err } - m := &TranscriptionModel{h: mh} + m := &TranscriptionModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *TranscriptionModel) { m.Close() }) return m, nil } @@ -469,11 +487,12 @@ func NewOpenAITranscriptionWithBase(apiKey, modelID, baseURL string) (*Transcrip // ── Files ─────────────────────────────────────────────────────────────────── // Files manages file uploads to providers. +// A Files value must not be copied after first use. type Files struct { - h *multimodalHandle + h multimodalHandle } -func (f *Files) Close() error { f.h.close(); return nil } +func (f *Files) Close() error { closeMultimodal(f, &f.h); return nil } // Upload uploads a file (base64-encoded) to the provider. func (f *Files) Upload(dataBase64, mediaType string, opts *UploadFileCallOptions) (string, error) { @@ -486,11 +505,10 @@ func (f *Files) Upload(dataBase64, mediaType string, opts *UploadFileCallOptions optsJSON = string(b) } - handle, release, err := f.h.acquire() + handle, err := f.h.load() if err != nil { return "", err } - defer release() ca, cb, cc, cleanup := cstringTriple(dataBase64, mediaType, optsJSON) defer cleanup() @@ -499,9 +517,11 @@ func (f *Files) Upload(dataBase64, mediaType string, opts *UploadFileCallOptions cOpts = cc } - return ffiString(func(cerr *C.AimuxError) *C.char { + result, callErr := ffiString(func(cerr *C.AimuxError) *C.char { return C.aimux_file_upload(C.uint64_t(handle), ca, cb, cOpts, cerr) }) + runtime.KeepAlive(f) + return result, callErr } func ParseUploadFileResult(jsonStr string) (*UploadFileResult, error) { @@ -528,7 +548,8 @@ func NewOpenAIFilesWithBase(apiKey, baseURL string) (*Files, error) { if err != nil { return nil, err } - f := &Files{h: mh} + f := &Files{} + f.h.handle.Store(mh) runtime.SetFinalizer(f, func(f *Files) { f.Close() }) return f, nil } @@ -536,11 +557,12 @@ func NewOpenAIFilesWithBase(apiKey, baseURL string) (*Files, error) { // ── RerankingModel ────────────────────────────────────────────────────────── // RerankingModel reranks documents by relevance to a query. +// A RerankingModel must not be copied after first use. type RerankingModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *RerankingModel) Close() error { m.h.close(); return nil } +func (m *RerankingModel) Close() error { closeMultimodal(m, &m.h); return nil } // Rerank reranks documents against a query. func (m *RerankingModel) Rerank(opts *RerankingCallOptions) (string, error) { @@ -552,7 +574,7 @@ func (m *RerankingModel) Rerank(opts *RerankingCallOptions) (string, error) { } optsJSON = string(b) } - return callFFIString(m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { + return callFFIString(m, &m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { cOpts := C.CString(optsJSON) defer C.free(unsafe.Pointer(cOpts)) return C.aimux_rerank(handle, cOpts, err) @@ -583,7 +605,8 @@ func NewCohereRerankingWithBase(apiKey, modelID, baseURL string) (*RerankingMode if err != nil { return nil, err } - m := &RerankingModel{h: mh} + m := &RerankingModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *RerankingModel) { m.Close() }) return m, nil } @@ -591,11 +614,12 @@ func NewCohereRerankingWithBase(apiKey, modelID, baseURL string) (*RerankingMode // ── VideoModel ────────────────────────────────────────────────────────────── // VideoModel generates videos from prompts. +// A VideoModel must not be copied after first use. type VideoModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *VideoModel) Close() error { m.h.close(); return nil } +func (m *VideoModel) Close() error { closeMultimodal(m, &m.h); return nil } func (m *VideoModel) Generate(opts *VideoCallOptions) (string, error) { optsJSON := "" @@ -606,7 +630,7 @@ func (m *VideoModel) Generate(opts *VideoCallOptions) (string, error) { } optsJSON = string(b) } - return callFFIString(m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { + return callFFIString(m, &m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { cOpts := C.CString(optsJSON) defer C.free(unsafe.Pointer(cOpts)) return C.aimux_video_generate(handle, cOpts, err) @@ -637,7 +661,8 @@ func NewGoogleVideoWithBase(apiKey, modelID, baseURL string) (*VideoModel, error if err != nil { return nil, err } - m := &VideoModel{h: mh} + m := &VideoModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *VideoModel) { m.Close() }) return m, nil } @@ -645,11 +670,12 @@ func NewGoogleVideoWithBase(apiKey, modelID, baseURL string) (*VideoModel, error // ── SearchModel ───────────────────────────────────────────────────────────── // SearchModel performs web search. +// A SearchModel must not be copied after first use. type SearchModel struct { - h *multimodalHandle + h multimodalHandle } -func (m *SearchModel) Close() error { m.h.close(); return nil } +func (m *SearchModel) Close() error { closeMultimodal(m, &m.h); return nil } func (m *SearchModel) Search(opts *SearchCallOptions) (string, error) { optsJSON := "" @@ -660,7 +686,7 @@ func (m *SearchModel) Search(opts *SearchCallOptions) (string, error) { } optsJSON = string(b) } - return callFFIString(m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { + return callFFIString(m, &m.h, func(handle C.uint64_t, err *C.AimuxError) *C.char { cOpts := C.CString(optsJSON) defer C.free(unsafe.Pointer(cOpts)) return C.aimux_search(handle, cOpts, err) @@ -700,7 +726,8 @@ func NewTavilySearchWithBase(apiKey, baseURL string) (*SearchModel, error) { if err != nil { return nil, err } - m := &SearchModel{h: mh} + m := &SearchModel{} + m.h.handle.Store(mh) runtime.SetFinalizer(m, func(m *SearchModel) { m.Close() }) return m, nil } @@ -723,14 +750,13 @@ func NewDeepSeek(apiKey, modelID string) (*Model, error) { // Push audio chunks with PushAudio, mark end-of-audio with InputDone, then // pull transcription parts (JSON TranscriptionStreamPart) with NextPart. // Close releases the session (safe and idempotent). +// +// A TranscriptionSession must not be copied after first use. Close atomically +// takes its native handle and immediately terminates the native session, so it +// can wake an in-flight NextPart(-1) or a backpressured PushAudio instead of +// waiting for that operation to return first. type TranscriptionSession struct { - // RWMutex: Push/NextPart/InputDone take the READ lock so a bidirectional - // session works from two goroutines (a blocking NextPart(-1) must not - // exclude a concurrent PushAudio — that would deadlock). Close takes the - // write lock, waiting for in-flight calls to finish. - mu sync.RWMutex - session uint64 - closed bool + session atomic.Uint64 } // InputAudioFormat is the input audio format for streaming transcription @@ -766,11 +792,13 @@ func StartTranscriptionSession(model *TranscriptionModel, opts *TranscriptionSes // StartTranscriptionSessionWithAbort is StartTranscriptionSession with an // abort handle (from AbortSignalNew); firing it aborts the session. func StartTranscriptionSessionWithAbort(model *TranscriptionModel, opts *TranscriptionSessionOpts, abortHandle uint64) (*TranscriptionSession, error) { - modelHandle, release, err := model.h.acquire() + if model == nil { + return nil, newError(CodeInvalidArgument, "aimux: transcription model is nil") + } + modelHandle, err := model.h.load() if err != nil { return nil, err } - defer release() optsJSON := cNullOrEmpty(opts) ca, cleanup := cstring1(optsJSON) @@ -784,10 +812,12 @@ func StartTranscriptionSessionWithAbort(model *TranscriptionModel, opts *Transcr ca, &cerr, ) + runtime.KeepAlive(model) if h == 0 { return nil, errorFromC(&cerr) } - s := &TranscriptionSession{session: uint64(h)} + s := &TranscriptionSession{} + s.session.Store(uint64(h)) runtime.SetFinalizer(s, func(s *TranscriptionSession) { s.Close() }) return s, nil } @@ -795,12 +825,10 @@ func StartTranscriptionSessionWithAbort(model *TranscriptionModel, opts *Transcr // PushAudio pushes one binary audio chunk. Blocks while the internal channel // is full (backpressure propagation). func (s *TranscriptionSession) PushAudio(audio []byte) error { - handle, release, err := s.acquire() + handle, err := s.load() if err != nil { return err } - defer release() - var ptr *C.uint8_t if len(audio) > 0 { ptr = (*C.uint8_t)(unsafe.Pointer(&audio[0])) @@ -813,6 +841,8 @@ func (s *TranscriptionSession) PushAudio(audio []byte) error { C.size_t(len(audio)), &cerr, ) + runtime.KeepAlive(audio) + runtime.KeepAlive(s) if rc == 0 { return errorFromC(&cerr) } @@ -821,15 +851,14 @@ func (s *TranscriptionSession) PushAudio(audio []byte) error { // InputDone signals end-of-audio (idempotent). func (s *TranscriptionSession) InputDone() error { - handle, release, err := s.acquire() + handle, err := s.load() if err != nil { return err } - defer release() - var cerr C.AimuxError C.aimux_error_clear(&cerr) rc := C.aimux_transcription_input_done(C.uint64_t(handle), &cerr) + runtime.KeepAlive(s) if rc == 0 { return errorFromC(&cerr) } @@ -841,12 +870,10 @@ func (s *TranscriptionSession) InputDone() error { // ErrTranscriptionTimeout when no part arrived in time (retryable). // timeoutMs: >0 wait at most; 0 immediate poll; <0 wait indefinitely. func (s *TranscriptionSession) NextPart(timeoutMs int64) (string, error) { - handle, release, err := s.acquire() + handle, err := s.load() if err != nil { return "", err } - defer release() - var cerr C.AimuxError C.aimux_error_clear(&cerr) ptr := C.aimux_transcription_next_part( @@ -854,6 +881,7 @@ func (s *TranscriptionSession) NextPart(timeoutMs int64) (string, error) { C.int64_t(timeoutMs), &cerr, ) + runtime.KeepAlive(s) if ptr != nil { defer C.aimux_free_string(ptr) return C.GoString(ptr), nil @@ -872,26 +900,19 @@ func (s *TranscriptionSession) NextPart(timeoutMs int64) (string, error) { // Close terminates and releases the session (aborts the driver; idempotent). func (s *TranscriptionSession) Close() { - s.mu.Lock() - defer s.mu.Unlock() - if s.closed { - return - } - s.closed = true - if s.session != 0 { - C.aimux_transcription_session_drop(C.uint64_t(s.session)) - s.session = 0 + if handle := s.session.Swap(0); handle != 0 { + C.aimux_transcription_session_drop(C.uint64_t(handle)) } + runtime.SetFinalizer(s, nil) + runtime.KeepAlive(s) } -func (s *TranscriptionSession) acquire() (uint64, func(), error) { - s.mu.RLock() - if s.closed { - s.mu.RUnlock() - return 0, nil, newError(CodeInvalidArgument, "aimux: transcription session already closed") +func (s *TranscriptionSession) load() (uint64, error) { + handle := s.session.Load() + if handle == 0 { + return 0, newError(CodeInvalidArgument, "aimux: transcription session already closed") } - h := s.session - return h, s.mu.RUnlock, nil + return handle, nil } func cstring1(a string) (*C.char, func()) { diff --git a/bindings/go/multimodal_test.go b/bindings/go/multimodal_test.go index f010c7b8..ff2c905f 100644 --- a/bindings/go/multimodal_test.go +++ b/bindings/go/multimodal_test.go @@ -7,9 +7,11 @@ package aimux import ( "encoding/json" + "errors" "net/http" "strings" "testing" + "time" ) // ── Constructor tests ──────────────────────────────────────────────────────── @@ -96,6 +98,214 @@ func TestSpeechDoubleClose(t *testing.T) { } } +func TestStartTranscriptionSessionRejectsNilModel(t *testing.T) { + _, err := StartTranscriptionSession(nil, nil) + var aimuxErr *Error + if !errors.As(err, &aimuxErr) || aimuxErr.Code != CodeInvalidArgument { + t.Fatalf("nil transcription model: expected flat InvalidArgument, got %v", err) + } +} + +func TestMultimodalCloseDoesNotWaitForInFlightCall(t *testing.T) { + server, requestStarted, releaseRequest := newBlockedHTTPServer( + t, "audio/mpeg", "test audio", + ) + defer releaseRequest() + + model, err := NewOpenAISpeechWithBase("sk-test-fake-key", "tts-1", server.URL) + if err != nil { + t.Fatalf("NewOpenAISpeechWithBase failed: %v", err) + } + defer model.Close() + + voice := "alloy" + format := "mp3" + generateDone := make(chan error, 1) + go func() { + _, err := model.Generate(&SpeechCallOptions{ + Text: "hello", + Voice: &voice, + OutputFormat: &format, + }) + generateDone <- err + }() + select { + case <-requestStarted: + case <-time.After(2 * time.Second): + t.Fatal("speech request did not start") + } + + closes := make([]func(), 32) + for i := range closes { + closes[i] = func() { + if err := model.Close(); err != nil { + t.Errorf("Close: %v", err) + } + } + } + runConcurrent(t, "multimodal Close during an in-flight call", closes...) + if got := model.h.handle.Load(); got != 0 { + t.Fatalf("multimodal handle after Close = %d, want 0", got) + } + + releaseRequest() + select { + case err := <-generateDone: + if err != nil { + t.Fatalf("in-flight Generate failed after Close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("in-flight Generate did not finish after releasing the response") + } + + _, err = model.Generate(&SpeechCallOptions{Text: "after close"}) + var aimuxErr *Error + if !errors.As(err, &aimuxErr) || aimuxErr.Code != CodeInvalidArgument { + t.Fatalf("Generate after Close: expected flat InvalidArgument, got %v", err) + } +} + +func TestTranscriptionSessionCloseUnblocksNextPart(t *testing.T) { + // Withhold the WebSocket handshake. The native driver cannot publish its + // first part, so NextPart(-1) is guaranteed to wait until Close aborts it. + server, handshakeStarted, releaseHandshake := newBlockedHTTPServer(t, "", "") + defer releaseHandshake() + + model, err := NewOpenAITranscriptionWithBase( + "sk-test-fake-key", "gpt-realtime-whisper", server.URL, + ) + if err != nil { + t.Fatalf("NewOpenAITranscriptionWithBase failed: %v", err) + } + defer model.Close() + session, err := StartTranscriptionSession(model, nil) + if err != nil { + t.Fatalf("StartTranscriptionSession failed: %v", err) + } + + select { + case <-handshakeStarted: + case <-time.After(2 * time.Second): + t.Fatal("WebSocket handshake did not start") + } + nextStarted := make(chan struct{}) + nextDone := make(chan error, 1) + go func() { + close(nextStarted) + _, err := session.NextPart(-1) + nextDone <- err + }() + <-nextStarted + select { + case err := <-nextDone: + session.Close() + t.Fatalf("NextPart(-1) returned before Close: %v", err) + case <-time.After(200 * time.Millisecond): + } + + closeDone := make(chan struct{}) + go func() { + session.Close() + close(closeDone) + }() + select { + case <-closeDone: + case <-time.After(2 * time.Second): + // Let the handshake fail so a lock-based implementation can unwind + // instead of leaking its blocked goroutines after the assertion. + releaseHandshake() + select { + case <-closeDone: + case <-time.After(2 * time.Second): + } + t.Fatal("Close did not wake NextPart(-1)") + } + releaseHandshake() + + select { + case err := <-nextDone: + var aimuxErr *Error + if !errors.As(err, &aimuxErr) || aimuxErr.Code != CodeAborted { + t.Fatalf("NextPart after Close: expected flat Aborted, got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("NextPart(-1) stayed blocked after Close returned") + } + // Idempotence remains part of the public session contract. + session.Close() +} + +func TestTranscriptionSessionCloseUnblocksBackpressuredPushAudio(t *testing.T) { + // Keeping the WebSocket handshake pending also keeps the driver from + // polling its 64-slot audio receiver, giving a deterministic full channel. + server, handshakeStarted, releaseHandshake := newBlockedHTTPServer(t, "", "") + defer releaseHandshake() + + model, err := NewOpenAITranscriptionWithBase( + "sk-test-fake-key", "gpt-realtime-whisper", server.URL, + ) + if err != nil { + t.Fatalf("NewOpenAITranscriptionWithBase failed: %v", err) + } + defer model.Close() + session, err := StartTranscriptionSession(model, nil) + if err != nil { + t.Fatalf("StartTranscriptionSession failed: %v", err) + } + + select { + case <-handshakeStarted: + case <-time.After(2 * time.Second): + t.Fatal("WebSocket handshake did not start") + } + for i := 0; i < 64; i++ { + if err := session.PushAudio([]byte{byte(i)}); err != nil { + session.Close() + t.Fatalf("PushAudio %d failed while filling the channel: %v", i, err) + } + } + + pushStarted := make(chan struct{}) + pushDone := make(chan error, 1) + go func() { + close(pushStarted) + pushDone <- session.PushAudio([]byte("backpressure")) + }() + <-pushStarted + select { + case err := <-pushDone: + session.Close() + t.Fatalf("PushAudio did not backpressure on a full channel: %v", err) + case <-time.After(200 * time.Millisecond): + } + + closeDone := make(chan struct{}) + go func() { + session.Close() + close(closeDone) + }() + select { + case <-closeDone: + case <-time.After(2 * time.Second): + releaseHandshake() + select { + case <-closeDone: + case <-time.After(2 * time.Second): + } + t.Fatal("Close did not wake backpressured PushAudio") + } + releaseHandshake() + + select { + case <-pushDone: + // The send may have linearized before teardown (success) or observe the + // closed receiver (error); this regression only requires bounded wakeup. + case <-time.After(2 * time.Second): + t.Fatal("backpressured PushAudio stayed blocked after Close returned") + } + session.Close() +} + // ── E2E: Embedding via mock server ─────────────────────────────────────────── func TestE2E_Embedding(t *testing.T) { diff --git a/bindings/go/trace.go b/bindings/go/trace.go index 0e6a09c1..e43d6508 100644 --- a/bindings/go/trace.go +++ b/bindings/go/trace.go @@ -12,6 +12,7 @@ package aimux import "C" import ( + "runtime" "unsafe" ) @@ -34,11 +35,11 @@ func (m *Model) TraceAudited(strict bool) (*Model, error) { } func (m *Model) traceWrap(audited bool, strict C.int) (*Model, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return nil, err } - defer release() + defer runtime.KeepAlive(m) var cerr C.AimuxError C.aimux_error_clear(&cerr) @@ -69,11 +70,11 @@ func (m *Model) TraceSessionChain(sessionId string) (string, error) { // TraceExportJsonl returns all probe records as JSONL (one TraceRecord per // line). func (m *Model) TraceExportJsonl() (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) return ffiString(func(cerr *C.AimuxError) *C.char { return C.aimux_trace_export_jsonl(C.uint64_t(handle), cerr) @@ -82,11 +83,11 @@ func (m *Model) TraceExportJsonl() (string, error) { // TraceClear drops all probe records of this model. func (m *Model) TraceClear() error { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return err } - defer release() + defer runtime.KeepAlive(m) if rc := C.aimux_trace_clear(C.uint64_t(handle)); rc != 0 { return newError(CodeInvalidArgument, "aimux: trace_clear failed (invalid handle)") @@ -96,11 +97,11 @@ func (m *Model) TraceClear() error { // traceQuery runs a query taking one C string argument. func (m *Model) traceQuery(arg string, call func(C.uint64_t, *C.char, *C.AimuxError) *C.char) (string, error) { - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { return "", err } - defer release() + defer runtime.KeepAlive(m) cArg := C.CString(arg) defer C.free(unsafe.Pointer(cArg)) diff --git a/bindings/go/typed.go b/bindings/go/typed.go index 995d8a1a..247c839e 100644 --- a/bindings/go/typed.go +++ b/bindings/go/typed.go @@ -111,18 +111,15 @@ type TypedStream struct { raw *Stream parts chan *StreamPart err error - done chan struct{} } // Parts returns a channel of parsed *StreamPart values. // The channel is closed when the stream ends. func (s *TypedStream) Parts() <-chan *StreamPart { return s.parts } -// Err returns any error that occurred during streaming. -func (s *TypedStream) Err() error { - <-s.done - return s.err -} +// Err returns any error that occurred during streaming. It must be called +// only after Parts has closed. +func (s *TypedStream) Err() error { return s.err } // Cancel stops this stream. It is safe to call more than once. func (s *TypedStream) Cancel() { s.raw.Cancel() } @@ -164,29 +161,36 @@ func (m *Model) StreamContext( ts := &TypedStream{ raw: rawStream, parts: make(chan *StreamPart, 256), - done: make(chan struct{}), } go func() { - defer close(ts.done) defer close(ts.parts) - for raw := range rawStream.Parts() { + for { + var raw string + select { + case part, ok := <-rawStream.Parts(): + if !ok { + ts.err = rawStream.Err() + return + } + raw = part + case <-rawStream.cancelCtx.Done(): + ts.err = context.Cause(rawStream.cancelCtx) + return + } sp, err := ParseStreamPart(raw) if err != nil { ts.err = err - rawStream.Cancel() + rawStream.cancel(err) return } select { case ts.parts <- sp: - case <-rawStream.entry.cancelled: - ts.err = rawStream.Err() + case <-rawStream.cancelCtx.Done(): + ts.err = context.Cause(rawStream.cancelCtx) return } } - if err := rawStream.Err(); err != nil { - ts.err = err - } }() return ts, nil @@ -224,18 +228,15 @@ type OpenAIStream struct { raw *Stream parts chan *ChatCompletionChunk err error - done chan struct{} } // Parts returns a channel of parsed *ChatCompletionChunk values. // The channel is closed when the stream ends. func (s *OpenAIStream) Parts() <-chan *ChatCompletionChunk { return s.parts } -// Err returns any error that occurred during streaming. -func (s *OpenAIStream) Err() error { - <-s.done - return s.err -} +// Err returns any error that occurred during streaming. It must be called +// only after Parts has closed. +func (s *OpenAIStream) Err() error { return s.err } // Cancel stops this stream. It is safe to call more than once. func (s *OpenAIStream) Cancel() { s.raw.Cancel() } @@ -278,29 +279,36 @@ func (m *Model) StreamAsOpenAIContext( ts := &OpenAIStream{ raw: rawStream, parts: make(chan *ChatCompletionChunk, 256), - done: make(chan struct{}), } go func() { - defer close(ts.done) defer close(ts.parts) - for raw := range rawStream.Parts() { + for { + var raw string + select { + case part, ok := <-rawStream.Parts(): + if !ok { + ts.err = rawStream.Err() + return + } + raw = part + case <-rawStream.cancelCtx.Done(): + ts.err = context.Cause(rawStream.cancelCtx) + return + } chunk, err := ParseChatCompletionChunk(raw) if err != nil { ts.err = err - rawStream.Cancel() + rawStream.cancel(err) return } select { case ts.parts <- chunk: - case <-rawStream.entry.cancelled: - ts.err = rawStream.Err() + case <-rawStream.cancelCtx.Done(): + ts.err = context.Cause(rawStream.cancelCtx) return } } - if err := rawStream.Err(); err != nil { - ts.err = err - } }() return ts, nil diff --git a/bindings/go/typed_test.go b/bindings/go/typed_test.go index de59056e..2852791c 100644 --- a/bindings/go/typed_test.go +++ b/bindings/go/typed_test.go @@ -222,16 +222,15 @@ func TestTypedStreamCancelUnblocksFullPartsChannel(t *testing.T) { t.Fatal("provider did not send the SSE burst") } waitForChannelFull(t, stream.parts) - select { - case <-stream.raw.entry.terminal: - t.Fatal("typed stream ended before cancellation") - default: - } stream.Cancel() stream.Cancel() done := make(chan error, 1) - go func() { done <- stream.Err() }() + go func() { + for range stream.Parts() { + } + done <- stream.Err() + }() select { case err := <-done: if !errors.Is(err, context.Canceled) { @@ -253,7 +252,7 @@ func TestDeepSeekFactory(t *testing.T) { if m == nil { t.Fatal("expected non-nil model") } - if m.handle == 0 { + if m.id.Load() == 0 { t.Fatal("expected non-zero handle") } } diff --git a/docs/api/go.md b/docs/api/go.md index 1afabced..4503320e 100644 --- a/docs/api/go.md +++ b/docs/api/go.md @@ -72,8 +72,16 @@ stream := model.StreamText(`"Write a haiku"`, "") for part := range stream.Parts() { fmt.Println(part) // StreamPart JSON } +if err := stream.Err(); err != nil { // only after Parts has closed + log.Fatal(err) +} ``` +Drain `Parts()` completely, or call `Cancel()` when stopping early. Context +variants connect cancellation automatically. `Cancel` releases callback +backpressure and aborts the native request; it never closes `Parts()` itself. +The producer closes `Parts()` after the blocking native call returns. + ## Providers All 250 registry-backed OpenAI-compatible providers are reachable by name; @@ -131,8 +139,16 @@ stream := model.StreamText(`"Write a haiku"`, "") for part := range stream.Parts() { fmt.Println(part) // StreamPart JSON } +if err := stream.Err(); err != nil { // only after Parts has closed + log.Fatal(err) +} ``` +Drain `Parts()` completely, or call `Cancel()` when stopping early. `Cancel` +aborts the native request and releases a callback blocked on a full channel; +the producer goroutine remains the only goroutine that closes `Parts()` and +publishes the terminal error. + > Stream part variants are documented in the [API overview](../API.md#streaming-generation). ## Vector Embedding @@ -356,7 +372,10 @@ fmt.Println(result.ProviderReference) // map["openai":"file-xxx"] All constructors come in two flavors: `NewXxx(...) (T, error)` (checked) and `Xxx(...) T` (unchecked, panics on failure). Every model type has a `Close()` -method that drops the underlying FFI handle. +method that atomically drops the underlying FFI handle. Handle wrappers must +not be copied after first use. `Close` is idempotent and does not wait for an +in-flight network call or stream; a racing call either enters Rust first and +continues with its cloned `Arc`, or receives an invalid-handle error. ### Constructors @@ -373,6 +392,25 @@ method that drops the underlying FFI handle. | `NewCohereReranking(key, modelID)` | `*RerankingModel` | | | `NewTavilySearch(key)` | `*SearchModel` | no model ID needed | | `NewOpenAIFiles(key)` | `*Files` | | +| `NewRouter(models []*Model, configJSON)` | `*Model` | RFC-0021 fallback router; `models` must be non-empty and may contain the same model more than once | +| `NewMoa(references []*Model, aggregator *Model, configJSON)` | `*Model` | RFC-0022 mixture-of-agents; references may be empty, repeated, or include the aggregator | + +#### Composite constructors and concurrency + +`NewRouter` and `NewMoa` snapshot each atomic handle in caller order and never +hold a Go lifecycle lock while calling C. Duplicate models require no special +case, and opposite caller orders cannot form an ABBA cycle because there is no +multi-lock protocol. Concurrent `Close` is resolved by the Rust registry: +construction either clones a model `Arc` or returns an invalid-handle error. + +A `nil` child or reference returns an error naming its position +(`aimux: models[i] is nil`); a `nil` MoA aggregator returns +`aimux: aggregator is nil`. Validation happens before the C call rather than +panicking on a nil pointer. + +Both constructors take a new reference to each child, so the caller keeps +ownership: closing a child afterwards does not invalidate the composite, and +the composite must be closed separately. ### Methods diff --git a/rfc/0011-golang-bindings.md b/rfc/0011-golang-bindings.md index d99fcc1e..acb3c992 100644 --- a/rfc/0011-golang-bindings.md +++ b/rfc/0011-golang-bindings.md @@ -105,7 +105,13 @@ aimux_stream_text → C callback on_part(json) → channel<- part → for - **Producer side (push)**: cgo calls the blocking `C.aimux_stream_text` in a separate goroutine; the C callback writes each StreamPart JSON to the Go channel - **Consumer side (pull)**: Go users consume with `for part := range ch`, conforming to Go conventions -- A buffered channel decouples the two in between, avoiding backpressure +- A bounded buffered channel decouples bursts while preserving backpressure; + callers that stop consuming must cancel the stream or its context +- The producer goroutine is the only terminal-state writer: after the blocking + C call returns (and all synchronous callbacks have returned), it stores the + terminal error and closes the parts channel. `Err` is read only after that + closure, so the channel close supplies the happens-before edge without a + separate mutex or done channel Compared to the streaming wrappers of other C ABI bindings: @@ -263,3 +269,4 @@ Reuse the shared JSON fixture [contract-tests/fixtures/wire-format.json](../cont | 2026-07-31 | v0.1 | **PoC landed**: `bindings/go/` complete implementation (aimux.go cgo declarations + Model + Generate/Stream; types.go typed JSON types; 19 tests all passing: 7 unit + 6 E2E + 6 contract subtests). Single binary empirically 8.7MB (after strip), statically linked `libaimux_ffi.a`, zero extra file dependencies | | 2026-07-31 | v0.1.1 | **Review fixes**: RWMutex handle lifecycle, stream fallback closeParts, error envelope JSON parsing, constructors return error, ParseStreamPart validation, numeric types aligned with Rust (uint32/uint64/float64), mock server io.ReadAll, contract default→Fatal | | 2026-07-31 | v0.2 | **Aligned with Node flagship coverage**: typed text API (Generate/Stream accept string\|[]ModelMessage + typed options, return *GenerateTextResult / typed StreamPart channel); 8-modality multimodal (Embedding/Speech/Image/Transcription/Files/Reranking/Video/Search) + factory functions + typed result types (aligned with ts-rs wire format); DeepSeek factory; 55 tests all passing (including -race) | +| 2026-08-18 | v0.3 | **Non-blocking lifecycle**: native handles use an atomic `0 = closed` owner instead of RWMutex/lock ordering; `Close` never waits for in-flight calls. Stream producers alone publish errors and close parts; explicit Cancel/context releases callback backpressure. | diff --git a/rfc/0028-transcription-streaming.md b/rfc/0028-transcription-streaming.md index 25d95332..adc3076b 100644 --- a/rfc/0028-transcription-streaming.md +++ b/rfc/0028-transcription-streaming.md @@ -235,7 +235,7 @@ FFI 层:mock provider 实现 `do_stream` → 验证 push/next_part/input_done/dr |---|---| | **Node**(native) | ~~streamTranscribe + ReadableStream/AsyncGenerator~~ → **实际落地为会话对象**(见 §10 D1):`startTranscriptionSession(model, optsJson?, bridge?)` 返回 `TranscriptionSession` 类(pushAudio/inputDone/nextPart/close) | | **Python**(native) | 同 D1:`start_transcription_session(model, opts_json?)` + `TranscriptionSession` pyclass(push_audio/input_done/next_part/close) | -| **Go/Swift/Java/Kotlin/Flutter**(C-ABI) | 各包一个 `TranscriptionSession` 类:start → pushAudio(bytes) → nextPart(timeoutMs) → inputDone → close。模式同各自的 Model wrapper(锁 + handle);Go 用 RWMutex 允许并发 push/pull(双向会话不死锁) | +| **Go/Swift/Java/Kotlin/Flutter**(C-ABI) | 各包一个 `TranscriptionSession` 类:start → pushAudio(bytes) → nextPart(timeoutMs) → inputDone → close。Go 以原子 handle（`0 = closed`）实现：操作只读取 handle 后直接进入 C，不持 Go 生命周期锁；`Close` 用 `Swap(0)` 取得所有权并立即调用 native session drop/abort，从而唤醒阻塞的 push/pull。其他四个 binding 使用各自的原生资源包装。 | Node/Python 的 native 路径不走 Phase 2 的 FFI 会话(直连 core,与 generateText 同架构);C-ABI 5 语言消费 Phase 2。 @@ -284,4 +284,5 @@ Node/Python 的 native 路径不走 Phase 2 的 FFI 会话(直连 core,与 gener 2. **D2 — `TranscriptionStreamOptions` 增加 `timeout` 字段**(R1 B1):设计稿的超时只在 ws 层;实现把 `Option` 提到 options(FFI opts_json 的 `"timeout"` 对象 / Node/Python opts 同),`first_chunk_ms` 为 connect+首事件**合并预算**(锚定在 connect 前,connect 后只花余额)。 3. **D3 — peer close 携带 code/reason**(R1 S3):设计稿只说"close frame → ApiCall 带 code";实现为 `"websocket closed by peer (code NNN: reason)"`。 4. **D4 — live-API smoke 未执行**:§3.3 的"用真实 key 跑一次最小会话"需要 API key,未在 CI/本地执行。wire 形状已由本地 WS 集成测试逐字段断言(含 session.update 嵌套结构与 turn_detection 嵌套位置);首次真实调用时如遇 shape 偏差请回报此 RFC。 -5. **D5 — 其余 review 修正**:ws close() 5s 有界;chunk-idle 窗口按 next() 调用计算(ping 不续命);FFI 驱动的 futures-mpsc `flush()` 不可用于断连检测(改 `try_send + is_disconnected`);五个 C-ABI 绑定的 timeout 路径先消费错误串(防泄漏);Go 会话 RWMutex(并发 push/pull);Python push/next 释放 GIL(allow_threads)。 +5. **D5 — 其余 review 修正**:ws close() 5s 有界;chunk-idle 窗口按 next() 调用计算(ping 不续命);FFI 驱动的 futures-mpsc `flush()` 不可用于断连检测(改 `try_send + is_disconnected`);五个 C-ABI 绑定的 timeout 路径先消费错误串(防泄漏);Python push/next 释放 GIL(allow_threads)。 +6. **D6 — Go 会话生命周期改为原子 owner**(2026-08-18):早期实现用 RWMutex 保护 handle，但 `NextPart(-1)` / 满 audio channel 的 `PushAudio` 会持读锁阻塞，导致 `Close` 无法取得写锁去触发恰好用于唤醒它们的 native abort。现改为 `atomic.Uint64`：方法只 snapshot handle，`Close` 原子清零并立即 session drop/abort；Rust registry 的 `Arc` clone 负责已进入调用的内存安全，不再存在 Go 锁等待环。