iop/apps/edge/internal/service/run_dispatch_internal_test.go

package service

import (
	"context"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func TestRunHandleCloseIsIdempotent(t *testing.T) {
	calls := 0
	handle := &RunHandle{
		close: func() {
			calls++
		},
	}

	handle.Close()
	handle.Close()

	if calls != 1 {
		t.Fatalf("close called %d times, want 1", calls)
	}
}

func waitForCondition(t *testing.T, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal(msg)
}

func inflightRunCount(q *modelQueueManager) int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return len(q.leaseByRun)
}

// TestRouteProviderTunnelFrameDoesNotPublishToBus verifies that tunnel frames
// reach only the request-bound tunnel subscriber and never the run event bus
// fanout (negative test for the passthrough/no-sideband boundary).
func TestRouteProviderTunnelFrameDoesNotPublishToBus(t *testing.T) {
	bus := edgeevents.NewBus()
	svc := New(edgenode.NewRegistry(), bus)

	runCh, unsubRun := bus.SubscribeAllRuns(8)
	defer unsubRun()

	tunnelCh, unsubTunnel := svc.tunnels.subscribe("tunnel-neg-1", 8)
	defer unsubTunnel()

	frame := &iop.ProviderTunnelFrame{
		RunId:    "run-neg-1",
		TunnelId: "tunnel-neg-1",
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
		Body:     []byte("data: raw provider bytes\n\n"),
	}
	svc.RouteProviderTunnelFrame(frame)

	select {
	case got := <-tunnelCh:
		if string(got.GetBody()) != string(frame.GetBody()) {
			t.Fatalf("tunnel frame body mismatch: %q", got.GetBody())
		}
	case <-time.After(time.Second):
		t.Fatal("tunnel subscriber did not receive the frame")
	}

	select {
	case e := <-runCh:
		t.Fatalf("tunnel frame leaked into run event bus: %+v", e)
	case <-time.After(100 * time.Millisecond):
	}
}

// TestRouteProviderTunnelFrameWithoutSubscriberIsDropped verifies routing a
// frame with no subscriber neither blocks nor panics.
func TestRouteProviderTunnelFrameWithoutSubscriberIsDropped(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{TunnelId: "no-sub"})
}

type providerTunnelTestEnv struct {
	svc        *Service
	nodeClient *toki.TcpClient
	capturedMu sync.Mutex
	captured   *iop.ProviderTunnelRequest
}

func (e *providerTunnelTestEnv) capturedRequest() *iop.ProviderTunnelRequest {
	e.capturedMu.Lock()
	defer e.capturedMu.Unlock()
	return e.captured
}

// newProviderTunnelTestEnv wires a provider-pool service to a fake node over
// net.Pipe and captures the ProviderTunnelRequest the node receives.
func newProviderTunnelTestEnv(t *testing.T) *providerTunnelTestEnv {
	t.Helper()
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	env := &providerTunnelTestEnv{nodeClient: nodeClient}
	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
		env.capturedMu.Lock()
		env.captured = req
		env.capturedMu.Unlock()
	})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pool",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm-01",
				Adapter:  "vllm-gpu",
				Type:     "vllm",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 1,
			},
		},
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-pool",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
	})

	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-vllm-01": "served-qwen"}},
	})
	env.svc = svc
	return env
}

// TestSubmitProviderTunnelProviderPoolSendsRequestAndReleasesSlotOnEnd covers
// the provider-pool tunnel path end to end: admission rewrites adapter/target,
// BuildBody receives the served target, ordered frames reach the handle, and
// the admission slot is released on the END frame.
func TestSubmitProviderTunnelProviderPoolSendsRequestAndReleasesSlotOnEnd(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	var builtTarget string
	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-001",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		BuildBody: func(target string) ([]byte, error) {
			builtTarget = target
			return []byte(`{"model":"` + target + `","stream":true}`), nil
		},
		Stream:       true,
		ProviderPool: true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	defer handle.Close()

	if builtTarget != "served-qwen" {
		t.Errorf("BuildBody target: got %q, want served-qwen", builtTarget)
	}
	dispatch := handle.Dispatch()
	if dispatch.Adapter != "vllm-gpu" || dispatch.Target != "served-qwen" {
		t.Errorf("dispatch adapter/target: got %q/%q", dispatch.Adapter, dispatch.Target)
	}

	waitForCondition(t, func() bool { return env.capturedRequest() != nil },
		"fake node did not receive ProviderTunnelRequest")
	captured := env.capturedRequest()
	if captured.GetAdapter() != "vllm-gpu" || captured.GetTarget() != "served-qwen" {
		t.Errorf("wire adapter/target: got %q/%q", captured.GetAdapter(), captured.GetTarget())
	}
	if !strings.Contains(string(captured.GetBody()), `"model":"served-qwen"`) {
		t.Errorf("wire body missing served model: %s", captured.GetBody())
	}
	if inflightRunCount(svc.queue) != 1 {
		t.Fatalf("expected 1 inflight run after dispatch, got %d", inflightRunCount(svc.queue))
	}

	tunnelID := captured.GetTunnelId()
	frames := []*iop.ProviderTunnelFrame{
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 0, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200},
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 1, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte("data: chunk-1\n\n")},
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 2, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true},
	}
	for _, f := range frames {
		svc.RouteProviderTunnelFrame(f)
	}

	stream := handle.Stream()
	for i, want := range frames {
		select {
		case got := <-stream.Frames:
			if got.GetSequence() != want.GetSequence() || got.GetKind() != want.GetKind() {
				t.Fatalf("frame %d: got seq=%d kind=%v", i, got.GetSequence(), got.GetKind())
			}
		case <-time.After(time.Second):
			t.Fatalf("frame %d not delivered", i)
		}
	}
	select {
	case _, open := <-stream.Frames:
		if open {
			t.Fatal("expected frame channel to close after END")
		}
	case <-time.After(time.Second):
		t.Fatal("frame channel not closed after END")
	}

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after END frame")
}

// TestSubmitProviderTunnelCloseReleasesSlot verifies the cancel path: closing
// the handle before a terminal frame releases the provider-pool slot.
func TestSubmitProviderTunnelCloseReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-cancel",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		Body:          []byte(`{"model":"qwen3.6:35b"}`),
		ProviderPool:  true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	if inflightRunCount(svc.queue) != 1 {
		t.Fatalf("expected 1 inflight run after dispatch, got %d", inflightRunCount(svc.queue))
	}

	handle.Close()

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after Close")
}

// TestProviderTunnelRouterBackpressureBlocksWithoutDrop verifies the router
// blocks the producer while the subscriber buffer is full instead of dropping
// frames, and delivers every frame in order once the subscriber drains.
func TestProviderTunnelRouterBackpressureBlocksWithoutDrop(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	ch, unsub := svc.tunnels.subscribe("tunnel-bp", 2)
	defer unsub()

	const total = 6
	routed := make(chan struct{})
	go func() {
		defer close(routed)
		for i := 0; i < total; i++ {
			svc.tunnels.route(&iop.ProviderTunnelFrame{
				TunnelId: "tunnel-bp",
				Sequence: int64(i),
				Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
			})
		}
	}()

	// With buffer 2 the producer cannot finish all routes before the
	// subscriber drains: full buffer must block, not drop.
	select {
	case <-routed:
		t.Fatal("route calls finished while the subscriber buffer was full; expected blocking backpressure")
	case <-time.After(100 * time.Millisecond):
	}

	for i := 0; i < total; i++ {
		select {
		case f := <-ch:
			if f.GetSequence() != int64(i) {
				t.Fatalf("frame %d: got sequence %d; frames must stay ordered and lossless under backpressure", i, f.GetSequence())
			}
		case <-time.After(time.Second):
			t.Fatalf("frame %d not delivered after draining", i)
		}
	}

	select {
	case <-routed:
	case <-time.After(time.Second):
		t.Fatal("producer still blocked after all frames were drained")
	}
}

// TestProviderTunnelRouterCloseUnblocksBlockedRoute verifies unsubscribe (the
// handle Close/cancel path) releases a route call blocked on a full buffer, so
// an abandoned tunnel can never wedge the transport frame reader.
func TestProviderTunnelRouterCloseUnblocksBlockedRoute(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	_, unsub := svc.tunnels.subscribe("tunnel-close", 1)

	frame := func(seq int64) *iop.ProviderTunnelFrame {
		return &iop.ProviderTunnelFrame{
			TunnelId: "tunnel-close",
			Sequence: seq,
			Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
		}
	}

	blocked := make(chan struct{})
	go func() {
		defer close(blocked)
		svc.tunnels.route(frame(0)) // fills the buffer
		svc.tunnels.route(frame(1)) // blocks until close releases it
	}()

	select {
	case <-blocked:
		t.Fatal("second route must block on the full subscriber buffer")
	case <-time.After(100 * time.Millisecond):
	}

	unsub()

	select {
	case <-blocked:
	case <-time.After(time.Second):
		t.Fatal("unsubscribe did not release the blocked route call")
	}

	// A late frame for the closed tunnel is dropped without blocking.
	svc.tunnels.route(frame(2))
}

// TestProviderTunnelReleaseExactlyOnceAcrossEndAndClose verifies the tracked
// provider-pool slot (capacity 1) is released exactly once when both the
// terminal END frame and handle Close fire, and that the freed slot admits the
// next tunnel immediately instead of queueing.
func TestProviderTunnelReleaseExactlyOnceAcrossEndAndClose(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	submit := func(runID string) ProviderTunnelResult {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		handle, err := svc.SubmitProviderTunnel(ctx, SubmitProviderTunnelRequest{
			RunID:         runID,
			ModelGroupKey: "qwen3.6:35b",
			Method:        "POST",
			Path:          "/v1/chat/completions",
			Body:          []byte(`{"model":"qwen3.6:35b"}`),
			ProviderPool:  true,
		})
		if err != nil {
			t.Fatalf("SubmitProviderTunnel(%s): %v", runID, err)
		}
		return handle
	}

	handle1 := submit("run-release-1")
	waitForCondition(t, func() bool {
		captured := env.capturedRequest()
		return captured != nil && captured.GetRunId() == "run-release-1"
	}, "fake node did not receive the first ProviderTunnelRequest")
	captured := env.capturedRequest()

	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{
		RunId:    captured.GetRunId(),
		TunnelId: captured.GetTunnelId(),
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
		End:      true,
	})
	stream := handle1.Stream()
	for {
		if _, open := <-stream.Frames; !open {
			break
		}
	}
	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after END frame")
	if got := leaseCount(svc.queue); got != 0 {
		t.Fatalf("live leases after END frame: got %d, want 0", got)
	}

	// Close after the terminal frame must stay a no-op release: the lease it
	// would target is already gone, so there is nothing left to decrement.
	handle1.Close()
	if got := inflightRunCount(svc.queue); got != 0 {
		t.Fatalf("inflight count after END+Close: got %d, want 0", got)
	}
	if got := leaseCount(svc.queue); got != 0 {
		t.Fatalf("live leases after END+Close: got %d, want 0", got)
	}

	// The capacity-1 slot is free again: the next tunnel is admitted without
	// waiting on the queue (a leaked slot would block this submit).
	handle2 := submit("run-release-2")
	defer handle2.Close()
	waitForCondition(t, func() bool {
		captured := env.capturedRequest()
		return captured != nil && captured.GetRunId() == "run-release-2"
	}, "fake node did not receive the second ProviderTunnelRequest after slot release")
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("inflight count after second dispatch: got %d, want 1", got)
	}
	// The second tunnel holds its own lease identity, so a late release carrying
	// the first tunnel's id can never free it.
	if got := leaseCount(svc.queue); got != 1 {
		t.Fatalf("live leases after second dispatch: got %d, want 1", got)
	}
}

// TestSubmitProviderTunnelErrorFrameReleasesSlot verifies the ERROR frame is
// terminal and releases the provider-pool slot like END.
func TestSubmitProviderTunnelErrorFrameReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-error",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		Body:          []byte(`{"model":"qwen3.6:35b"}`),
		ProviderPool:  true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	defer handle.Close()

	waitForCondition(t, func() bool { return env.capturedRequest() != nil },
		"fake node did not receive ProviderTunnelRequest")
	captured := env.capturedRequest()

	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{
		RunId:    captured.GetRunId(),
		TunnelId: captured.GetTunnelId(),
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
		Error:    "provider unavailable",
	})

	stream := handle.Stream()
	select {
	case got := <-stream.Frames:
		if got.GetKind() != iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR {
			t.Fatalf("expected ERROR frame, got %v", got.GetKind())
		}
	case <-time.After(time.Second):
		t.Fatal("ERROR frame not delivered")
	}

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after ERROR frame")
}

// TestSubmitProviderPoolSelectsTunnelProviderAndReleasesSlot verifies that
// SubmitProviderPool with an OpenAI-compatible provider (tunnel path) performs
// a single admission, sends a ProviderTunnelRequest, and releases the slot on
// tunnel END frame.
func TestSubmitProviderPoolSelectsTunnelProviderAndReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	// Catalog entry with a vLLM provider (tunnel path).
	catalog := []config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-vllm-01": "served-qwen"}},
	}
	svc.SetModelCatalog(catalog)

	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "qwen3.6:35b",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}
	if result.Path != ProviderPoolPathTunnel {
		t.Fatalf("expected path=tunnel, got %q", result.Path)
	}
	if result.Tunnel == nil {
		t.Fatal("expected tunnel handle for tunnel path")
	}
	if result.Run != nil {
		t.Fatal("expected nil run for tunnel path")
	}

	// Verify inflight after dispatch.
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight after dispatch, got %d", got)
	}
}

// TestSubmitProviderPoolSelectsNormalizedProviderAndDoesNotOpenTunnel verifies
// that SubmitProviderPool with an Ollama provider (normalized path) performs
// a single admission, sends a RunRequest (not ProviderTunnelRequest), and
// the tunnel handle is nil.
func TestSubmitProviderPoolSelectsNormalizedProviderAndDoesNotOpenTunnel(t *testing.T) {
	// Build a provider-pool service with an Ollama provider.
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	// Capture the RunRequest received by the fake node.
	var capturedRunReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedRunReq = req
		capturedMu.Unlock()
	})

	// Build NodeStore with an Ollama provider (normalized path).
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ollama",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-ollama-01",
				Adapter:  "ollama",
				Type:     "ollama",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Build catalog with Ollama provider.
	catalog := []config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-ollama-01": "served-qwen"}},
	}

	// Build registry with the fake node.
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-ollama",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
	})

	// Create Service with queue and catalog.
	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// SubmitProviderPool with Ollama provider → normalized path.
	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "qwen3.6:35b",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}
	if result.Path != ProviderPoolPathNormalized {
		t.Fatalf("expected path=normalized, got %q", result.Path)
	}
	if result.Run == nil {
		t.Fatal("expected run handle for normalized path")
	}
	if result.Tunnel != nil {
		t.Fatal("expected nil tunnel for normalized path")
	}

	// Wait for the fake node to receive the RunRequest.
	time.Sleep(50 * time.Millisecond)

	capturedMu.Lock()
	defer capturedMu.Unlock()

	if capturedRunReq == nil {
		t.Fatal("no RunRequest captured from fake node; SubmitProviderPool did not send")
	}

	// Verify that the adapter and target were rewritten from the provider-pool candidate.
	if capturedRunReq.GetAdapter() != "ollama" {
		t.Errorf("adapter: got %q, want %q", capturedRunReq.GetAdapter(), "ollama")
	}
	if capturedRunReq.GetTarget() != "served-qwen" {
		t.Errorf("target: got %q, want %q", capturedRunReq.GetTarget(), "served-qwen")
	}

	// Verify inflight after dispatch.
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight after dispatch, got %d", got)
	}
}

// TestSubmitProviderPoolUsesSingleAdmissionForMixedCandidates verifies that
// SubmitProviderPool performs exactly one queue admission when presented with
// mixed tunnel+normalized candidates (e.g., vLLM + Ollama). The selected
// candidate (by in-flight/priority rotation) determines the path.
func TestSubmitProviderPoolUsesSingleAdmissionForMixedCandidates(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	var capturedRunReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedRunReq = req
		capturedMu.Unlock()
	})

	// Build NodeStore with two providers: vLLM (tunnel) and Ollama (normalized).
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-mixed",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm",
				Adapter:  "vllm-gpu",
				Type:     "vllm",
				Models:   []string{"served-model"},
				Health:   "available",
				Capacity: 2,
			},
			{
				ID:       "prov-ollama",
				Adapter:  "ollama",
				Type:     "ollama",
				Models:   []string{"served-model"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Catalog with both providers.
	catalog := []config.ModelCatalogEntry{
		{ID: "mixed-model", Providers: map[string]string{"prov-vllm": "served-model", "prov-ollama": "served-model"}},
	}

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-mixed",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
	})

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// SubmitProviderPool → should do exactly one admission.
	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "mixed-model",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}

	// The selected provider depends on admission order (vLLM or Ollama by
	// providerID alphabetical tie-break). Verify only one request was sent.
	time.Sleep(50 * time.Millisecond)

	capturedMu.Lock()
	defer capturedMu.Unlock()

	if capturedRunReq == nil {
		t.Fatal("expected exactly one RunRequest for normalized path; admission must be single")
	}

	// Verify the request target matches the catalog served model.
	if capturedRunReq.GetTarget() != "served-model" {
		t.Errorf("target: got %q, want %q", capturedRunReq.GetTarget(), "served-model")
	}

	// Verify inflight is exactly 1 (single admission).
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight (single admission), got %d", got)
	}
}

// TestDispatchRejectsStaleGenerationBeforeSend pins the pre-send generation
// fence on both execution paths: when the connection a lease was admitted under
// is superseded by a reconnect before dispatch, the normalized path releases the
// lease without ever sending a RunRequest, and the tunnel path releases it
// without ever sending a ProviderTunnelRequest, to the stale (dead) client. Both
// request parsers are wired to the stale connection so a leaked send of either
// type is caught directly.
func TestDispatchRejectsStaleGenerationBeforeSend(t *testing.T) {
	t.Run("normalized", func(t *testing.T) {
		staleGenerationFenceCase(t, providerExecutionPathNormalized)
	})
	t.Run("tunnel", func(t *testing.T) {
		staleGenerationFenceCase(t, providerExecutionPathTunnel)
	})
}

// staleGenerationFenceCase drives one execution path through its provider-pool
// dispatch entry point with a lease admitted under a connection a reconnect has
// already superseded, and asserts the fence releases the lease and fails with no
// request of any type reaching the stale connection.
func staleGenerationFenceCase(t *testing.T, path providerExecutionPath) {
	t.Helper()

	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	staleClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	var capturedMu sync.Mutex
	captured := 0
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {
		capturedMu.Lock()
		captured++
		capturedMu.Unlock()
	})
	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(*iop.ProviderTunnelRequest) {
		capturedMu.Lock()
		captured++
		capturedMu.Unlock()
	})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-stale",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{{Name: "ollama", Enabled: true}},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-stale", Adapter: "ollama", Type: "ollama", Models: []string{"served"}, Health: "available", Capacity: 1},
		},
	})

	reg := edgenode.NewRegistry()
	// entry1 is the connection admission selected; it holds the stale client.
	entry1 := &edgenode.NodeEntry{NodeID: "node-stale", LifecycleState: edgenode.LifecycleConnected, Client: staleClient}
	reg.Register(entry1)
	staleGeneration := entry1.ConnectionGeneration
	// A reconnect supersedes entry1: the registry's current owner is a strictly
	// higher generation, so entry1's lease is now stale.
	reg.Register(&edgenode.NodeEntry{NodeID: "node-stale", LifecycleState: edgenode.LifecycleConnected, Client: &toki.TcpClient{}})
	if reg.IsCurrentOwnerGeneration("node-stale", staleGeneration) {
		t.Fatal("reconnect should have superseded the stale generation")
	}

	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{
		{ID: "stale-model", Providers: map[string]string{"prov-stale": "served"}},
	})

	// Admit a lease bound to entry1's (now stale) generation on the path under test.
	selectedCand := candidateNode{
		entry:         entry1,
		providerID:    "prov-stale",
		servedTarget:  "served",
		generation:    staleGeneration,
		capacity:      1,
		adapter:       "ollama",
		providerType:  "ollama",
		executionPath: path,
	}
	admitted, _, err := svc.queue.admitWithReason(context.Background(), "stale-model", "ollama", "served", []candidateNode{selectedCand}, groupPolicy{}, nil, false, false)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}
	if leaseCount(svc.queue) != 1 {
		t.Fatalf("expected 1 live lease after admit, got %d", leaseCount(svc.queue))
	}
	reservation := newQueueReservation(svc.queue, admitted)

	switch path {
	case providerExecutionPathNormalized:
		req := SubmitRunRequest{ModelGroupKey: "stale-model", ProviderPool: true, Background: true}
		result, err := svc.dispatchProviderPoolRun(context.Background(), req, "ollama", "served", admitted, "dispatched", reservation)
		if err == nil {
			t.Fatal("expected stale-generation dispatch to be rejected before send")
		}
		if result != nil {
			t.Fatal("expected nil result on a fenced dispatch")
		}
		if !strings.Contains(err.Error(), "connection changed before dispatch") {
			t.Fatalf("unexpected error: %v", err)
		}
	case providerExecutionPathTunnel:
		req := ProviderPoolDispatchRequest{
			Run: SubmitRunRequest{ModelGroupKey: "stale-model", ProviderPool: true},
			Tunnel: SubmitProviderTunnelRequest{
				ModelGroupKey: "stale-model",
				ProviderPool:  true,
				Method:        "POST",
				Path:          "/v1/chat/completions",
			},
		}
		result, err := svc.dispatchProviderPoolTunnel(req, "ollama", "served", admitted, "dispatched", reservation)
		if err == nil {
			t.Fatal("expected stale-generation tunnel dispatch to be rejected before send")
		}
		if result != nil {
			t.Fatal("expected nil result on a fenced tunnel dispatch")
		}
		if !strings.Contains(err.Error(), "connection changed before dispatch") {
			t.Fatalf("unexpected error: %v", err)
		}
	default:
		t.Fatalf("unhandled execution path %q", path)
	}

	// The fenced dispatch must have released its lease and never sent a request.
	if got := leaseCount(svc.queue); got != 0 {
		t.Fatalf("expected the fenced dispatch to release its lease, got %d live leases", got)
	}
	time.Sleep(50 * time.Millisecond)
	capturedMu.Lock()
	sent := captured
	capturedMu.Unlock()
	if sent != 0 {
		t.Fatalf("a request reached the stale connection %d time(s); want none", sent)
	}
}

// TestSubmitProviderPoolDispatchInfoObservation verifies that the provider-pool
// one-shot dispatch carries selected provider id, provider type, and execution
// path in RunDispatch on both tunnel and normalized paths (SURFACE_OBS-1).
func TestSubmitProviderPoolDispatchInfoObservation(t *testing.T) {
	for _, tc := range []struct {
		name         string
		providerType string
		wantPath     providerPoolPath
		wantExecPath string
	}{
		{
			name:         "vllm_tunnel",
			providerType: "vllm",
			wantPath:     ProviderPoolPathTunnel,
			wantExecPath: "provider_tunnel",
		},
		{
			name:         "ollama_normalized",
			providerType: "ollama",
			wantPath:     ProviderPoolPathNormalized,
			wantExecPath: "normalized",
		},
		{
			name:         "cli_normalized",
			providerType: "cli",
			wantPath:     ProviderPoolPathNormalized,
			wantExecPath: "normalized",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			edgeConn, nodeConn := net.Pipe()
			t.Cleanup(func() {
				_ = edgeConn.Close()
				_ = nodeConn.Close()
			})

			parserMap := toki.ParserMap{
				toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.RunRequest{}
					return m, proto.Unmarshal(b, m)
				},
				toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.ProviderTunnelRequest{}
					return m, proto.Unmarshal(b, m)
				},
			}

			edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
			_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)

			store := edgenode.NewNodeStore()
			store.Add(&edgenode.NodeRecord{
				ID:      "node-obs",
				Runtime: config.RuntimeConf{Concurrency: 2},
				Adapters: config.AdaptersConf{
					VllmInstances:         []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
					OllamaInstances:       []config.OllamaInstanceConf{{Name: "ollama", Enabled: true}},
					OpenAICompatInstances: []config.OpenAICompatInstanceConf{{Name: "openai", Enabled: true}},
				},
				Providers: []config.NodeProviderConf{
					{
						ID:       "prov-obs",
						Adapter:  "vllm-gpu",
						Type:     tc.providerType,
						Models:   []string{"served-model"},
						Health:   "available",
						Capacity: 2,
					},
				},
			})

			catalog := []config.ModelCatalogEntry{
				{ID: "obs-model", Providers: map[string]string{"prov-obs": "served-model"}},
			}

			reg := edgenode.NewRegistry()
			reg.Register(&edgenode.NodeEntry{
				NodeID:         "node-obs",
				LifecycleState: edgenode.LifecycleConnected,
				Client:         edgeClient,
			})

			bus := edgeevents.NewBus()
			svc := New(reg, bus)
			svc.SetNodeStore(store)
			svc.SetModelCatalog(catalog)

			result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run: SubmitRunRequest{
					ModelGroupKey: "obs-model",
					ProviderPool:  true,
					Background:    true,
				},
			})
			if err != nil {
				t.Fatalf("SubmitProviderPool: %v", err)
			}
			if result == nil {
				t.Fatal("expected non-nil result")
			}
			if result.Path != tc.wantPath {
				t.Fatalf("path: got %q, want %q", result.Path, tc.wantPath)
			}

			disp := result.DispatchInfo
			if disp.ProviderID != "prov-obs" {
				t.Errorf("provider_id: got %q, want %q", disp.ProviderID, "prov-obs")
			}
			if disp.ProviderType != tc.providerType {
				t.Errorf("provider_type: got %q, want %q", disp.ProviderType, tc.providerType)
			}
			if disp.ExecutionPath != tc.wantExecPath {
				t.Errorf("execution_path: got %q, want %q", disp.ExecutionPath, tc.wantExecPath)
			}
			if disp.QueueReason != "dispatched" {
				t.Errorf("queue_reason: got %q, want %q", disp.QueueReason, "dispatched")
			}
			if disp.Adapter != "vllm-gpu" {
				t.Errorf("adapter: got %q, want %q", disp.Adapter, "vllm-gpu")
			}
			if disp.Target != "served-model" {
				t.Errorf("target: got %q, want %q", disp.Target, "served-model")
			}

			// Also verify tunnel/normalized handle DispatchInfo matches.
			switch result.Path {
			case ProviderPoolPathTunnel:
				if result.Tunnel == nil {
					t.Fatal("expected tunnel handle for tunnel path")
				}
				hDisp := result.Tunnel.Dispatch()
				if hDisp.ProviderID != "prov-obs" {
					t.Errorf("tunnel handle provider_id: got %q, want %q", hDisp.ProviderID, "prov-obs")
				}
				if hDisp.ProviderType != tc.providerType {
					t.Errorf("tunnel handle provider_type: got %q, want %q", hDisp.ProviderType, tc.providerType)
				}
				if hDisp.ExecutionPath != tc.wantExecPath {
					t.Errorf("tunnel handle execution_path: got %q, want %q", hDisp.ExecutionPath, tc.wantExecPath)
				}
				// Close releases the inflight slot.
				result.Tunnel.Close()
				if got := inflightRunCount(svc.queue); got != 0 {
					t.Fatalf("expected 0 inflight after tunnel close, got %d", got)
				}
			case ProviderPoolPathNormalized:
				if result.Run == nil {
					t.Fatal("expected run handle for normalized path")
				}
				hDisp := result.Run.Dispatch()
				if hDisp.ProviderID != "prov-obs" {
					t.Errorf("run handle provider_id: got %q, want %q", hDisp.ProviderID, "prov-obs")
				}
				if hDisp.ProviderType != tc.providerType {
					t.Errorf("run handle provider_type: got %q, want %q", hDisp.ProviderType, tc.providerType)
				}
				if hDisp.ExecutionPath != tc.wantExecPath {
					t.Errorf("run handle execution_path: got %q, want %q", hDisp.ExecutionPath, tc.wantExecPath)
				}
				// Inflight should be 1 for normalized (run hasn't completed yet).
				if got := inflightRunCount(svc.queue); got != 1 {
					t.Fatalf("expected 1 inflight for normalized, got %d", got)
				}
				result.Run.Close()
			}
		})
	}
}

// TestPendingClientDirectSendFence asserts the pending→ready dispatch fence at
// the granularity of the exact protobuf each direct wire carries. While the
// connection is pending every direct send path fails and nothing reaches the
// node. After the ready transition each wire is delivered exactly once, and its
// identity fields — not merely its type — are asserted against a cloned copy of
// the received message, so a request that dispatched with an empty or wrong field
// can no longer pass by type count alone.
func TestPendingClientDirectSendFence(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.CancelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.CancelRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	// Each listener clones the received protobuf so its exact identity fields can
	// be asserted; fire-and-forget wires arrive on buffered channels, and the
	// synchronous NodeCommandRequests accumulate under a lock keyed by type.
	runReqs := make(chan *iop.RunRequest, 4)
	tunnelReqs := make(chan *iop.ProviderTunnelRequest, 4)
	cancelReqs := make(chan *iop.CancelRequest, 4)
	var cmdMu sync.Mutex
	cmdReqs := make(map[iop.NodeCommandType][]*iop.NodeCommandRequest)

	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		runReqs <- proto.Clone(req).(*iop.RunRequest)
	})
	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
		tunnelReqs <- proto.Clone(req).(*iop.ProviderTunnelRequest)
	})
	toki.AddListenerTyped[*iop.CancelRequest](&nodeClient.Communicator, func(req *iop.CancelRequest) {
		cancelReqs <- proto.Clone(req).(*iop.CancelRequest)
	})
	toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		cmdMu.Lock()
		cmdReqs[req.GetType()] = append(cmdReqs[req.GetType()], proto.Clone(req).(*iop.NodeCommandRequest))
		cmdMu.Unlock()
		return &iop.NodeCommandResponse{
			RequestId: req.GetRequestId(),
			Type:      req.GetType(),
			Result:    map[string]string{"result": "ok", "status_code": "200"},
		}, nil
	})

	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID:        "node-fence-1",
		Alias:         "fence-alias",
		AgentKind:     config.AgentKindGenericNode,
		Client:        edgeClient,
		DispatchReady: false, // start pending
	}
	reg.RegisterIfAbsent(entry)
	svc := New(reg, nil)
	ctx := context.Background()

	// 1. Observability lookups still see the pending connection.
	if _, err := svc.ResolveNode("node-fence-1"); err != nil {
		t.Errorf("ResolveNode failed for pending: %v", err)
	}
	snap, err := svc.ResolveNodeSnapshot("node-fence-1")
	if err != nil {
		t.Errorf("ResolveNodeSnapshot failed for pending: %v", err)
	}
	if !snap.Connected {
		t.Errorf("expected Connected=true for pending snapshot, got %v", snap.Connected)
	}

	// 2. Every direct send path fails while pending and puts nothing on the wire.
	if _, err := svc.SubmitRun(ctx, SubmitRunRequest{NodeRef: "node-fence-1", Adapter: "cli", Target: "codex"}); err == nil {
		t.Error("expected SubmitRun on pending node to fail")
	}
	if _, err := svc.submitProviderTunnelDirect(SubmitProviderTunnelRequest{NodeRef: "node-fence-1", Adapter: "cli", Target: "codex"}); err == nil {
		t.Error("expected submitProviderTunnelDirect on pending node to fail")
	}
	if _, err := svc.CancelRun(ctx, CancelRunRequest{NodeRef: "node-fence-1", RunID: "run-1"}); err == nil {
		t.Error("expected CancelRun on pending node to fail")
	}
	if _, err := svc.TerminateSession(ctx, TerminateSessionRequest{NodeRef: "node-fence-1"}); err == nil {
		t.Error("expected TerminateSession on pending node to fail")
	}
	if _, err := svc.UsageStatus(ctx, UsageStatusRequest{NodeRef: "node-fence-1"}); err == nil {
		t.Error("expected UsageStatus on pending node to fail")
	}
	pendingResp, _ := svc.ExecuteCommand(ctx, &iop.EdgeCommandRequest{
		CommandId:      "cmd-pending",
		Operation:      "agent.command",
		TargetSelector: "node-fence-1",
		Parameters:     map[string]string{"command": "capabilities"},
	}, func(*iop.EdgeCommandEvent) {})
	if pendingResp == nil || pendingResp.Status != "error" {
		t.Errorf("expected ExecuteCommand on pending node to error, got %v", pendingResp.GetStatus())
	}

	if n := len(runReqs) + len(tunnelReqs) + len(cancelReqs); n != 0 {
		t.Fatalf("pending node received %d fire-and-forget wire messages, want 0", n)
	}
	cmdMu.Lock()
	pendingCmds := len(cmdReqs)
	cmdMu.Unlock()
	if pendingCmds != 0 {
		t.Fatalf("pending node received %d command wire types, want 0", pendingCmds)
	}

	// 3. Ready opens dispatch.
	if _, transitioned, ok := reg.MarkDispatchReadyIfClient("node-fence-1", edgeClient); !ok || !transitioned {
		t.Fatalf("failed to transition node to ready: ok=%v transitioned=%v", ok, transitioned)
	}
	snap2, err := svc.ResolveNodeSnapshot("node-fence-1")
	if err != nil {
		t.Errorf("ResolveNodeSnapshot failed after ready: %v", err)
	}
	if !snap2.Connected {
		t.Errorf("expected Connected=true after ready, got %v", snap2.Connected)
	}

	// 4. Each direct wire now dispatches exactly once with its exact payload.
	if _, err := svc.SubmitRun(ctx, SubmitRunRequest{
		NodeRef: "node-fence-1", RunID: "run-fence-run", Adapter: "cli", Target: "codex",
		SessionID: "sess-run", Background: true, Prompt: "hello-fence",
		Metadata: map[string]string{"source": "fence-test"},
	}); err != nil {
		t.Errorf("SubmitRun after ready: %v", err)
	}
	if _, err := svc.submitProviderTunnelDirect(SubmitProviderTunnelRequest{
		NodeRef: "node-fence-1", RunID: "run-fence-tunnel", Adapter: "cli", Target: "codex",
		SessionID: "sess-tunnel", Method: "POST", Path: "/v1/chat", Body: []byte("tunnel-body"),
	}); err != nil {
		t.Errorf("submitProviderTunnelDirect after ready: %v", err)
	}
	if _, err := svc.CancelRun(ctx, CancelRunRequest{
		NodeRef: "node-fence-1", RunID: "run-fence-cancel", Adapter: "cli", Target: "codex", SessionID: "sess-cancel",
	}); err != nil {
		t.Errorf("CancelRun after ready: %v", err)
	}
	if _, err := svc.TerminateSession(ctx, TerminateSessionRequest{
		NodeRef: "node-fence-1", Adapter: "cli", Target: "codex", SessionID: "sess-term",
	}); err != nil {
		t.Errorf("TerminateSession after ready: %v", err)
	}
	if _, err := svc.UsageStatus(ctx, UsageStatusRequest{
		NodeRef: "node-fence-1", Adapter: "cli", Target: "codex", SessionID: "sess-usage",
	}); err != nil {
		t.Errorf("UsageStatus after ready: %v", err)
	}
	if _, err := svc.SessionList(ctx, NodeCommandRequestSpec{
		NodeRef: "node-fence-1", Adapter: "cli", Target: "codex", SessionID: "sess-list",
	}); err != nil {
		t.Errorf("SessionList after ready: %v", err)
	}
	if _, err := svc.TransportStatus(ctx, NodeCommandRequestSpec{
		NodeRef: "node-fence-1", Adapter: "cli", Target: "codex", SessionID: "sess-transport",
	}); err != nil {
		t.Errorf("TransportStatus after ready: %v", err)
	}
	if _, err := svc.OllamaAPI(ctx, OllamaAPIRequest{
		NodeRef: "node-fence-1", Adapter: "ollama", Target: "llama", Method: "POST", Path: "/api/chat", Body: "ollama-body",
	}); err != nil {
		t.Errorf("OllamaAPI after ready: %v", err)
	}
	readyResp, _ := svc.ExecuteCommand(ctx, &iop.EdgeCommandRequest{
		CommandId:      "cmd-ready",
		Operation:      "agent.command",
		TargetSelector: "node-fence-1",
		Parameters:     map[string]string{"command": "capabilities"},
	}, func(*iop.EdgeCommandEvent) {})
	if readyResp == nil || readyResp.Status == "error" {
		t.Errorf("ExecuteCommand after ready failed: %v", readyResp.GetError())
	}

	// 4.1 RunRequest identity.
	run := recvWire(t, runReqs, "RunRequest")
	if got := run.GetRunId(); got != "run-fence-run" {
		t.Errorf("RunRequest.run_id = %q, want run-fence-run", got)
	}
	if got := run.GetAdapter(); got != "cli" {
		t.Errorf("RunRequest.adapter = %q, want cli", got)
	}
	if got := run.GetTarget(); got != "codex" {
		t.Errorf("RunRequest.target = %q, want codex", got)
	}
	if !run.GetBackground() {
		t.Errorf("RunRequest.background = false, want true")
	}
	if got := run.GetSessionId(); got != "sess-run" {
		t.Errorf("RunRequest.session_id = %q, want sess-run", got)
	}
	if got := run.GetInput().GetFields()["prompt"].GetStringValue(); got != "hello-fence" {
		t.Errorf("RunRequest.input.prompt = %q, want hello-fence", got)
	}
	if got := run.GetMetadata()["source"]; got != "fence-test" {
		t.Errorf("RunRequest.metadata.source = %q, want fence-test", got)
	}
	if run.GetSessionMode() != iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING {
		t.Errorf("RunRequest.session_mode = %v, want CREATE_IF_MISSING", run.GetSessionMode())
	}
	assertNoExtra(t, runReqs, "RunRequest")

	// 4.2 ProviderTunnelRequest identity.
	tunnel := recvWire(t, tunnelReqs, "ProviderTunnelRequest")
	if got := tunnel.GetRunId(); got != "run-fence-tunnel" {
		t.Errorf("ProviderTunnelRequest.run_id = %q, want run-fence-tunnel", got)
	}
	if got := tunnel.GetTunnelId(); got != "run-fence-tunnel-tunnel" {
		t.Errorf("ProviderTunnelRequest.tunnel_id = %q, want run-fence-tunnel-tunnel", got)
	}
	if got := tunnel.GetAdapter(); got != "cli" {
		t.Errorf("ProviderTunnelRequest.adapter = %q, want cli", got)
	}
	if got := tunnel.GetTarget(); got != "codex" {
		t.Errorf("ProviderTunnelRequest.target = %q, want codex", got)
	}
	if got := tunnel.GetMethod(); got != "POST" {
		t.Errorf("ProviderTunnelRequest.method = %q, want POST", got)
	}
	if got := tunnel.GetPath(); got != "/v1/chat" {
		t.Errorf("ProviderTunnelRequest.path = %q, want /v1/chat", got)
	}
	if got := string(tunnel.GetBody()); got != "tunnel-body" {
		t.Errorf("ProviderTunnelRequest.body = %q, want tunnel-body", got)
	}
	if got := tunnel.GetSessionId(); got != "sess-tunnel" {
		t.Errorf("ProviderTunnelRequest.session_id = %q, want sess-tunnel", got)
	}
	assertNoExtra(t, tunnelReqs, "ProviderTunnelRequest")

	// 4.3 CancelRequest: one CANCEL_RUN and one TERMINATE_SESSION.
	firstCancel := recvWire(t, cancelReqs, "CancelRequest")
	secondCancel := recvWire(t, cancelReqs, "CancelRequest")
	cancels := map[iop.CancelAction]*iop.CancelRequest{
		firstCancel.GetAction():  firstCancel,
		secondCancel.GetAction(): secondCancel,
	}
	if len(cancels) != 2 {
		t.Fatalf("expected distinct CANCEL_RUN and TERMINATE_SESSION, got actions %v and %v", firstCancel.GetAction(), secondCancel.GetAction())
	}
	if cancelRun := cancels[iop.CancelAction_CANCEL_ACTION_CANCEL_RUN]; cancelRun == nil {
		t.Error("missing CANCEL_RUN cancel request")
	} else {
		if got := cancelRun.GetRunId(); got != "run-fence-cancel" {
			t.Errorf("CANCEL_RUN.run_id = %q, want run-fence-cancel", got)
		}
		if got := cancelRun.GetSessionId(); got != "sess-cancel" {
			t.Errorf("CANCEL_RUN.session_id = %q, want sess-cancel", got)
		}
		if got := cancelRun.GetAdapter(); got != "cli" {
			t.Errorf("CANCEL_RUN.adapter = %q, want cli", got)
		}
		if got := cancelRun.GetTarget(); got != "codex" {
			t.Errorf("CANCEL_RUN.target = %q, want codex", got)
		}
	}
	if term := cancels[iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION]; term == nil {
		t.Error("missing TERMINATE_SESSION cancel request")
	} else {
		if got := term.GetSessionId(); got != "sess-term" {
			t.Errorf("TERMINATE_SESSION.session_id = %q, want sess-term", got)
		}
		if got := term.GetRunId(); got != "" {
			t.Errorf("TERMINATE_SESSION.run_id = %q, want empty", got)
		}
	}
	assertNoExtra(t, cancelReqs, "CancelRequest")

	// 4.4 NodeCommandRequests: each type exactly once with its exact fields.
	cmdMu.Lock()
	defer cmdMu.Unlock()
	expectCommand := func(typ iop.NodeCommandType, wantAdapter, wantTarget, wantSession, wantIDPrefix string, extra func(*iop.NodeCommandRequest)) {
		reqs := cmdReqs[typ]
		if len(reqs) != 1 {
			t.Errorf("command %v received %d times, want 1", typ, len(reqs))
			return
		}
		req := reqs[0]
		if got := req.GetAdapter(); got != wantAdapter {
			t.Errorf("%v adapter = %q, want %q", typ, got, wantAdapter)
		}
		if got := req.GetTarget(); got != wantTarget {
			t.Errorf("%v target = %q, want %q", typ, got, wantTarget)
		}
		if got := req.GetSessionId(); got != wantSession {
			t.Errorf("%v session_id = %q, want %q", typ, got, wantSession)
		}
		if got := req.GetRequestId(); !strings.HasPrefix(got, wantIDPrefix) {
			t.Errorf("%v request_id = %q, want prefix %q", typ, got, wantIDPrefix)
		}
		if extra != nil {
			extra(req)
		}
	}
	expectCommand(iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS, "cli", "codex", "sess-usage", "status-", nil)
	expectCommand(iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, "cli", "codex", "sess-list", "sessions-", nil)
	expectCommand(iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "cli", "codex", "sess-transport", "transport-", nil)
	expectCommand(iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API, "ollama", "llama", "default", "ollama-", func(req *iop.NodeCommandRequest) {
		if got := req.GetMetadata()["ollama_method"]; got != "POST" {
			t.Errorf("OLLAMA_API metadata.ollama_method = %q, want POST", got)
		}
		if got := req.GetMetadata()["ollama_path"]; got != "/api/chat" {
			t.Errorf("OLLAMA_API metadata.ollama_path = %q, want /api/chat", got)
		}
		if got := req.GetMetadata()["ollama_body"]; got != "ollama-body" {
			t.Errorf("OLLAMA_API metadata.ollama_body = %q, want ollama-body", got)
		}
	})
	expectCommand(iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "", "", "default", "caps-", nil)
	if len(cmdReqs) != 5 {
		t.Errorf("received %d distinct command types, want 5", len(cmdReqs))
	}
}

// recvWire returns the next captured wire of type T or fails after a timeout.
func recvWire[T proto.Message](t *testing.T, ch chan T, name string) T {
	t.Helper()
	select {
	case v := <-ch:
		return v
	case <-time.After(2 * time.Second):
		t.Fatalf("timed out waiting for %s wire", name)
		var zero T
		return zero
	}
}

// assertNoExtra fails if any additional wire of type T is buffered. No further
// send was issued after the expected ones, so a non-blocking read is a
// deterministic exactly-once guard.
func assertNoExtra[T proto.Message](t *testing.T, ch chan T, name string) {
	t.Helper()
	select {
	case v := <-ch:
		t.Errorf("unexpected extra %s wire: %v", name, v)
	default:
	}
}

// gateConn blocks the first (and every subsequent) Write until release is
// closed, signalling entered once the write is in progress. The provider-pool
// dispatch performs its wire send inside WithCurrentDispatchOwner while holding
// the registry lock, so blocking the write here holds that lock and lets a test
// schedule a competing disconnect deterministically against it.
type gateConn struct {
	net.Conn
	entered   chan struct{}
	release   chan struct{}
	enterOnce sync.Once
}

func (c *gateConn) Write(b []byte) (int, error) {
	c.enterOnce.Do(func() { close(c.entered) })
	<-c.release
	return c.Conn.Write(b)
}

// newProviderPoolRaceService builds a Service whose registry holds one ready
// node with a single capacity-1 provider serving model-1. provType selects the
// dispatch path: an OpenAI-compatible type (e.g. "vllm") routes to the tunnel
// path, everything else (e.g. "openai") to the normalized run path.
func newProviderPoolRaceService(t *testing.T, nodeID, provType string, edgeClient *toki.TcpClient) (*Service, *edgenode.Registry, *edgenode.NodeEntry) {
	t.Helper()
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID:        nodeID,
		Alias:         nodeID + "-alias",
		AgentKind:     config.AgentKindGenericNode,
		Client:        edgeClient,
		DispatchReady: true,
	}
	reg.Register(entry)
	svc := New(reg, edgeevents.NewBus())
	store, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: nodeID, Alias: nodeID + "-alias", Token: "token", AgentKind: config.AgentKindGenericNode, Providers: []config.NodeProviderConf{
			{ID: "prov-1", Type: provType, Models: []string{"model-1"}, Health: "available", Capacity: 1},
		}},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{{ID: "model-1", Providers: map[string]string{"prov-1": "model-1"}}})
	return svc, reg, entry
}

// assertQueueSettled asserts every lease is released and every provider/group
// in-flight counter is back to zero, so a race can never leak a slot.
func assertQueueSettled(t *testing.T, q *modelQueueManager) {
	t.Helper()
	q.mu.Lock()
	defer q.mu.Unlock()
	if len(q.leases) != 0 {
		t.Errorf("leases not settled: %d remain", len(q.leases))
	}
	if len(q.leaseByRun) != 0 {
		t.Errorf("leaseByRun not settled: %d remain", len(q.leaseByRun))
	}
	for key, res := range q.resources {
		if res.inFlight != 0 || res.longInFlight != 0 {
			t.Errorf("provider resource %v not settled: inFlight=%d longInFlight=%d", key, res.inFlight, res.longInFlight)
		}
	}
	for gk, g := range q.groups {
		for slot, n := range g.inflight {
			if n != 0 {
				t.Errorf("group %q slot %q inflight=%d, want 0", gk, slot, n)
			}
		}
		for slot, n := range g.longInflight {
			if n != 0 {
				t.Errorf("group %q slot %q longInflight=%d, want 0", gk, slot, n)
			}
		}
	}
}

// assertTrackedLease proves the send-winning barrier was reached only after the
// exact run identity had been bound to one live lease for the expected owner
// generation. The later disconnect/terminal assertions then prove that same
// identity is removed without leaving a duplicate decrement behind.
func assertTrackedLease(t *testing.T, q *modelQueueManager, runID, nodeID string, generation uint64) {
	t.Helper()
	q.mu.Lock()
	defer q.mu.Unlock()
	leaseID, ok := q.leaseByRun[runID]
	if !ok || leaseID == 0 {
		t.Fatalf("run %q has no tracked lease", runID)
	}
	lease := q.leases[leaseID]
	if lease == nil {
		t.Fatalf("run %q maps to missing lease %d", runID, leaseID)
	}
	if lease.nodeID != nodeID || lease.generation != generation {
		t.Fatalf("run %q lease owner=(node=%q generation=%d), want (%q,%d)",
			runID, lease.nodeID, lease.generation, nodeID, generation)
	}
}

func disconnectCurrentOwner(svc *Service, reg *edgenode.Registry, entry *edgenode.NodeEntry, reason string) (uint64, bool) {
	generation, ok := reg.UnregisterIfClient(entry.NodeID, entry.Client)
	if ok {
		svc.HandleNodeDisconnect(entry.NodeID, generation, reason)
	}
	return generation, ok
}

// TestProviderPoolDispatchRunDisconnectRace pins the normalized send/disconnect
// race in two deterministic directions instead of inferring order from sleeps.
//
//   - send-winning: the dispatch enters its wire send under the owner gate
//     (registry lock held); a disconnect racing in blocks on that lock, the wire
//     is written exactly once, and the authoritative disconnect then releases the
//     handed-off lease exactly once, leaving every counter at zero.
//   - disconnect-winning: the disconnect settles (unregister + fence) while the
//     dispatch is parked at its post-admission PrepareRun barrier, before the
//     owner gate. The gate then fences the stale generation, so no wire and no
//     run handle escape and the reserved lease is released exactly once.
func TestProviderPoolDispatchRunDisconnectRace(t *testing.T) {
	runParser := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	t.Run("send-winning", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		defer edgeConn.Close()
		defer nodeConn.Close()
		gate := &gateConn{Conn: edgeConn, entered: make(chan struct{}), release: make(chan struct{})}
		edgeClient := toki.NewTcpClient(gate, 0, 0, runParser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, runParser)
		runReqs := make(chan *iop.RunRequest, 4)
		toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
			runReqs <- proto.Clone(req).(*iop.RunRequest)
		})

		svc, reg, entry := newProviderPoolRaceService(t, "node-race-run", "openai", edgeClient)

		dispatchDone := make(chan struct{})
		var result *ProviderPoolDispatchResult
		var dispatchErr error
		go func() {
			result, dispatchErr = svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run: SubmitRunRequest{NodeRef: "node-race-run", RunID: "run-send-win", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true, Background: true},
			})
			close(dispatchDone)
		}()

		select {
		case <-gate.entered:
		case <-time.After(3 * time.Second):
			t.Fatal("dispatch did not reach the wire send")
		}
		assertTrackedLease(t, svc.queue, "run-send-win", "node-race-run", entry.ConnectionGeneration)

		type disconnectOutcome struct {
			generation uint64
			ok         bool
		}
		disconnectDone := make(chan disconnectOutcome, 1)
		go func() {
			generation, ok := disconnectCurrentOwner(svc, reg, entry, "test-disconnect")
			disconnectDone <- disconnectOutcome{generation: generation, ok: ok}
		}()

		close(gate.release)
		<-dispatchDone
		disconnected := <-disconnectDone
		if !disconnected.ok || disconnected.generation != entry.ConnectionGeneration {
			t.Fatalf("send-winning disconnect=(generation=%d ok=%v), want (%d,true)",
				disconnected.generation, disconnected.ok, entry.ConnectionGeneration)
		}

		if dispatchErr != nil {
			t.Fatalf("send-winning dispatch must succeed, got %v", dispatchErr)
		}
		if result == nil || result.Run == nil {
			t.Fatal("send-winning must return a run handle")
		}
		wire := recvWire(t, runReqs, "RunRequest")
		if got := wire.GetRunId(); got != "run-send-win" {
			t.Errorf("wire run_id = %q, want run-send-win", got)
		}
		assertNoExtra(t, runReqs, "RunRequest")
		assertQueueSettled(t, svc.queue)

		// A terminal callback losing to the authoritative disconnect is an
		// idempotent no-op for the same tracked lease; closing the handle must not
		// reopen or double-release it either.
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wire.GetRunId(), Type: "complete", NodeId: "node-race-run"})
		result.Run.Close()
		assertQueueSettled(t, svc.queue)
	})

	t.Run("disconnect-winning", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		defer edgeConn.Close()
		defer nodeConn.Close()
		edgeClient := toki.NewTcpClient(edgeConn, 0, 0, runParser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, runParser)
		runReqs := make(chan *iop.RunRequest, 4)
		toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
			runReqs <- proto.Clone(req).(*iop.RunRequest)
		})

		svc, reg, entry := newProviderPoolRaceService(t, "node-race-run", "openai", edgeClient)

		reachedPrepare := make(chan struct{})
		releasePrepare := make(chan struct{})
		dispatchDone := make(chan struct{})
		var result *ProviderPoolDispatchResult
		var dispatchErr error
		go func() {
			result, dispatchErr = svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run: SubmitRunRequest{NodeRef: "node-race-run", RunID: "run-disc-win", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true, Background: true},
				PrepareRun: func(req SubmitRunRequest) (SubmitRunRequest, error) {
					close(reachedPrepare)
					<-releasePrepare
					return req, nil
				},
			})
			close(dispatchDone)
		}()

		<-reachedPrepare
		if got := leaseCount(svc.queue); got != 1 {
			t.Fatalf("disconnect-winning reserved leases=%d before fence, want 1", got)
		}
		disconnectedGen, ok := disconnectCurrentOwner(svc, reg, entry, "test-disconnect")
		if !ok || disconnectedGen != entry.ConnectionGeneration {
			t.Fatalf("disconnect-winning disconnect=(generation=%d ok=%v), want (%d,true)",
				disconnectedGen, ok, entry.ConnectionGeneration)
		}
		close(releasePrepare)
		<-dispatchDone

		if dispatchErr == nil {
			t.Fatal("disconnect-winning dispatch must fail")
		}
		if result != nil {
			t.Fatalf("disconnect-winning must not return a result, got %+v", result)
		}
		assertNoExtra(t, runReqs, "RunRequest")
		assertQueueSettled(t, svc.queue)
	})
}

// TestProviderPoolDispatchTunnelDisconnectRace mirrors the normalized race for
// the tunnel/passthrough path: a send-winning subtest that writes exactly one
// ProviderTunnelRequest then releases the lease on the authoritative disconnect,
// and a disconnect-winning subtest whose PrepareTunnel barrier lets the fence
// win so no tunnel frame and no handle escape.
func TestProviderPoolDispatchTunnelDisconnectRace(t *testing.T) {
	tunnelParser := toki.ParserMap{
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	t.Run("send-winning", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		defer edgeConn.Close()
		defer nodeConn.Close()
		gate := &gateConn{Conn: edgeConn, entered: make(chan struct{}), release: make(chan struct{})}
		edgeClient := toki.NewTcpClient(gate, 0, 0, tunnelParser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, tunnelParser)
		tunnelReqs := make(chan *iop.ProviderTunnelRequest, 4)
		toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
			tunnelReqs <- proto.Clone(req).(*iop.ProviderTunnelRequest)
		})

		svc, reg, entry := newProviderPoolRaceService(t, "node-race-tunnel", "vllm", edgeClient)

		dispatchDone := make(chan struct{})
		var result *ProviderPoolDispatchResult
		var dispatchErr error
		go func() {
			result, dispatchErr = svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run:    SubmitRunRequest{NodeRef: "node-race-tunnel", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true},
				Tunnel: SubmitProviderTunnelRequest{NodeRef: "node-race-tunnel", RunID: "run-tunnel-send-win", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true},
			})
			close(dispatchDone)
		}()

		select {
		case <-gate.entered:
		case <-time.After(3 * time.Second):
			t.Fatal("tunnel dispatch did not reach the wire send")
		}
		assertTrackedLease(t, svc.queue, "run-tunnel-send-win", "node-race-tunnel", entry.ConnectionGeneration)

		type disconnectOutcome struct {
			generation uint64
			ok         bool
		}
		disconnectDone := make(chan disconnectOutcome, 1)
		go func() {
			generation, ok := disconnectCurrentOwner(svc, reg, entry, "test-disconnect")
			disconnectDone <- disconnectOutcome{generation: generation, ok: ok}
		}()

		close(gate.release)
		<-dispatchDone
		disconnected := <-disconnectDone
		if !disconnected.ok || disconnected.generation != entry.ConnectionGeneration {
			t.Fatalf("send-winning tunnel disconnect=(generation=%d ok=%v), want (%d,true)",
				disconnected.generation, disconnected.ok, entry.ConnectionGeneration)
		}

		if dispatchErr != nil {
			t.Fatalf("send-winning tunnel dispatch must succeed, got %v", dispatchErr)
		}
		if result == nil || result.Tunnel == nil {
			t.Fatal("send-winning must return a tunnel handle")
		}
		wire := recvWire(t, tunnelReqs, "ProviderTunnelRequest")
		if got := wire.GetRunId(); got != "run-tunnel-send-win" {
			t.Errorf("wire run_id = %q, want run-tunnel-send-win", got)
		}
		assertNoExtra(t, tunnelReqs, "ProviderTunnelRequest")
		assertQueueSettled(t, svc.queue)

		// Deliver the competing tunnel terminal after disconnect. The handle sees
		// exactly one END and closes, while the already-settled lease/counters stay
		// at zero when Close races afterward.
		stream := result.Tunnel.Stream()
		svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{
			RunId: wire.GetRunId(), TunnelId: wire.GetTunnelId(),
			Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true,
		})
		select {
		case terminal, ok := <-stream.Frames:
			if !ok || terminal.GetKind() != iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END || !terminal.GetEnd() {
				t.Fatalf("tunnel terminal=%v open=%v, want one END frame", terminal, ok)
			}
		case <-time.After(2 * time.Second):
			t.Fatal("timed out waiting for tunnel END frame")
		}
		select {
		case extra, ok := <-stream.Frames:
			if ok {
				t.Fatalf("unexpected extra tunnel terminal/frame: %v", extra)
			}
		case <-time.After(2 * time.Second):
			t.Fatal("tunnel stream did not close after END")
		}
		result.Tunnel.Close()
		assertQueueSettled(t, svc.queue)
	})

	t.Run("disconnect-winning", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		defer edgeConn.Close()
		defer nodeConn.Close()
		edgeClient := toki.NewTcpClient(edgeConn, 0, 0, tunnelParser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, tunnelParser)
		tunnelReqs := make(chan *iop.ProviderTunnelRequest, 4)
		toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
			tunnelReqs <- proto.Clone(req).(*iop.ProviderTunnelRequest)
		})

		svc, reg, entry := newProviderPoolRaceService(t, "node-race-tunnel", "vllm", edgeClient)

		reachedPrepare := make(chan struct{})
		releasePrepare := make(chan struct{})
		dispatchDone := make(chan struct{})
		var result *ProviderPoolDispatchResult
		var dispatchErr error
		go func() {
			result, dispatchErr = svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run:    SubmitRunRequest{NodeRef: "node-race-tunnel", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true},
				Tunnel: SubmitProviderTunnelRequest{NodeRef: "node-race-tunnel", RunID: "run-tunnel-disc-win", ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true},
				PrepareTunnel: func(req SubmitProviderTunnelRequest) (SubmitProviderTunnelRequest, error) {
					close(reachedPrepare)
					<-releasePrepare
					return req, nil
				},
			})
			close(dispatchDone)
		}()

		<-reachedPrepare
		if got := leaseCount(svc.queue); got != 1 {
			t.Fatalf("disconnect-winning tunnel reserved leases=%d before fence, want 1", got)
		}
		disconnectedGen, ok := disconnectCurrentOwner(svc, reg, entry, "test-disconnect")
		if !ok || disconnectedGen != entry.ConnectionGeneration {
			t.Fatalf("disconnect-winning tunnel disconnect=(generation=%d ok=%v), want (%d,true)",
				disconnectedGen, ok, entry.ConnectionGeneration)
		}
		close(releasePrepare)
		<-dispatchDone

		if dispatchErr == nil {
			t.Fatal("disconnect-winning tunnel dispatch must fail")
		}
		if result != nil {
			t.Fatalf("disconnect-winning must not return a result, got %+v", result)
		}
		assertNoExtra(t, tunnelReqs, "ProviderTunnelRequest")
		assertQueueSettled(t, svc.queue)
	})
}

// providerResourceState reads a provider resource's in-flight count and orphan
// flag under the manager lock for white-box assertions.
func providerResourceInflight(q *modelQueueManager, nodeID, providerID string) (inFlight int, orphan, ok bool) {
	q.mu.Lock()
	defer q.mu.Unlock()
	res, exists := q.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
	if !exists {
		return 0, false, false
	}
	return res.inFlight, res.orphan, true
}

func queuedItemCount(q *modelQueueManager) int {
	q.mu.Lock()
	defer q.mu.Unlock()
	n := 0
	for _, g := range q.groups {
		n += len(g.queue)
	}
	return n
}

// TestReconnectActivationLinearizesAgainstDisconnect drives the real queue
// through an actual reconnect and proves the ownership linearization the fix
// establishes: a stranded provider-pool waiter is pumped exactly once, and only
// by the current reconnect generation. A stale (superseded) generation's
// HandleNodeConnect is a no-op — it neither re-activates the orphaned resource
// nor pumps the waiter — while the current generation activates and dispatches it
// exactly once, and all counters settle to zero afterward.
func TestReconnectActivationLinearizesAgainstDisconnect(t *testing.T) {
	parser := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	altEdge, altNode := net.Pipe()
	defer altEdge.Close()
	defer altNode.Close()
	reconEdge, reconNode := net.Pipe()
	defer reconEdge.Close()
	defer reconNode.Close()

	altEdgeClient := toki.NewTcpClient(altEdge, 0, 0, parser)
	altNodeClient := toki.NewTcpClient(altNode, 0, 0, parser)
	reconEdgeClient := toki.NewTcpClient(reconEdge, 0, 0, parser)
	reconNodeClient := toki.NewTcpClient(reconNode, 0, 0, parser)

	altRuns := make(chan *iop.RunRequest, 4)
	reconRuns := make(chan *iop.RunRequest, 4)
	toki.AddListenerTyped[*iop.RunRequest](&altNodeClient.Communicator, func(req *iop.RunRequest) {
		altRuns <- proto.Clone(req).(*iop.RunRequest)
	})
	toki.AddListenerTyped[*iop.RunRequest](&reconNodeClient.Communicator, func(req *iop.RunRequest) {
		reconRuns <- proto.Clone(req).(*iop.RunRequest)
	})

	reg := edgenode.NewRegistry()
	altEntry := &edgenode.NodeEntry{NodeID: "node-alt", Alias: "alt", AgentKind: config.AgentKindGenericNode, Client: altEdgeClient, DispatchReady: true}
	reg.Register(altEntry)
	reconEntry := &edgenode.NodeEntry{NodeID: "node-recon", Alias: "recon", AgentKind: config.AgentKindGenericNode, Client: reconEdgeClient, DispatchReady: true}
	reg.Register(reconEntry)
	gen1 := reconEntry.ConnectionGeneration

	svc := New(reg, edgeevents.NewBus())
	store, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "node-alt", Alias: "alt", Token: "alt-token", AgentKind: config.AgentKindGenericNode, Providers: []config.NodeProviderConf{
			{ID: "prov-alt", Type: "openai", Models: []string{"model-1"}, Health: "available", Capacity: 1},
		}},
		{ID: "node-recon", Alias: "recon", Token: "recon-token", AgentKind: config.AgentKindGenericNode, Providers: []config.NodeProviderConf{
			{ID: "prov-recon", Type: "openai", Models: []string{"model-1"}, Health: "available", Capacity: 1},
		}},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{{ID: "model-1", Providers: map[string]string{"prov-alt": "model-1", "prov-recon": "model-1"}}})

	svc.HandleNodeConnect("node-alt", altEntry.ConnectionGeneration)
	svc.HandleNodeConnect("node-recon", gen1)

	ctx := context.Background()

	// Occupy prov-alt with a background run. With both providers idle the tie-break
	// (prov-alt < prov-recon) sends it to prov-alt, leaving prov-recon the only
	// candidate a later waiter can reach.
	if _, err := svc.SubmitRun(ctx, SubmitRunRequest{ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true, Background: true}); err != nil {
		t.Fatalf("submit blocking run: %v", err)
	}
	occupy := recvWire(t, altRuns, "RunRequest(alt)")
	if got, _, _ := providerResourceInflight(svc.queue, "node-alt", "prov-alt"); got != 1 {
		t.Fatalf("prov-alt in-flight = %d after blocking run, want 1", got)
	}

	// Disconnect node-recon: its provider resource is orphaned.
	reg.UnregisterIfClient("node-recon", reconEntry.Client)
	svc.HandleNodeDisconnect("node-recon", gen1, "test-disconnect")
	if _, orphan, ok := providerResourceInflight(svc.queue, "node-recon", "prov-recon"); !ok || !orphan {
		t.Fatalf("prov-recon must be orphaned after disconnect (ok=%v orphan=%v)", ok, orphan)
	}

	// Submit the waiter: alt is full and recon is offline, so it strands in queue.
	type outcome struct {
		res RunResult
		err error
	}
	waiterDone := make(chan outcome, 1)
	go func() {
		res, err := svc.SubmitRun(ctx, SubmitRunRequest{ModelGroupKey: "model-1", Adapter: "openai", Target: "model-1", ProviderPool: true, Background: true})
		waiterDone <- outcome{res: res, err: err}
	}()
	waitForCondition(t, func() bool { return queuedItemCount(svc.queue) == 1 }, "waiter did not enqueue")

	// Reconnect node-recon under a strictly higher generation (pending→ready in the
	// real flow; Register here mints the new generation and makes it the owner).
	reconEntry2 := &edgenode.NodeEntry{NodeID: "node-recon", Alias: "recon", AgentKind: config.AgentKindGenericNode, Client: reconEdgeClient, DispatchReady: true}
	reg.Register(reconEntry2)
	gen2 := reconEntry2.ConnectionGeneration
	if gen2 <= gen1 {
		t.Fatalf("reconnect generation %d must exceed %d", gen2, gen1)
	}

	// Stale activation: the superseded generation must be a no-op — no resource
	// re-activation, no pump, waiter stays queued.
	svc.HandleNodeConnect("node-recon", gen1)
	if _, orphan, ok := providerResourceInflight(svc.queue, "node-recon", "prov-recon"); !ok || !orphan {
		t.Fatalf("stale activation must leave prov-recon orphaned (ok=%v orphan=%v)", ok, orphan)
	}
	if n := queuedItemCount(svc.queue); n != 1 {
		t.Fatalf("stale activation pumped the waiter: queued=%d, want 1", n)
	}
	select {
	case o := <-waiterDone:
		t.Fatalf("stale activation dispatched the waiter: err=%v", o.err)
	default:
	}
	assertNoExtra(t, reconRuns, "RunRequest(recon)")

	// Current activation: the reconnect generation activates the resource and
	// pumps the stranded waiter to it exactly once.
	svc.HandleNodeConnect("node-recon", gen2)
	var waiter outcome
	select {
	case waiter = <-waiterDone:
	case <-time.After(3 * time.Second):
		t.Fatal("current activation did not dispatch the waiter")
	}
	if waiter.err != nil {
		t.Fatalf("waiter dispatch failed: %v", waiter.err)
	}
	dispatch := waiter.res.Dispatch()
	if dispatch.NodeID != "node-recon" || dispatch.ProviderID != "prov-recon" {
		t.Fatalf("waiter dispatch=(node=%q provider=%q), want node-recon/prov-recon", dispatch.NodeID, dispatch.ProviderID)
	}
	reconWire := recvWire(t, reconRuns, "RunRequest(recon)")
	if reconWire.GetRunId() != dispatch.RunID {
		t.Fatalf("recon wire run_id=%q, want %q", reconWire.GetRunId(), dispatch.RunID)
	}
	assertNoExtra(t, reconRuns, "RunRequest(recon)")
	if got, _, _ := providerResourceInflight(svc.queue, "node-recon", "prov-recon"); got != 1 {
		t.Fatalf("prov-recon in-flight = %d after activation, want 1", got)
	}

	// Settle both runs and confirm no counter leaked.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: occupy.GetRunId(), Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch.RunID, Type: "complete"})
	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 }, "runs did not settle")
	assertQueueSettled(t, svc.queue)
}