Make aggregate ONNX gates execute for real

Reference-count the process-global ONNX Runtime across embedder and routing-head sessions, make close idempotent, and require named proof that both aggregate routing gates executed rather than self-skipped (V-716). Owner explicitly requested direct commits to master.

internal/router/eval/eval_test.go

@@ -197,6 +197,7 @@ func TestONNXBaseline(t *testing.T) {
 		t.Skipf("onnx embedder unavailable: %v", err)
 	}
 	defer emb.Close()
+	recordONNXGateExecuted(t)
 
 	f, err := Load()
 	if err != nil {

internal/router/eval/heads_test.go

@@ -54,6 +54,7 @@ func TestONNXRoutingHeads(t *testing.T) {
 		t.Skipf("routing heads unavailable: %v", err)
 	}
 	defer h.Close()
+	recordONNXGateExecuted(t)
 
 	f, err := Load()
 	if err != nil {

internal/router/eval/onnx_testmain_test.go

@@ -0,0 +1,99 @@
+package eval
+
+import (
+	"fmt"
+	"os"
+	"sort"
+	"strconv"
+	"sync"
+	"testing"
+
+	"github.com/kami/maven/internal/router"
+)
+
+const requiredONNXGatesEnv = "MAVEN_ONNX_REQUIRED_GATES"
+
+var executedONNXGates struct {
+	sync.Mutex
+	names []string
+}
+
+// recordONNXGateExecuted is called only after a gate has loaded every model it
+// needs. TestMain uses the names as proof that an aggregate command did not
+// turn a second initialization error into a green SKIP.
+func recordONNXGateExecuted(t *testing.T) {
+	t.Helper()
+	executedONNXGates.Lock()
+	executedONNXGates.names = append(executedONNXGates.names, t.Name())
+	executedONNXGates.Unlock()
+}
+
+func TestMain(m *testing.M) {
+	required, err := requiredONNXGateCount()
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(2)
+	}
+
+	lease, err := configuredONNXTestRuntime(required > 0)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(2)
+	}
+
+	code := m.Run()
+	if required > 0 {
+		executedONNXGates.Lock()
+		names := append([]string(nil), executedONNXGates.names...)
+		executedONNXGates.Unlock()
+		sort.Strings(names)
+		fmt.Fprintf(os.Stderr, "ONNX aggregate proof: verified %d required model gates: %v\n", len(names), names)
+		if len(names) < required {
+			fmt.Fprintf(os.Stderr, "ONNX aggregate gate failed: executed %d, require at least %d\n", len(names), required)
+			code = 1
+		}
+	}
+
+	if lease != nil {
+		if err := lease.Close(); err != nil {
+			fmt.Fprintf(os.Stderr, "ONNX test runtime cleanup: %v\n", err)
+			code = 1
+		}
+	}
+	os.Exit(code)
+}
+
+func requiredONNXGateCount() (int, error) {
+	raw := os.Getenv(requiredONNXGatesEnv)
+	if raw == "" {
+		return 0, nil
+	}
+	n, err := strconv.Atoi(raw)
+	if err != nil || n < 0 {
+		return 0, fmt.Errorf("%s must be a non-negative integer, got %q", requiredONNXGatesEnv, raw)
+	}
+	return n, nil
+}
+
+func configuredONNXTestRuntime(required bool) (*router.ONNXRuntimeLease, error) {
+	lib := os.Getenv("MAVEN_ONNX_LIB")
+	if lib == "" {
+		if required {
+			return nil, fmt.Errorf("ONNX aggregate gate requires MAVEN_ONNX_LIB")
+		}
+		return nil, nil
+	}
+	if _, err := os.Stat(lib); err != nil {
+		if required {
+			return nil, fmt.Errorf("ONNX aggregate gate runtime %s: %w", lib, err)
+		}
+		// Preserve the ordinary portable test path: individual model tests
+		// report a skip when optional local model dependencies are absent.
+		return nil, nil
+	}
+	lease, err := router.AcquireONNXRuntime(lib)
+	if err != nil {
+		return nil, fmt.Errorf("initialize shared ONNX test runtime: %w", err)
+	}
+	return lease, nil
+}

internal/router/heads.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"os"
 	"path/filepath"
+	"sync"
 
 	ort "github.com/yalue/onnxruntime_go"
 )
@@ -54,9 +55,12 @@ const (
 type RouterHeads struct {
 	tokenizer *unigramTokenizer
 	session   *ort.DynamicSession[int64, float32]
+	runtime   *ONNXRuntimeLease
 	intents   []Intent
 	sources   []Source
 	threshold float64
+	closeOnce sync.Once
+	closeErr  error
 }
 
 // headsMeta — router_heads.json, written beside the weights by the exporter.
@@ -70,8 +74,9 @@ type headsMeta struct {
 // NewRouterHeads loads the graph and its label order. modelPath points at the
 // .onnx; the external weights and router_heads.json sit beside it.
 //
-// It assumes the ONNX environment is already initialised, because the embedder
-// does that at startup and the runtime allows it once.
+// It shares the process-global ONNX environment with the embedder. The runtime
+// lease is independent so shutdown order cannot unload the library while this
+// graph's session is still alive.
 func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
 	metaPath := filepath.Join(filepath.Dir(modelPath), "router_heads.json")
 	raw, err := os.ReadFile(metaPath)
@@ -106,18 +111,24 @@ func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
 	if err != nil {
 		return nil, fmt.Errorf("heads: tokenizer: %w", err)
 	}
+	runtime, err := AcquireONNXRuntime("")
+	if err != nil {
+		return nil, fmt.Errorf("heads: %w", err)
+	}
 	session, err := ort.NewDynamicSession[int64, float32](
 		modelPath,
 		[]string{"input_ids", "attention_mask"},
 		[]string{"intent", "source", "slots", "clarify"},
 	)
 	if err != nil {
+		_ = runtime.Close()
 		return nil, fmt.Errorf("heads: create session: %w", err)
 	}
 
 	return &RouterHeads{
 		tokenizer: tok,
 		session:   session,
+		runtime:   runtime,
 		intents:   intents,
 		sources:   sources,
 		threshold: headsThreshold,
@@ -128,8 +139,14 @@ func (h *RouterHeads) Close() error {
 	if h == nil {
 		return nil
 	}
-	h.session.Destroy()
-	return nil
+	h.closeOnce.Do(func() {
+		if h.session != nil {
+			h.session.Destroy()
+			h.session = nil
+		}
+		h.closeErr = h.runtime.Close()
+	})
+	return h.closeErr
 }
 
 // headsResult — one forward pass, read back.

internal/router/onnxembedder.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"os"
 	"strings"
+	"sync"
 
 	ort "github.com/yalue/onnxruntime_go"
 	"golang.org/x/text/unicode/norm"
@@ -33,17 +34,21 @@ const (
 type onnxEmbedder struct {
 	tokenizer *unigramTokenizer
 	session   *ort.DynamicSession[int64, float32]
+	runtime   *ONNXRuntimeLease
 	id        string
+	closeOnce sync.Once
+	closeErr  error
 }
 
 func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
-	ort.SetSharedLibraryPath(libPath)
-	if err := ort.InitializeEnvironment(); err != nil {
-		return nil, fmt.Errorf("onnx: init environment: %w", err)
+	runtime, err := AcquireONNXRuntime(libPath)
+	if err != nil {
+		return nil, err
 	}
 
 	tok, err := newUnigramTokenizer(tokenizerPath)
 	if err != nil {
+		_ = runtime.Close()
 		return nil, fmt.Errorf("tokenizer: %w", err)
 	}
 
@@ -53,12 +58,14 @@ func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, e
 		[]string{"last_hidden_state"},
 	)
 	if err != nil {
+		_ = runtime.Close()
 		return nil, fmt.Errorf("onnx: create session: %w", err)
 	}
 
 	return &onnxEmbedder{
 		tokenizer: tok,
 		session:   session,
+		runtime:   runtime,
 		id:        modelIDFromPath(modelPath),
 	}, nil
 }
@@ -148,8 +155,17 @@ func (e *onnxEmbedder) embed(ctx context.Context, text string) ([]float32, error
 }
 
 func (e *onnxEmbedder) Close() error {
-	e.session.Destroy()
-	return nil
+	if e == nil {
+		return nil
+	}
+	e.closeOnce.Do(func() {
+		if e.session != nil {
+			e.session.Destroy()
+			e.session = nil
+		}
+		e.closeErr = e.runtime.Close()
+	})
+	return e.closeErr
 }
 
 func meanPool(hidden []float32, mask []int64, seqLen, dim int) []float32 {

internal/router/onnxruntime.go

@@ -0,0 +1,124 @@
+package router
+
+import (
+	"fmt"
+	"path/filepath"
+	"sync"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+// ONNXRuntimeLease keeps the process-global ONNX Runtime environment alive.
+//
+// onnxruntime_go exposes the environment as package-global state: it may be
+// initialized only once, and destroying it while any session is alive is
+// unsafe. A lease makes that ownership explicit. Model wrappers hold one for
+// the lifetime of their session; TestMain may hold an additional lease so a
+// package's model-aware gates share one environment instead of repeatedly
+// loading and unloading the shared library.
+type ONNXRuntimeLease struct {
+	once sync.Once
+	err  error
+}
+
+var sharedONNXRuntime = struct {
+	sync.Mutex
+	refs    int
+	libPath string
+	managed bool
+}{}
+
+// AcquireONNXRuntime acquires the process-global ONNX Runtime environment.
+// The first caller initializes it; later callers share it. A non-empty library
+// path must agree with the path that initialized the active environment.
+//
+// The external-environment branch is intentional. It keeps this package safe
+// when an executable initialized onnxruntime_go directly before constructing
+// a router model; in that case the lease never destroys state it does not own.
+func AcquireONNXRuntime(libPath string) (*ONNXRuntimeLease, error) {
+	libPath = cleanONNXLibraryPath(libPath)
+
+	sharedONNXRuntime.Lock()
+	defer sharedONNXRuntime.Unlock()
+
+	if sharedONNXRuntime.refs > 0 {
+		if err := compatibleONNXLibraryPath(sharedONNXRuntime.libPath, libPath); err != nil {
+			return nil, err
+		}
+		sharedONNXRuntime.refs++
+		return &ONNXRuntimeLease{}, nil
+	}
+
+	if ort.IsInitialized() {
+		// Some other component owns the already-live environment. Attach to it,
+		// but never unload its shared library when our last lease closes.
+		sharedONNXRuntime.refs = 1
+		sharedONNXRuntime.libPath = libPath
+		sharedONNXRuntime.managed = false
+		return &ONNXRuntimeLease{}, nil
+	}
+
+	if libPath != "" {
+		ort.SetSharedLibraryPath(libPath)
+	}
+	if err := ort.InitializeEnvironment(); err != nil {
+		return nil, fmt.Errorf("onnx: init environment: %w", err)
+	}
+	sharedONNXRuntime.refs = 1
+	sharedONNXRuntime.libPath = libPath
+	sharedONNXRuntime.managed = true
+	return &ONNXRuntimeLease{}, nil
+}
+
+// Close releases one environment lease. The last lease destroys an
+// environment initialized by this package, after every model session that
+// held a lease has already been destroyed.
+func (l *ONNXRuntimeLease) Close() error {
+	if l == nil {
+		return nil
+	}
+	l.once.Do(func() {
+		l.err = releaseONNXRuntime()
+	})
+	return l.err
+}
+
+func releaseONNXRuntime() error {
+	sharedONNXRuntime.Lock()
+	defer sharedONNXRuntime.Unlock()
+
+	if sharedONNXRuntime.refs == 0 {
+		return fmt.Errorf("onnx: release environment without a lease")
+	}
+	sharedONNXRuntime.refs--
+	if sharedONNXRuntime.refs != 0 {
+		return nil
+	}
+
+	managed := sharedONNXRuntime.managed
+	sharedONNXRuntime.libPath = ""
+	sharedONNXRuntime.managed = false
+	if !managed {
+		return nil
+	}
+	if err := ort.DestroyEnvironment(); err != nil {
+		return fmt.Errorf("onnx: destroy environment: %w", err)
+	}
+	return nil
+}
+
+func cleanONNXLibraryPath(path string) string {
+	if path == "" {
+		return ""
+	}
+	return filepath.Clean(path)
+}
+
+func compatibleONNXLibraryPath(active, requested string) error {
+	// Empty means "use the environment that is already active". RouterHeads
+	// intentionally has no second copy of the embedder's library-path config.
+	if active == "" || requested == "" || active == requested {
+		return nil
+	}
+	return fmt.Errorf("onnx: runtime already uses shared library %q, cannot also use %q", active, requested)
+}

internal/router/onnxruntime_test.go

@@ -0,0 +1,39 @@
+package router
+
+import (
+	"os"
+	"testing"
+)
+
+// TestONNXRuntimeCanBeReacquired pins the process-global lifecycle itself.
+// The model gates exercise live sessions; this smaller check proves that the
+// last lease unloads a managed environment and that a later gate can acquire
+// it again. It is opt-in for the same reason as every model measurement.
+func TestONNXRuntimeCanBeReacquired(t *testing.T) {
+	lib := os.Getenv("MAVEN_ONNX_LIB")
+	if lib == "" {
+		t.Skip("MAVEN_ONNX_LIB unset — see AGENTS.md § Embedder model for intent routing")
+	}
+	if _, err := os.Stat(lib); err != nil {
+		t.Skipf("missing %s: %v", lib, err)
+	}
+
+	first, err := AcquireONNXRuntime(lib)
+	if err != nil {
+		t.Fatalf("first acquire: %v", err)
+	}
+	if err := first.Close(); err != nil {
+		t.Fatalf("first close: %v", err)
+	}
+	if err := first.Close(); err != nil {
+		t.Fatalf("idempotent close: %v", err)
+	}
+
+	second, err := AcquireONNXRuntime(lib)
+	if err != nil {
+		t.Fatalf("reacquire after last close: %v", err)
+	}
+	if err := second.Close(); err != nil {
+		t.Fatalf("second close: %v", err)
+	}
+}