From ef6fb876eb8f6c71249177d16278ae3a0a38be00 Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Tue, 18 Aug 2026 18:14:59 +0800 Subject: [PATCH 1/3] fix(ffi): resolve router/moa members under one registry lock; reject any invalid handle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit aimux_router_new and aimux_moa_new looked each handle up separately and skipped the ones they could not find, failing only when *all* were unknown. A concurrent drop between two lookups — or a caller passing a dead handle — therefore built a composite with fewer members than requested, silently. get_models() now takes the registry lock once, clones every requested Arc, and the constructor fails on the first unknown handle. Membership is all-or-nothing. MoA also rejects a NULL reference pointer with ref_len > 0 instead of reading it as "no references". --- aimux-ffi/aimux-ffi.h | 7 ++-- aimux-ffi/src/lib.rs | 90 ++++++++++++++++++++++++++++--------------- 2 files changed, 64 insertions(+), 33 deletions(-) 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). From 83dd482bfbde7fce4730d00be720a9bc7455124e Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Tue, 18 Aug 2026 18:18:26 +0800 Subject: [PATCH 2/3] fix(go)!: atomic handle owners; Close never waits for an in-flight C call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Go binding held sync.RWMutex read locks across blocking C calls, which formed real wait cycles: TranscriptionSession.Close could not reach the native abort while NextPart(-1) or a backpressured PushAudio held the read lock, and a full stream channel pinned the C stream call under the read lock while Model.Close waited for the write lock. Router/MoA had the opposite defect: they read each child's handle under the lock and released it before the C call, so a concurrent Close in that window freed a handle mid-construction and the FFI silently dropped that child. Every native owner (Model, ProviderHandle, the multimodal models, TranscriptionSession) is now an atomic.Uint64 with 0 = closed. A call snapshots the id and KeepAlives the owner until C returns; Close is Swap(0) + drop and never waits for Go state — a call that already entered the Rust registry owns a cloned Arc and finishes, a later one gets the closed error. Router/MoA snapshot ids in caller order and rely on the FFI's all-or-nothing lookup; no multi-lock protocol. The stream producer keeps its structure and merely stops holding a lock across the C call. Breaking: Close is no longer a join; owners must not be copied after first use (go vet now flags it); Router/MoA reject any invalid member instead of dropping it. Finalizers stay a best-effort leak fallback. --- bindings/go/aimux.go | 190 +++++++-------- bindings/go/aimux_test.go | 350 +++++++++++++++++++++++++++- bindings/go/e2e_test.go | 12 +- bindings/go/multimodal.go | 239 ++++++++++--------- bindings/go/multimodal_test.go | 210 +++++++++++++++++ bindings/go/trace.go | 17 +- bindings/go/typed_test.go | 13 +- docs/api/go.md | 40 +++- rfc/0011-golang-bindings.md | 9 +- rfc/0028-transcription-streaming.md | 5 +- 10 files changed, 853 insertions(+), 232 deletions(-) diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index a6e842c8..b07e4eab 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -19,7 +19,7 @@ package aimux #include "aimux-ffi.h" -// Go-side callback trampolines (//export below). stream_ctx is the stream id. +// Go-side callback trampoline (//export below). stream_ctx is a cgo.Handle. extern void goStreamPart(int64_t id, char* json); extern void goStreamDone(int64_t id); @@ -62,6 +62,7 @@ import ( "io" "runtime" "sync" + "sync/atomic" "unsafe" ) @@ -72,43 +73,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 +362,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 +396,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 +456,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 +653,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 +737,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 +772,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 +799,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 +826,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 +853,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)) @@ -1027,12 +1031,12 @@ func (m *Model) StreamTextContext(ctx context.Context, promptJson, optsJson stri entry.closeParts() }() - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { entry.markTerminal(err, false) return } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) @@ -1105,12 +1109,12 @@ func (m *Model) StreamTextAsOpenAIContext(ctx context.Context, promptJson, optsJ entry.closeParts() }() - handle, release, err := m.acquireHandle() + handle, err := m.handle() if err != nil { entry.markTerminal(err, false) return } - defer release() + defer runtime.KeepAlive(m) cPrompt := C.CString(promptJson) defer C.free(unsafe.Pointer(cPrompt)) diff --git a/bindings/go/aimux_test.go b/bindings/go/aimux_test.go index 54df5e90..d34868a7 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,41 @@ func TestStreamTextContextAlreadyCanceled(t *testing.T) { stream.Cancel() } +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 +221,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 +244,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 +333,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..1da70dc4 100644 --- a/bindings/go/e2e_test.go +++ b/bindings/go/e2e_test.go @@ -512,14 +512,12 @@ func TestStreamCancelUnblocksFullPartsChannel(t *testing.T) { 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: - } - 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_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 锁等待环。 From 80cebbc5c128072b2aee395c799890c7a13a9334 Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Tue, 18 Aug 2026 18:18:26 +0800 Subject: [PATCH 3/3] =?UTF-8?q?refactor(go):=20stream=20internals=20?= =?UTF-8?q?=E2=80=94=20single=20writer,=20cgo.Handle,=20context=20cancel?= =?UTF-8?q?=20cause?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No behaviour change beyond preserving the caller's cancellation cause. The producer goroutine is now the only writer of err and the only closer of Parts, so the channel close is the publication point and Err needs no lock; goStreamDone becomes a no-op (the C call returns right after on_done). The handwritten id → entry registry and its mutex are replaced by runtime/cgo.Handle; the cancelled channel + once + cause fields by context.WithCancelCause; the watcher goroutine by context.AfterFunc. The producer re-checks context.Cause after the C call returns so a cancellation that lands during the last part is not lost. Typed and OpenAI adapters follow the same shape. Stream/TypedStream/OpenAIStream.Err must be read after Parts closes. --- bindings/go/aimux.go | 338 ++++++++++++-------------------------- bindings/go/aimux_test.go | 15 ++ bindings/go/e2e_test.go | 2 +- bindings/go/typed.go | 68 ++++---- 4 files changed, 163 insertions(+), 260 deletions(-) diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index b07e4eab..22289832 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -20,36 +20,34 @@ package aimux #include "aimux-ffi.h" // Go-side callback trampoline (//export below). stream_ctx is a cgo.Handle. -extern void goStreamPart(int64_t id, char* json); -extern void goStreamDone(int64_t id); +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,7 @@ import ( "fmt" "io" "runtime" - "sync" + "runtime/cgo" "sync/atomic" "unsafe" ) @@ -875,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. @@ -969,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. @@ -1002,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, err := m.handle() - if err != nil { - entry.markTerminal(err, false) - return - } - 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)) - } - - 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 @@ -1081,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, err := m.handle() - 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 runtime.KeepAlive(m) + 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). @@ -1192,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 d34868a7..7a6202e3 100644 --- a/bindings/go/aimux_test.go +++ b/bindings/go/aimux_test.go @@ -172,6 +172,21 @@ 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) { diff --git a/bindings/go/e2e_test.go b/bindings/go/e2e_test.go index 1da70dc4..40ef89b5 100644 --- a/bindings/go/e2e_test.go +++ b/bindings/go/e2e_test.go @@ -511,7 +511,7 @@ 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) + waitForChannelFull(t, stream.parts) tasks := []func(){func() { m.Close() }} for i := 0; i < 32; i++ { 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