diff --git a/CHANGELOG.md b/CHANGELOG.md index 07c46f22ce..25da0ddfd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- **Host aborts** - The C ABI accepts an `abort` reply, and the Go binding adds + `Abort(err)` to stop routing when the host refuses a model or decision call, + even if the algorithm recovers from ordinary call failures. - **VGR routing (planned for 0.4.0)** — adds `type = "vgr"` deployment routes and the `AlgorithmSpec::Vgr` variant. **Rust source-breaking change:** callers that match exhaustively on `AlgorithmSpec` must handle the new variant. diff --git a/bindings/go/README.md b/bindings/go/README.md index 56b4ec1792..39ee6b6bdf 100644 --- a/bindings/go/README.md +++ b/bindings/go/README.md @@ -4,17 +4,16 @@ NVIDIA NeMo Switchyard is a library for routing LLM requests between models. These Go bindings run Switchyard algorithms inside your application. Your application owns model eligibility, credentials, provider calls, and retries. -A private WebAssembly runtime exercises translation codecs through wazero. -Routing and public translation functions still use the native library. +The default backend runs libsy in WebAssembly through wazero, with no cgo +dependency. Native builds can use the C ABI instead. ## Installation -Requires Go 1.25+, cgo, a C compiler, and the repository's Rust toolchain. This is -a source-only module. From a checkout of this repository, build the native -library for the same platform as your Go application, then build the WASM guest: +Requires Go 1.25+ and the repository's Rust toolchain to build the WASM guest. +This is a source-only module. From a checkout of this repository, build the guest +before compiling the Go package: ```sh -cargo build -p switchyard-ffi --release --locked rustup target add wasm32-wasip1 cargo build -p switchyard-ffi --target wasm32-wasip1 --release --locked --lib mkdir -p bindings/go/internal/wasm @@ -32,13 +31,19 @@ go mod edit -replace=github.com/NVIDIA-NeMo/Switchyard/bindings/go=/absolute/pat go get github.com/NVIDIA-NeMo/Switchyard/bindings/go ``` +For the native backend, build the Rust library for the same platform as your Go +application. This path also requires cgo and a C compiler: + +```sh +cargo build -p switchyard-ffi --release --locked +``` + Native builds link `target/release/libswitchyard_ffi.a` into the Go host artifact. The native path still requires cgo and platform system libraries. Native artifacts are not downloaded or published by this package. -The `switchyard_native` build tag excludes the private WASM runtime and its -tests, so that build needs only the native library. The package still requires -cgo; the private runtime is not a complete WASM routing backend. +The native backend uses the `switchyard_native` build tag and does not need the +WASM binary. These commands build locally; they do not publish artifacts. ## Use @@ -63,6 +68,9 @@ outcome, err := router.Route(ctx, switchyard.Input{ Import `encoding/json` and `github.com/NVIDIA-NeMo/Switchyard/bindings/go`. +The backend is selected at build time: `cgo` with `switchyard_native` uses the +native library; other builds use WASM. Both use `New(config)` and the same routing API. + Supply the model groups needed by the configured algorithm. Composite needs `any` for eligible answer targets, plus `capable`, `efficient`, and `judge`. Random and signal-only Stage need no host callback. For Composite, supply a @@ -73,6 +81,11 @@ Random and signal-only Stage need no host callback. For Composite, supply a - `call_decision`: send `Work.Request` to `Work.Model` and return its normalized decision response. `ServedModel` does not apply to decision calls. +A `HostCall` error is a failed call. The algorithm's error policy decides what +happens next, so a `fail_open` classifier can still route. Return `Abort(err)` +when the host refuses the call, such as for a budget or authorization check. +`Route` then stops and returns `err`, which keeps its Go error type. + Requests and replies use Switchyard's protocol JSON. `DecodeRequest`, `EncodeRequest`, and `DecodeResponse` use the existing translation codecs for `openai_chat`, `openai_responses`, and `anthropic_messages`. Typed decision @@ -87,12 +100,27 @@ the final selected provider response separately. `Route` waits in its calling goroutine. A host that needs immediate return can call it in a goroutine; waiting for routing steps does not occupy a blocked -native-reader thread. Concurrent routes share one router. +native-reader thread. Each WASM router uses a pool of `GOMAXPROCS` guest instances. +A session ID keeps the parent and all its child agents on the same instance. +Without a session ID, the first user message is used, or requests are assigned round-robin. +Guest calls are serialized within each instance; host callbacks can overlap. +A trap fails runs sharing that instance. It is replaced and its retained state is lost. + +Each instance has its own algorithm state. VGR's circuit breaker is not shared +across the pool. Advisor's review budget is also per instance when no session ID +is supplied. Enforce gateway-wide spending limits in the host. +Message-based assignment does not give requests without session IDs persistent +session state. + +While a WASM run is idle, Go re-enters it every 5 ms to advance Rust timers and +session waiters. Deadline handling can be delayed by that interval and scheduling load. Host callbacks may run concurrently and must honor their context. Cancellation -stops the native run and cancels the host context, but cannot forcibly stop Go -code that ignores it. Late host results cannot access the freed native run. +stops routing and cancels the host context, but cannot forcibly stop Go code that +ignores it. WASM checks cancellation between guest calls. Late host results cannot +access a freed run. Host errors cross the boundary as messages, not concrete Go error types. +Only an aborted call returns the host's original error from `Route`. `Close` prevents new runs; existing runs retain their algorithm and finish under their own contexts. Cancel those contexts to stop them. Router value copies @@ -100,22 +128,23 @@ share closed state, and repeated `Close` calls are safe. Native callback handles are released only after native cleanup has waited for callbacks to return. Call `Close` when finished; Go garbage collection does not release the router. -`New` checks ABI generation 1 against the linked library. Scope session and agent +`New` checks ABI generation 1 against the native library or WASM module. Scope session and agent IDs to the host's tenant and conversation. `Input.Metadata` and `Input.SubagentModels` carry routing context separately from provider requests. ## Test -From the repository root, after building the native library and WASM guest: +From the repository root, after building the WASM guest: ```sh cd bindings/go +CGO_ENABLED=0 go test ./... +CGO_ENABLED=0 go vet ./... CGO_ENABLED=1 go test -race ./... -CGO_ENABLED=1 GOEXPERIMENT=cgocheck2 go test -race ./... -CGO_ENABLED=1 go vet ./... ``` -For a native-only build, skip the guest build and use the build tag: +For a native build, build the native library above instead of the guest and use +the build tag: ```sh cd bindings/go diff --git a/bindings/go/backend_native.go b/bindings/go/backend_native.go index 2864c70e14..380d62ad9e 100644 --- a/bindings/go/backend_native.go +++ b/bindings/go/backend_native.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//go:build cgo +//go:build cgo && switchyard_native package switchyard @@ -150,17 +150,7 @@ func (backend *nativeBackend) route(ctx context.Context, input Input, call HostC return Outcome{}, err } case reply := <-replies: - envelope := struct { - Response json.RawMessage `json:"response,omitempty"` - ServedModel string `json:"served_model,omitempty"` - Error *string `json:"error,omitempty"` - }{Response: reply.reply.Response, ServedModel: reply.reply.ServedModel} - if reply.err != nil { - envelope.Response = nil - message := reply.err.Error() - envelope.Error = &message - } - data, err := json.Marshal(envelope) + data, err := hostReplyJSON(reply.reply, reply.err) if err != nil { return Outcome{}, err } @@ -170,6 +160,9 @@ func (backend *nativeBackend) route(ctx context.Context, input Input, call HostC value := C.CString(string(data)) err = nativeError(C.sy_run_respond(run, C.uint64_t(reply.id), value)) C.free(unsafe.Pointer(value)) + if abort := hostAbort(reply.err); abort != nil { + return Outcome{}, abort + } if err != nil { return Outcome{}, err } diff --git a/bindings/go/backend_test.go b/bindings/go/backend_test.go index 324c8a2723..4e7ae75f9d 100644 --- a/bindings/go/backend_test.go +++ b/bindings/go/backend_test.go @@ -84,7 +84,7 @@ func testRoutingDeadline(t *testing.T, router *Router) { } } -// Host calls must overlap without blocking other requests. +// Waiting for a provider must leave the shared guest available to other requests. func TestConcurrentHostCalls(t *testing.T) { router := newRouter(t, composite) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) diff --git a/bindings/go/backend_wazero.go b/bindings/go/backend_wazero.go new file mode 100644 index 0000000000..30ffa6150c --- /dev/null +++ b/bindings/go/backend_wazero.go @@ -0,0 +1,554 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !cgo || !switchyard_native + +package switchyard + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "hash/fnv" + "runtime" + "strings" + "sync" + "sync/atomic" + "time" +) + +type wasmBackend struct { + runtime *wasmRuntime + slots []*wasmSlot + next atomic.Uint64 + active atomic.Int64 + closed atomic.Bool + lifecycle sync.Mutex + finalize sync.Once +} + +type wasmSlot struct { + mu sync.Mutex + owner *wasmBackend + config []byte + gen *wasmGeneration + err error +} + +type wasmGeneration struct { + guest *guestModule + router uint32 + runs map[*wasmRun]struct{} + done chan struct{} + errMu sync.RWMutex + err error + poisoned bool +} + +func (generation *wasmGeneration) error() error { + generation.errMu.RLock() + defer generation.errMu.RUnlock() + return generation.err +} + +func (generation *wasmGeneration) setError(err error) { + generation.errMu.Lock() + generation.err = err + generation.errMu.Unlock() +} + +type wasmRun struct { + slot *wasmSlot + gen *wasmGeneration + ptr uint32 + ctx context.Context + cancel context.CancelFunc +} + +func newBackend(config json.RawMessage) (engineBackend, error) { + return newWasmBackendWithPoolSize(config, runtime.GOMAXPROCS(0)) +} + +func newWasmBackendWithPoolSize(config json.RawMessage, size int) (*wasmBackend, error) { + if size < 1 { + return nil, errors.New("Switchyard wasm instance pool must have at least one instance") + } + runtime, err := newWasmRuntime() + if err != nil { + return nil, err + } + backend := &wasmBackend{runtime: runtime, slots: make([]*wasmSlot, size)} + for index := range backend.slots { + slot := &wasmSlot{owner: backend, config: bytesClone(config)} + slot.gen, err = backend.newGeneration(context.Background(), slot.config) + if err != nil { + backend.close() + return nil, fmt.Errorf("initialize Switchyard wasm instance %d: %w", index, err) + } + backend.slots[index] = slot + } + return backend, nil +} + +func bytesClone(value []byte) []byte { + return append([]byte(nil), value...) +} + +func (backend *wasmBackend) newGeneration(ctx context.Context, config []byte) (*wasmGeneration, error) { + guest, err := backend.runtime.newModule(ctx) + if err != nil { + return nil, err + } + gen := &wasmGeneration{ + guest: guest, + runs: make(map[*wasmRun]struct{}), + done: make(chan struct{}), + } + gen.router, err = gen.newRouter(ctx, config) + if err != nil { + _ = guest.module.Close(context.Background()) + return nil, err + } + return gen, nil +} + +func (generation *wasmGeneration) newRouter(ctx context.Context, config []byte) (uint32, error) { + inputPtr, inputLen, err := generation.guest.cstring(ctx, config) + if err != nil { + return 0, err + } + defer generation.guest.free(context.Background(), inputPtr, inputLen) + outPtr, err := generation.guest.alloc(ctx, make([]byte, 4)) + if err != nil { + return 0, err + } + defer generation.guest.free(context.Background(), outPtr, 4) + values, err := generation.guest.call(ctx, "sy_router_new", uint64(inputPtr), uint64(outPtr)) + if err != nil { + return 0, err + } + if len(values) != 1 { + return 0, errors.New("Switchyard wasm router creation returned an invalid result") + } + if err := generation.guest.guestError(uint32(values[0])); err != nil { + return 0, err + } + ptr, ok := generation.guest.module.Memory().ReadUint32Le(outPtr) + if !ok || ptr == 0 { + return 0, errors.New("Switchyard wasm returned no router handle") + } + return ptr, nil +} + +func (generation *wasmGeneration) newRun(ctx context.Context, router uint32, input []byte) (uint32, error) { + inputPtr, inputLen, err := generation.guest.cstring(ctx, input) + if err != nil { + return 0, err + } + defer generation.guest.free(context.Background(), inputPtr, inputLen) + outPtr, err := generation.guest.alloc(ctx, make([]byte, 4)) + if err != nil { + return 0, err + } + defer generation.guest.free(context.Background(), outPtr, 4) + values, err := generation.guest.call(ctx, "sy_run_new", uint64(router), uint64(inputPtr), uint64(outPtr)) + if err != nil { + return 0, err + } + if len(values) != 1 { + return 0, errors.New("Switchyard wasm run creation returned an invalid result") + } + if err := generation.guest.guestError(uint32(values[0])); err != nil { + return 0, err + } + ptr, ok := generation.guest.module.Memory().ReadUint32Le(outPtr) + if !ok || ptr == 0 { + return 0, errors.New("Switchyard wasm returned no run handle") + } + return ptr, nil +} + +func (generation *wasmGeneration) next(ctx context.Context, run uint32) ([]byte, error) { + outPtr, err := generation.guest.alloc(ctx, make([]byte, 4)) + if err != nil { + return nil, err + } + defer generation.guest.free(context.Background(), outPtr, 4) + values, err := generation.guest.call(ctx, "sy_run_next", uint64(run), uint64(outPtr)) + if err != nil { + return nil, err + } + if len(values) != 1 { + return nil, errors.New("Switchyard wasm next returned an invalid result") + } + if err := generation.guest.guestError(uint32(values[0])); err != nil { + return nil, err + } + ptr, ok := generation.guest.module.Memory().ReadUint32Le(outPtr) + if !ok { + return nil, errors.New("Switchyard wasm next step is outside guest memory") + } + if ptr == 0 { + // The runtime is idle. + return nil, nil + } + output, err := generation.guest.readCString(ptr) + if err != nil { + return nil, err + } + if _, err := generation.guest.call(context.Background(), "sy_string_free", uint64(ptr)); err != nil { + return nil, err + } + return output, nil +} + +func (generation *wasmGeneration) respond(ctx context.Context, run uint32, id uint64, reply []byte) error { + replyPtr, replyLen, err := generation.guest.cstring(ctx, reply) + if err != nil { + return err + } + defer generation.guest.free(context.Background(), replyPtr, replyLen) + values, err := generation.guest.call(ctx, "sy_run_respond", uint64(run), id, uint64(replyPtr)) + if err != nil { + return err + } + if len(values) != 1 { + return errors.New("Switchyard wasm respond returned an invalid result") + } + return generation.guest.guestError(uint32(values[0])) +} + +func (backend *wasmBackend) close() { + backend.lifecycle.Lock() + if !backend.closed.CompareAndSwap(false, true) { + backend.lifecycle.Unlock() + return + } + finalize := backend.active.Load() == 0 + backend.lifecycle.Unlock() + if finalize { + backend.finalizeClose() + } +} + +func (backend *wasmBackend) finalizeClose() { + backend.finalize.Do(func() { + for _, slot := range backend.slots { + if slot == nil { + continue + } + slot.mu.Lock() + slot.closeLocked(errors.New("Switchyard router is closed")) + slot.mu.Unlock() + } + backend.runtime.closeAll() + }) +} + +func (slot *wasmSlot) closeLocked(cause error) { + gen := slot.gen + if gen == nil || gen.poisoned { + return + } + gen.poisoned = true + gen.setError(cause) + close(gen.done) + for run := range gen.runs { + run.cancel() + } + if gen.router != 0 { + _, _ = gen.guest.call(context.Background(), "sy_router_free", uint64(gen.router)) + gen.router = 0 + } + _ = gen.guest.module.Close(context.Background()) +} + +func (slot *wasmSlot) poisonLocked(gen *wasmGeneration, cause error) { + if slot.gen != gen || gen.poisoned { + return + } + gen.poisoned = true + state := "trapped" + if errors.Is(cause, context.Canceled) || errors.Is(cause, context.DeadlineExceeded) { + state = "closed" + } + gen.setError(fmt.Errorf("Switchyard wasm instance %s: %w", state, cause)) + close(gen.done) + for run := range gen.runs { + run.cancel() + } + _ = gen.guest.module.Close(context.Background()) + if slot.owner.closed.Load() { + return + } + replacement, err := slot.owner.newGeneration(context.Background(), slot.config) + if err != nil { + slot.err = fmt.Errorf("replace trapped Switchyard wasm instance: %w", err) + return + } + slot.gen = replacement + slot.err = nil +} + +func (backend *wasmBackend) route(ctx context.Context, input Input, call HostCall) (Outcome, error) { + if backend.closed.Load() { + return Outcome{}, errors.New("Switchyard router is closed") + } + data, err := json.Marshal(input) + if err != nil { + return Outcome{}, err + } + run, err := backend.slotFor(input).startRun(ctx, data) + if err != nil { + return Outcome{}, err + } + defer run.close() + // Host calls run concurrently, as on the native backend; replies are delivered + // whenever the guest reports that it needs one to progress. + replies := make(chan wasmHostReply) + var wake <-chan time.Time + for { + if err := ctx.Err(); err != nil { + return Outcome{}, err + } + stepData, err := run.next(ctx) + if err != nil { + if ctx.Err() != nil { + return Outcome{}, ctx.Err() + } + return Outcome{}, err + } + if stepData == nil { + if wake == nil { + // Re-enter the guest so Rust timers and shared session state can progress. + timer := time.NewTicker(5 * time.Millisecond) + defer timer.Stop() + wake = timer.C + } + select { + case reply := <-replies: + if err := run.deliver(ctx, reply); err != nil { + return Outcome{}, err + } + case <-wake: + case <-run.ctx.Done(): + if err := run.gen.error(); err != nil { + return Outcome{}, err + } + return Outcome{}, run.ctx.Err() + case <-run.gen.done: + return Outcome{}, run.gen.error() + } + continue + } + var step struct { + Type string `json:"type"` + ID uint64 `json:"id"` + Models []string `json:"models"` + Model string `json:"model"` + Request json.RawMessage `json:"request"` + } + if err := json.Unmarshal(stepData, &step); err != nil { + return Outcome{}, err + } + switch step.Type { + case "done": + var outcome Outcome + if err := json.Unmarshal(stepData, &outcome); err != nil { + return Outcome{}, err + } + if err := ctx.Err(); err != nil { + return Outcome{}, err + } + return outcome, nil + case "call_model", "call_decision": + if call == nil { + return Outcome{}, errors.New("Switchyard algorithm requires a host model client") + } + work := Work{Type: step.Type, ID: step.ID, Models: step.Models, Model: step.Model, Request: step.Request} + go run.hostCall(call, work, replies) + default: + return Outcome{}, errors.New("unknown Switchyard step: " + step.Type) + } + } +} + +func (backend *wasmBackend) slotFor(input Input) *wasmSlot { + if metadata := input.Metadata; metadata != nil && metadata.SessionID != "" { + return backend.slotForKey("session\x00" + metadata.SessionID) + } + // Without a session ID, keep message-keyed retention on one instance. + if text, ok := firstUserText(input.Request); ok { + return backend.slotForKey("message\x00" + text) + } + index := backend.next.Add(1) - 1 + return backend.slots[index%uint64(len(backend.slots))] +} + +func (backend *wasmBackend) slotForKey(key string) *wasmSlot { + hash := fnv.New64a() + _, _ = hash.Write([]byte(key)) + return backend.slots[hash.Sum64()%uint64(len(backend.slots))] +} + +// firstUserText matches libsy's Message::text_content("") on the first user message. +func firstUserText(request json.RawMessage) (string, bool) { + var parsed struct { + Messages []struct { + Role string `json:"role"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } `json:"messages"` + } + if json.Unmarshal(request, &parsed) != nil { + return "", false + } + for _, message := range parsed.Messages { + if message.Role != "user" { + continue + } + var text strings.Builder + found := false + for _, block := range message.Content { + if block.Type == "text" || block.Type == "refusal" { + text.WriteString(block.Text) + found = true + } + } + return text.String(), found + } + return "", false +} + +func (slot *wasmSlot) startRun(ctx context.Context, input []byte) (*wasmRun, error) { + slot.mu.Lock() + defer slot.mu.Unlock() + slot.owner.lifecycle.Lock() + defer slot.owner.lifecycle.Unlock() + if slot.owner.closed.Load() { + return nil, errors.New("Switchyard router is closed") + } + if err := ctx.Err(); err != nil { + return nil, err + } + if slot.err != nil { + replacement, err := slot.owner.newGeneration(context.Background(), slot.config) + if err != nil { + slot.err = fmt.Errorf("replace trapped Switchyard wasm instance: %w", err) + return nil, slot.err + } + slot.gen, slot.err = replacement, nil + } + gen := slot.gen + if gen == nil || gen.poisoned { + return nil, errors.New("Switchyard wasm instance is unavailable") + } + // Guest calls never see the caller's context: wazero closes the whole shared module + // when a call's context ends. The route loop checks ctx between steps instead. + ptr, err := gen.newRun(context.Background(), gen.router, input) + if err != nil { + if gen.guest.module.IsClosed() || isWasmTrap(err) { + slot.poisonLocked(gen, err) + } + if isWasmTrap(err) { + return nil, gen.error() + } + return nil, err + } + runCtx, cancel := context.WithCancel(ctx) + run := &wasmRun{slot: slot, gen: gen, ptr: ptr, ctx: runCtx, cancel: cancel} + gen.runs[run] = struct{}{} + slot.owner.active.Add(1) + return run, nil +} + +func (run *wasmRun) next(ctx context.Context) ([]byte, error) { + slot := run.slot + slot.mu.Lock() + defer slot.mu.Unlock() + if slot.gen != run.gen || run.gen.poisoned { + return nil, run.gen.error() + } + if err := ctx.Err(); err != nil { + return nil, err + } + output, err := run.gen.next(context.Background(), run.ptr) + if err != nil { + if run.gen.guest.module.IsClosed() || isWasmTrap(err) { + slot.poisonLocked(run.gen, err) + return nil, run.gen.error() + } + return nil, err + } + return output, nil +} + +func (run *wasmRun) respond(ctx context.Context, id uint64, reply []byte) error { + slot := run.slot + slot.mu.Lock() + defer slot.mu.Unlock() + if slot.gen != run.gen || run.gen.poisoned { + return run.gen.error() + } + if err := ctx.Err(); err != nil { + return err + } + if err := run.gen.respond(context.Background(), run.ptr, id, reply); err != nil { + if run.gen.guest.module.IsClosed() || isWasmTrap(err) { + slot.poisonLocked(run.gen, err) + return run.gen.error() + } + return err + } + return nil +} + +type wasmHostReply struct { + id uint64 + reply Reply + err error +} + +func (run *wasmRun) hostCall(call HostCall, work Work, replies chan<- wasmHostReply) { + reply, err := call(run.ctx, work) + select { + case replies <- wasmHostReply{work.ID, reply, err}: + case <-run.ctx.Done(): + } +} + +func (run *wasmRun) deliver(ctx context.Context, reply wasmHostReply) error { + data, err := hostReplyJSON(reply.reply, reply.err) + if err != nil { + return err + } + err = run.respond(ctx, reply.id, data) + if abort := hostAbort(reply.err); abort != nil { + return abort + } + if err != nil { + if reply.err != nil { + return fmt.Errorf("host call failed: %w", reply.err) + } + return err + } + return nil +} + +func (run *wasmRun) close() { + run.cancel() + slot := run.slot + slot.mu.Lock() + delete(run.gen.runs, run) + if slot.gen == run.gen && !run.gen.poisoned { + _, _ = run.gen.guest.call(context.Background(), "sy_run_cancel", uint64(run.ptr)) + _, _ = run.gen.guest.call(context.Background(), "sy_run_free", uint64(run.ptr)) + } + slot.mu.Unlock() + if slot.owner.active.Add(-1) == 0 && slot.owner.closed.Load() { + slot.owner.finalizeClose() + } +} diff --git a/bindings/go/callbacks.go b/bindings/go/callbacks.go index 31eb616781..97a15fc2c6 100644 --- a/bindings/go/callbacks.go +++ b/bindings/go/callbacks.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//go:build cgo +//go:build cgo && switchyard_native package switchyard diff --git a/bindings/go/switchyard.go b/bindings/go/switchyard.go index 18b0c82a50..eddcbdf47c 100644 --- a/bindings/go/switchyard.go +++ b/bindings/go/switchyard.go @@ -59,8 +59,52 @@ type Reply struct { } // HostCall executes model or decision work. It must honor ctx and may run concurrently. +// Return an error to let the algorithm apply its error policy, or Abort(err) to stop routing. type HostCall func(context.Context, Work) (Reply, error) +type abortError struct{ err error } + +func (e abortError) Error() string { return e.err.Error() } +func (e abortError) Unwrap() error { return e.err } + +// Abort marks a host call refusal. Route stops and returns err, even when the +// algorithm would recover from a failed call. A nil error remains nil. +func Abort(err error) error { + if err == nil { + return nil + } + return abortError{err} +} + +// hostReplyJSON encodes a host result for sy_run_respond. +func hostReplyJSON(reply Reply, err error) ([]byte, error) { + envelope := struct { + Response json.RawMessage `json:"response,omitempty"` + ServedModel string `json:"served_model,omitempty"` + Error *string `json:"error,omitempty"` + Abort *string `json:"abort,omitempty"` + }{Response: reply.Response, ServedModel: reply.ServedModel} + if err != nil { + envelope.Response = nil + message := err.Error() + if errors.As(err, new(abortError)) { + envelope.Abort = &message + } else { + envelope.Error = &message + } + } + return json.Marshal(envelope) +} + +// hostAbort returns the host's original error when it aborted the call. +func hostAbort(err error) error { + var abort abortError + if errors.As(err, &abort) { + return abort.err + } + return nil +} + // Outcome retains the ordered candidates, prepared requests, and any routing-produced answer. type Outcome struct { SelectedModelIDs []string `json:"selected_model_ids"` @@ -97,7 +141,7 @@ func (router *Router) Close() { } // Route waits for an outcome while the host services model and decision calls. -// Host errors follow the algorithm's configured error policy. +// Host errors follow the algorithm's configured error policy. Aborted calls stop routing. func (router *Router) Route(ctx context.Context, input Input, call HostCall) (Outcome, error) { if router == nil || router.state == nil { return Outcome{}, errors.New("Switchyard router is closed") diff --git a/bindings/go/switchyard_test.go b/bindings/go/switchyard_test.go index c1d91a1422..443a229493 100644 --- a/bindings/go/switchyard_test.go +++ b/bindings/go/switchyard_test.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "io/fs" "runtime" "slices" "strings" @@ -200,6 +201,33 @@ func TestHostFailureUsesAlgorithmPolicy(t *testing.T) { } } +// An aborted host call stops routing even when the algorithm recovers from host errors. +func TestHostAbortStopsRouting(t *testing.T) { + if Abort(nil) != nil { + t.Fatal("nil error should stay nil") + } + errBudget := errors.New("budget exhausted") + refused := &fs.PathError{Op: "route", Path: "judge", Err: errBudget} + original := fmt.Errorf("judge call refused: %w", refused) + for _, algorithm := range []string{composite, decision} { + router := newRouter(t, strings.Replace(algorithm, `"fail_open":false`, `"fail_open":true`, 1)) + calls := 0 + var hostCtx context.Context + _, err := router.Route(context.Background(), input(), func(ctx context.Context, _ Work) (Reply, error) { + calls++ + hostCtx = ctx + return Reply{}, Abort(original) + }) + var typed *fs.PathError + if calls != 1 || err != original || !errors.Is(err, errBudget) || !errors.As(err, &typed) || typed != refused { + t.Fatalf("calls=%d err=%v", calls, err) + } + if hostCtx.Err() != context.Canceled { + t.Fatal("abort did not cancel host context") + } + } +} + // Cancellation releases routing work even if a host ignores its context and replies later. func TestCancellationAndLateHostReply(t *testing.T) { router := newRouter(t, composite) diff --git a/bindings/go/translation_native.go b/bindings/go/translation_native.go index b4a4df02fc..71eb2cf50e 100644 --- a/bindings/go/translation_native.go +++ b/bindings/go/translation_native.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//go:build cgo +//go:build cgo && switchyard_native package switchyard diff --git a/bindings/go/translation_wazero.go b/bindings/go/translation_wazero.go new file mode 100644 index 0000000000..4f884fd77d --- /dev/null +++ b/bindings/go/translation_wazero.go @@ -0,0 +1,12 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !cgo || !switchyard_native + +package switchyard + +import "encoding/json" + +func translate(export string, input []byte) (json.RawMessage, error) { + return wasmTranslate(export, input) +} diff --git a/bindings/go/wasm_test.go b/bindings/go/wasm_test.go index 1c8d6db3af..aad517456f 100644 --- a/bindings/go/wasm_test.go +++ b/bindings/go/wasm_test.go @@ -8,8 +8,12 @@ package switchyard import ( "context" "encoding/json" + "errors" "strings" "testing" + "time" + + "github.com/tetratelabs/wazero/api" ) // Missing startup exports must return errors instead of panicking. @@ -37,39 +41,314 @@ func TestWasmMissingStartupExport(t *testing.T) { } } -// The private runtime uses the same JSON codec ABI as the native backend. -func TestWasmTranslation(t *testing.T) { - request, err := wasmTranslate("sy_decode_request", []byte(`{"format":"openai_chat","body":{"messages":[{"role":"user","content":"hi"}]}}`)) +type cancelAfterNextModule struct { + api.Module + cancel context.CancelFunc +} + +func (module cancelAfterNextModule) ExportedFunction(name string) api.Function { + function := module.Module.ExportedFunction(name) + if name == "sy_run_next" { + return cancelAfterCallFunction{Function: function, cancel: module.cancel} + } + return function +} + +type cancelAfterCallFunction struct { + api.Function + cancel context.CancelFunc +} + +func (function cancelAfterCallFunction) Call(ctx context.Context, args ...uint64) ([]uint64, error) { + values, err := function.Function.Call(ctx, args...) + if err == nil { + function.cancel() + } + return values, err +} + +// Cancellation after the final guest call must take precedence over its outcome. +func TestWasmCancellationBeforeOutcome(t *testing.T) { + backend, err := newWasmBackendWithPoolSize(testConfig(`{"type":"stage_router","picker":"efficient_first",`+stage+`}`), 1) if err != nil { t.Fatal(err) } - var decoded struct { - Messages []struct { - Role string - Content []struct{ Text string } + defer backend.close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + guest := backend.slots[0].gen.guest + guest.module = cancelAfterNextModule{Module: guest.module, cancel: cancel} + if _, err := backend.route(ctx, input(), nil); !errors.Is(err, context.Canceled) { + t.Fatalf("cancelled route returned %v", err) + } +} + +// Parent and child identities must share the guest that owns their session state. +func TestWasmSessionAffinity(t *testing.T) { + backend := &wasmBackend{slots: []*wasmSlot{{}, {}}} + request := input() + request.Metadata = &RequestMetadata{SessionID: "shared"} + parent := backend.slotFor(request) + request.Request = json.RawMessage(`{"messages":[{"role":"user","content":[{"type":"text","text":"child task"}]}]}`) + for _, agent := range []string{"a", "b", ""} { + request.Metadata.IsSubagent = true + request.Metadata.AgentID = agent + if backend.slotFor(request) != parent { + t.Errorf("child %q was assigned outside its parent's session guest", agent) } } - if err := json.Unmarshal(request, &decoded); err != nil || len(decoded.Messages) != 1 || decoded.Messages[0].Role != "user" || len(decoded.Messages[0].Content) != 1 || decoded.Messages[0].Content[0].Text != "hi" { - t.Fatalf("request=%s err=%v", request, err) +} + +func TestWasmTrapIsolatesAndReplacesInstance(t *testing.T) { + config := testConfig(`{"type":"stage_router","picker":"efficient_first",` + stage + `}`) + backend, err := newWasmBackendWithPoolSize(config, 2) + if err != nil { + t.Fatal(err) + } + defer backend.close() + input, err := json.Marshal(input()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + first, err := backend.slots[0].startRun(ctx, input) + if err != nil { + t.Fatal(err) + } + defer first.close() + second, err := backend.slots[0].startRun(ctx, input) + if err != nil { + t.Fatal(err) + } + defer second.close() + other, err := backend.slots[1].startRun(ctx, input) + if err != nil { + t.Fatal(err) + } + defer other.close() + + backend.slots[0].mu.Lock() + _, allocationErr := first.gen.guest.call(ctx, "sy_alloc", uint64(^uint32(0))) + backend.slots[0].mu.Unlock() + if !isWasmTrap(allocationErr) { + t.Fatalf("forced guest allocation panic returned %v", allocationErr) + } + first.ptr = ^uint32(0) + if _, err := first.next(ctx); err == nil || !strings.Contains(err.Error(), "trapped") { + t.Fatalf("forced guest trap returned %v", err) + } + if _, err := second.next(ctx); err == nil || !strings.Contains(err.Error(), "trapped") { + t.Fatalf("run sharing the trapped instance returned %v", err) + } + step, err := other.next(ctx) + if err != nil { + t.Fatalf("unrelated instance failed: %v", err) + } + if !strings.Contains(string(step), `"type":"done"`) { + t.Fatalf("unrelated run returned %s", step) + } + + replacement, err := backend.slots[0].startRun(ctx, input) + if err != nil { + t.Fatalf("pool did not replace the trapped instance: %v", err) + } + defer replacement.close() + step, err = replacement.next(ctx) + if err != nil || !strings.Contains(string(step), `"type":"done"`) { + t.Fatalf("replacement returned step=%s err=%v", step, err) } } -func TestWasmTrapIncludesGuestPanicMessage(t *testing.T) { - runtime, err := newWasmRuntime() +// A failed replacement must not prevent a later request from recovering the slot. +func TestWasmFailedReplacementIsRetried(t *testing.T) { + backend, err := newWasmBackendWithPoolSize(testConfig(`{"type":"stage_router","picker":"efficient_first",`+stage+`}`), 1) if err != nil { t.Fatal(err) } - defer runtime.closeAll() - guest, err := runtime.newModule(context.Background()) + defer backend.close() + request := input() + data, err := json.Marshal(request) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + slot := backend.slots[0] + old, err := slot.startRun(ctx, data) + if err != nil { + t.Fatal(err) + } + defer old.close() + slot.mu.Lock() + config := slot.config + slot.config = []byte(`{`) + slot.poisonLocked(old.gen, errors.New("test trap")) + failed := slot.err != nil + slot.config = config + slot.mu.Unlock() + if !failed { + t.Fatal("replacement with invalid config succeeded") + } + cancelled, cancel := context.WithCancel(ctx) + cancel() + if _, err := slot.startRun(cancelled, data); !errors.Is(err, context.Canceled) || slot.gen != old.gen || slot.err == nil { + t.Fatalf("cancelled request attempted replacement: %v", err) + } + outcome, err := backend.route(ctx, request, nil) + if err != nil || len(outcome.SelectedModelIDs) == 0 || outcome.SelectedModelIDs[0] != "weak-model" { + t.Fatalf("recovered route returned outcome=%+v err=%v", outcome, err) + } + if slot.err != nil || slot.gen == old.gen { + t.Fatal("recovery did not replace the failed generation and clear its error") + } + if _, err := old.next(ctx); err == nil { + t.Fatal("recovery revived a failed run") + } +} + +func TestWasmClosedInstanceIsReplaced(t *testing.T) { + config := testConfig(`{"type":"stage_router","picker":"efficient_first",` + stage + `}`) + backend, err := newWasmBackendWithPoolSize(config, 2) if err != nil { t.Fatal(err) } - _, err = guest.call(context.Background(), "sy_alloc", uint64(^uint32(0))) + defer backend.close() + input, err := json.Marshal(input()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + closedRun, err := backend.slots[0].startRun(ctx, input) + if err != nil { + t.Fatal(err) + } + defer closedRun.close() + otherRun, err := backend.slots[1].startRun(ctx, input) + if err != nil { + t.Fatal(err) + } + defer otherRun.close() + if err := closedRun.gen.guest.module.Close(ctx); err != nil { + t.Fatal(err) + } + if _, err := closedRun.next(ctx); err == nil { + t.Fatal("call on a closed instance succeeded") + } + step, err := otherRun.next(ctx) + if err != nil || !strings.Contains(string(step), `"type":"done"`) { + t.Fatalf("unrelated instance returned step=%s err=%v", step, err) + } + replacement, err := backend.slots[0].startRun(ctx, input) + if err != nil { + t.Fatalf("closed instance was not replaced: %v", err) + } + defer replacement.close() + step, err = replacement.next(ctx) + if err != nil || !strings.Contains(string(step), `"type":"done"`) { + t.Fatalf("replacement returned step=%s err=%v", step, err) + } +} + +// A caller's expired context must not close an instance that other runs share. +func TestWasmCancelledRunKeepsSharedInstance(t *testing.T) { + backend, err := newWasmBackendWithPoolSize(json.RawMessage(`{"algorithm":`+composite+`,"targets":{"weak":"weak-model","strong":"strong-model","judge":"judge-model"}}`), 1) + if err != nil { + t.Fatal(err) + } + defer backend.close() + + entered := make(chan struct{}) + release := make(chan struct{}) + result := make(chan error, 1) + go func() { + _, err := backend.route(context.Background(), input(), func(ctx context.Context, _ Work) (Reply, error) { + close(entered) + select { + case <-release: + return Reply{Response: json.RawMessage(judgeAnswer)}, nil + case <-ctx.Done(): + return Reply{}, ctx.Err() + } + }) + result <- err + }() + <-entered + + large := input() + large.Request = json.RawMessage(`{"messages":[{"role":"user","content":[{"type":"text","text":"` + strings.Repeat("x", 32<<20) + `"}]}]}`) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Millisecond) + defer cancel() + if _, err := backend.route(ctx, large, nil); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expired run returned %v", err) + } + + close(release) + if err := <-result; err != nil { + t.Fatalf("run sharing the instance failed: %v", err) + } +} + +func TestWasmTrapIncludesGuestPanicMessage(t *testing.T) { + backend, err := newWasmBackendWithPoolSize(testConfig(`{"type":"stage_router","picker":"efficient_first",`+stage+`}`), 1) + if err != nil { + t.Fatal(err) + } + defer backend.close() + slot := backend.slots[0] + slot.mu.Lock() + _, err = slot.gen.guest.call(context.Background(), "sy_alloc", uint64(^uint32(0))) + slot.mu.Unlock() if !isWasmTrap(err) || !strings.Contains(err.Error(), "panicked") { t.Fatalf("guest panic returned %v", err) } } +// Pending calls and same-session waiters yield without blocking other runs. +func TestWasmNextWhileHostCallPending(t *testing.T) { + backend, err := newWasmBackendWithPoolSize(json.RawMessage(`{"algorithm":`+composite+`,"targets":{"weak":"weak-model","strong":"strong-model","judge":"judge-model"}}`), 1) + if err != nil { + t.Fatal(err) + } + defer backend.close() + request := input() + request.Metadata = &RequestMetadata{SessionID: "shared"} + data, err := json.Marshal(request) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + run, err := backend.slots[0].startRun(ctx, data) + if err != nil { + t.Fatal(err) + } + defer run.close() + step, err := run.next(ctx) + if err != nil || !strings.Contains(string(step), `"type":"call_model"`) { + t.Fatalf("first step=%s err=%v", step, err) + } + if step, err := run.next(ctx); err != nil || step != nil { + t.Fatalf("read with an unanswered call returned step=%s err=%v", step, err) + } + waiting, err := backend.slots[0].startRun(ctx, data) + if err != nil { + t.Fatal(err) + } + if step, err := waiting.next(ctx); err != nil || step != nil { + waiting.close() + t.Fatalf("same-session waiter returned step=%s err=%v", step, err) + } + waiting.close() + var work Work + if err := json.Unmarshal(step, &work); err != nil { + t.Fatal(err) + } + if err := run.respond(ctx, work.ID, []byte(`{"response":`+judgeAnswer+`}`)); err != nil { + t.Fatal(err) + } + if step, err := run.next(ctx); err != nil || !strings.Contains(string(step), `"type":"done"`) { + t.Fatalf("surviving run returned step=%s err=%v", step, err) + } +} + func TestWasmTranslatorIsReplacedAfterFailure(t *testing.T) { input := []byte(`{"format":"openai_chat","body":{"messages":[{"role":"user","content":"hi"}]}}`) if _, err := wasmTranslate("sy_decode_request", input); err != nil { diff --git a/crates/switchyard-ffi/include/switchyard.h b/crates/switchyard-ffi/include/switchyard.h index 5c53473c9b..fb3c333860 100644 --- a/crates/switchyard-ffi/include/switchyard.h +++ b/crates/switchyard-ffi/include/switchyard.h @@ -35,6 +35,8 @@ char *sy_run_new(const SyRouter *router, const char *input_json, SyRun **out); // Callbacks must return promptly, not unwind, and not call sy_run_respond or free handles. char *sy_run_next_async(const SyRun *run, SyNextCallback callback, uintptr_t user_data); #endif +// reply_json is {"response": ...}, {"error": "..."} for the algorithm's error policy, +// or {"abort": "..."} to fail the run. char *sy_run_respond(const SyRun *run, uint64_t call_id, const char *reply_json); char *sy_decode_request(const char *input_json, char **out_json); diff --git a/crates/switchyard-ffi/src/lib.rs b/crates/switchyard-ffi/src/lib.rs index 72e238df42..30eedf79f7 100644 --- a/crates/switchyard-ffi/src/lib.rs +++ b/crates/switchyard-ffi/src/lib.rs @@ -97,6 +97,9 @@ enum CallReply { Error { error: String, }, + Abort { + abort: String, + }, Response { response: serde_json::Value, served_model: Option, @@ -629,6 +632,7 @@ pub unsafe extern "C" fn sy_run_next(run: *const SyRun, out_json: *mut *mut c_ch } /// Delivers a buffered host reply to its pending call. +/// An `abort` reply fails the run, even when the algorithm recovers from `error` replies. /// # Safety /// The run must remain live and JSON must be NUL-terminated. May overlap an async read. /// Do not call from a completion callback or Tokio task. @@ -647,6 +651,7 @@ pub unsafe extern "C" fn sy_run_respond( .pending .remove(&call_id) .ok_or("unknown call ID")?; + let is_abort = matches!(reply, CallReply::Abort { .. }); let result = match call { Call::Model(call) => { let response = match reply { @@ -654,6 +659,10 @@ pub unsafe extern "C" fn sy_run_respond( "host model call", std::io::Error::other(error), ))), + CallReply::Abort { abort } => Ok(Err(LibsyError::external( + "host routing", + std::io::Error::other(abort), + ))), CallReply::Response { response, served_model, @@ -675,9 +684,13 @@ pub unsafe extern "C" fn sy_run_respond( return Err(error.into()); } }; - run.engine - .runtime - .block_on(call.respond(std::future::ready(response))) + match response { + Err(error) if is_abort => call.fail(error), + response => run + .engine + .runtime + .block_on(call.respond(std::future::ready(response))), + } } Call::Decision(call) => { let response = match reply { @@ -685,6 +698,10 @@ pub unsafe extern "C" fn sy_run_respond( "host decision call", std::io::Error::other(error), ))), + CallReply::Abort { abort } => Ok(Err(LibsyError::external( + "host routing", + std::io::Error::other(abort), + ))), CallReply::Response { response, .. } => { serde_json::from_value(response).map(Ok) } @@ -696,9 +713,13 @@ pub unsafe extern "C" fn sy_run_respond( return Err(error.into()); } }; - run.engine - .runtime - .block_on(call.respond(std::future::ready(response))) + match response { + Err(error) if is_abort => call.fail(error), + response => run + .engine + .runtime + .block_on(call.respond(std::future::ready(response))), + } } }; if result.is_err() { diff --git a/crates/switchyard-ffi/src/tests.rs b/crates/switchyard-ffi/src/tests.rs index f59e987408..65ceed9258 100644 --- a/crates/switchyard-ffi/src/tests.rs +++ b/crates/switchyard-ffi/src/tests.rs @@ -342,6 +342,79 @@ fn invalid_decision_replies_preserve_pending_calls() { assert_eq!(done["selected_model_ids"], json!(["capable", "efficient"])); } +fn fail_open_judge_router(decision: bool) -> Box { + let mut algorithm = json!({ + "type": "llm_classifier", "classifier_target": "judge", + "strong_target": "capable", "weak_target": "efficient", "fail_open": true, + }); + if decision { + algorithm["decision"] = json!({ + "cutoff": 0.5, "candidates": {"a": "capable", "b": "efficient"}, "evidence": {}, + }); + } else { + algorithm["base_threshold"] = json!(0.5); + } + router(algorithm) +} + +fn judge_input() -> Value { + json!({ + "request": {}, + "models": { + "any": ["capable", "efficient"], "capable": ["capable"], + "efficient": ["efficient"], "judge": ["judge"], + }, + }) +} + +// An abort reply stops the run even when the algorithm recovers from an error reply. +#[test] +fn abort_replies_fail_runs_that_recover_from_errors() { + for (decision, kind) in [(false, "call_model"), (true, "call_decision")] { + let router = fail_open_judge_router(decision); + let recovered = run(&router, judge_input()); + let call = next(&recovered); + assert_eq!(call["type"], kind); + let id = call["id"].as_u64().unwrap(); + check(unsafe { + sy_run_respond( + &*recovered, + id, + c"{\"error\":\"judge unavailable\"}".as_ptr(), + ) + }); + assert_eq!(next(&recovered)["type"], "done", "{kind}"); + + let aborted = run(&router, judge_input()); + let id = next(&aborted)["id"].as_u64().unwrap(); + let (sender, receiver) = mpsc::channel(); + next_async(&aborted, &sender).unwrap(); + let error = + unsafe { sy_run_respond(&*aborted, id, c"{\"abort\":\"budget exhausted\"}".as_ptr()) }; + assert_eq!( + error_message(error), + "host routing failed: budget exhausted", + "{kind}" + ); + assert_eq!( + receiver + .recv_timeout(Duration::from_secs(5)) + .unwrap() + .unwrap_err() + .to_str() + .unwrap(), + "routing failed" + ); + assert_eq!( + error_message(unsafe { + sy_run_respond(&*aborted, id, c"{\"abort\":\"budget exhausted\"}".as_ptr()) + }), + "routing failed" + ); + assert_eq!(next_async(&aborted, &sender).unwrap_err(), "routing failed"); + } +} + // Translation exports return owned JSON strings that the caller releases through the ABI. #[test] fn translation_exports_return_owned_json() {