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

package service_test

import (
	"context"
	"errors"
	"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"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func TestSubmitRunReturnsDispatchMetadata(t *testing.T) {
	dispatch := edgeservice.RunDispatch{
		RunID:         "run-x",
		NodeID:        "node-1",
		NodeLabel:     "alpha",
		ModelGroupKey: "codex-model",
		Adapter:       "cli",
		Target:        "codex",
		SessionID:     "session-a",
		Background:    true,
		TimeoutSec:    30,
	}
	handle := &edgeservice.RunHandle{RunDispatch: dispatch}
	if handle.RunID != dispatch.RunID || handle.NodeLabel != dispatch.NodeLabel {
		t.Fatalf("RunHandle does not expose embedded RunDispatch fields: %+v", handle)
	}
	if handle.ModelGroupKey != "codex-model" {
		t.Errorf("RunHandle ModelGroupKey: got %q want codex-model", handle.ModelGroupKey)
	}
	if handle.Background != true || handle.TimeoutSec != 30 {
		t.Errorf("RunHandle dispatch fields wrong: %+v", handle)
	}
}

// TestSubmitRunModelQueueDispatchesQueuedRunAfterDisconnectToLiveNode verifies
// that after a node disconnect, a queued run is dispatched to a remaining live
// node once that node's slot becomes available.
func TestSubmitRunModelQueueDispatchesQueuedRunAfterDisconnectToLiveNode(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "dc-node-1", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "dc-node-2", Client: edgeClient2})

	catalog := []config.ModelCatalogEntry{
		{
			ID: "dc-group",
			Providers: map[string]string{
				"prov-dc-1": "dc-model",
				"prov-dc-2": "dc-model",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "dc-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-dc-1", Adapter: "mock", Models: []string{"dc-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID: "dc-node-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-dc-2", Adapter: "mock", Models: []string{"dc-model"}, Health: "available", Capacity: 1},
		},
	})

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

	// Submit two background runs to fill both nodes.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "dc-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "dc-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2: %v", err)
	}
	defer res2.Close()

	dispatch1 := res1.Dispatch()
	dispatch2 := res2.Dispatch()
	if dispatch1.NodeID == "" || dispatch2.NodeID == "" || dispatch1.NodeID == dispatch2.NodeID {
		t.Fatalf("expected two different nodes; got %q and %q", dispatch1.NodeID, dispatch2.NodeID)
	}

	// Start run3 in a goroutine — will queue because both nodes are full.
	var (
		res3 edgeservice.RunResult
		err3 error
		wg   sync.WaitGroup
	)
	wg.Add(1)
	go func() {
		defer wg.Done()
		res3, err3 = svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
			ModelGroupKey: "dc-group",
			ProviderPool:  true,
			Background:    true,
		})
	}()

	// Allow run3 to enter the queue.
	time.Sleep(30 * time.Millisecond)

	// Identify which node is nd1 (to disconnect) and nd2 (to receive run3).
	nd1NodeID := dispatch1.NodeID
	nd2RunID := dispatch2.RunID
	nd2NodeID := dispatch2.NodeID

	// nd1 disconnects — its candidate is removed from run3's queue item. The
	// transport drives these hooks directly; the event bus is observability only.
	// Generation 0 fences the node unconditionally (whole-node release).
	svc.HandleNodeDisconnect(nd1NodeID, 0, "disconnected")

	// nd2's run terminates — now nd2 has capacity, run3 dispatches to nd2.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: nd2RunID, Type: "complete"})

	wg.Wait()

	if err3 != nil {
		t.Fatalf("run3 error: %v", err3)
	}
	if res3 == nil {
		t.Fatal("run3: expected non-nil result")
	}
	defer res3.Close()

	if got := res3.Dispatch().NodeID; got != nd2NodeID {
		t.Errorf("run3 dispatched to %q, want %q (the live node; nd1=%q disconnected)", got, nd2NodeID, nd1NodeID)
	}
}

// TestSubmitRunModelQueueContextCancelRemovesQueuedItem verifies that cancelling
// the SubmitRun context removes the item from the queue and returns context.Canceled.
func TestSubmitRunModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "ctx-group",
			Providers: map[string]string{"prov-ctx-1": "ctx-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "ctx-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-ctx-1", Adapter: "mock", Models: []string{"ctx-model"}, Health: "available", Capacity: 1},
		},
	})

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

	// Fill the node capacity with run1.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "ctx-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()

	// Start run2 with a cancellable context (will queue).
	ctx, cancel := context.WithCancel(context.Background())
	var run2Err error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		_, run2Err = svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
			ModelGroupKey: "ctx-group",
			ProviderPool:  true,
			Background:    true,
		})
	}()

	// Give run2 time to enter the queue.
	time.Sleep(30 * time.Millisecond)
	cancel()

	wg.Wait()

	if run2Err == nil {
		t.Fatal("expected error from cancelled context, got nil")
	}
	if !errors.Is(run2Err, context.Canceled) {
		t.Errorf("expected context.Canceled, got: %v", run2Err)
	}
}

// TestSubmitRunModelQueueDispatchesAcrossNodes verifies that concurrent
// SubmitRun calls with the same ModelGroupKey are dispatched across multiple
// nodes when each node has capacity 1.
func TestSubmitRunModelQueueDispatchesAcrossNodes(t *testing.T) {
	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.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)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)

	// Node-side clients consume RunRequest messages without responding.
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "mq-node-1", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "mq-node-2", Client: edgeClient2})

	catalog := []config.ModelCatalogEntry{
		{
			ID: "test-group",
			Providers: map[string]string{
				"prov-mq-1": "served-model",
				"prov-mq-2": "served-model",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:    "mq-node-1",
		Token: "tok-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-mq-1", Adapter: "mock", Models: []string{"served-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID:    "mq-node-2",
		Token: "tok-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-mq-2", Adapter: "mock", Models: []string{"served-model"}, Health: "available", Capacity: 1},
		},
	})

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

	var (
		mu      sync.Mutex
		nodeIDs []string
		errs    []error
	)
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			res, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
				ModelGroupKey: "test-group",
				ProviderPool:  true,
				Background:    true,
			})
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, err)
				return
			}
			nodeIDs = append(nodeIDs, res.Dispatch().NodeID)
			res.Close()
		}()
	}
	wg.Wait()

	for _, err := range errs {
		t.Fatalf("SubmitRun error: %v", err)
	}

	if len(nodeIDs) != 2 {
		t.Fatalf("expected 2 dispatches, got %d", len(nodeIDs))
	}

	usedNodes := map[string]int{}
	for _, id := range nodeIDs {
		usedNodes[id]++
	}
	if usedNodes["mq-node-1"] != 1 || usedNodes["mq-node-2"] != 1 {
		t.Errorf("expected each node used exactly once: %v", usedNodes)
	}
}

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

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "local-group",
			Providers: map[string]string{"prov-local-1": "local-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-1",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-local-1", Adapter: "ollama-local", Models: []string{"local-model"}, Health: "available", Capacity: 2, MaxQueue: 3},
		},
	})

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

	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2 error: %v", err)
	}
	defer res2.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err3 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err3 == nil {
		t.Fatal("expected third run to block and timeout (capacity=2 exceeded)")
	}
	if !errors.Is(err3, context.DeadlineExceeded) {
		t.Errorf("expected deadline exceeded, got: %v", err3)
	}
}

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

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "fallback-group",
			Providers: map[string]string{"prov-fallback-1": "fallback-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-fallback",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			// capacity=2, maxQueue=3 come from provider config (provider-pool always uses provider policy).
			{ID: "prov-fallback-1", Adapter: "ollama-local", Models: []string{"fallback-model"}, Health: "available", Capacity: 2, MaxQueue: 3},
		},
	})

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

	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2 error: %v", err)
	}
	defer res2.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err3 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err3 == nil {
		t.Fatal("expected third run to block and timeout (capacity=2 exceeded)")
	}
	if !errors.Is(err3, context.DeadlineExceeded) {
		t.Errorf("expected deadline exceeded, got: %v", err3)
	}
}

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

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "timeout-group",
			Providers: map[string]string{"prov-timeout-1": "timeout-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-timeout",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			// QueueTimeoutMS=10 comes from provider config; provider-pool ignores request-level QueueTimeoutMS.
			{ID: "prov-timeout-1", Adapter: "ollama-local", Models: []string{"timeout-model"}, Health: "available", Capacity: 1, MaxQueue: 5, QueueTimeoutMS: 10},
		},
	})

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

	// Set the canonical provider-pool policy with QueueTimeoutMS=10.
	// The root policy owns the timeout; provider-level QueueTimeoutMS is no
	// longer consulted for provider-pool admission.
	svc.SetRuntimeConfig(store, catalog, edgeservice.NewGroupPolicy(16, 10*time.Millisecond))

	// run 1: fills capacity; root policy sets QueueTimeoutMS=10.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "timeout-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	// run 2: enters queue; times out after ~10ms per provider policy.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	start := time.Now()
	_, err2 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "timeout-group",
		ProviderPool:  true,
		Background:    true,
	})
	elapsed := time.Since(start)

	if err2 == nil {
		t.Fatal("expected second run to fail with queue timeout")
	}
	if !strings.Contains(err2.Error(), "queue timeout") {
		t.Errorf("expected queue timeout error, got: %v", err2)
	}
	if elapsed > 1*time.Second {
		t.Errorf("expected timeout to happen quickly, took: %v", elapsed)
	}
}

// TestSubmitRunLegacyModelGroupKeyGoesDirectNotQueued verifies that a request
// with ModelGroupKey but ProviderPool=false bypasses the queue and goes directly
// to the single registered node. This proves the queue admission gate on
// req.ProviderPool: without ProviderPool=true, ModelGroupKey is decorative only
// and the service falls back to direct dispatch without touching the model catalog.
func TestSubmitRunLegacyModelGroupKeyGoesDirectNotQueued(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	// No catalog set — if ProviderPool were true, resolveProviderPoolCandidates would
	// return "not found in catalog". Since ProviderPool=false, direct dispatch succeeds.

	res, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "legacy-group",
		Adapter:       "cli",
		Target:        "codex",
		Background:    true,
		// ProviderPool: false (default)
	})
	if err != nil {
		t.Fatalf("expected direct dispatch to succeed, got: %v", err)
	}
	defer res.Close()

	if res.Dispatch().NodeID != "leg-node-1" {
		t.Errorf("dispatch node: got %q, want leg-node-1", res.Dispatch().NodeID)
	}
}

// TestGlobalPumpNodeExclusionFallsBackAcrossGroups verifies the exclusion +
// cross-group path end to end: a queued request whose preferred node disconnects
// must fall back to the remaining live provider, and the release that frees that
// provider happens in a *different* model group. Before the global pump, the
// waiter's group saw no event of its own and the request sat until timeout.
func TestGlobalPumpNodeExclusionFallsBackAcrossGroups(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "xg-node-1", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "xg-node-2", Client: edgeClient2})

	// Two model groups share the provider on node 2; only group A can also use
	// the provider on node 1.
	catalog := []config.ModelCatalogEntry{
		{
			ID: "xg-group-a",
			Providers: map[string]string{
				"prov-xg-1": "xg-model-a",
				"prov-xg-2": "xg-model-a",
			},
		},
		{
			ID: "xg-group-b",
			Providers: map[string]string{
				"prov-xg-2": "xg-model-b",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "xg-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-xg-1", Adapter: "mock", Models: []string{"xg-model-a"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID: "xg-node-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-xg-2", Adapter: "mock", Models: []string{"xg-model-a", "xg-model-b"}, Health: "available", Capacity: 1},
		},
	})

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

	// Fill node 1 via group A and node 2 via group B.
	resA, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "xg-group-a",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("group A run: %v", err)
	}
	defer resA.Close()
	resB, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "xg-group-b",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("group B run: %v", err)
	}
	defer resB.Close()

	dispatchA := resA.Dispatch()
	dispatchB := resB.Dispatch()
	if dispatchA.NodeID != "xg-node-1" {
		t.Fatalf("group A run landed on %q, want xg-node-1", dispatchA.NodeID)
	}
	if dispatchB.NodeID != "xg-node-2" {
		t.Fatalf("group B run landed on %q, want xg-node-2", dispatchB.NodeID)
	}

	// A second group A request queues: both providers it can use are busy.
	var (
		queued edgeservice.RunResult
		errQ   error
		wg     sync.WaitGroup
	)
	wg.Add(1)
	go func() {
		defer wg.Done()
		queued, errQ = svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
			ModelGroupKey: "xg-group-a",
			ProviderPool:  true,
			Background:    true,
		})
	}()
	time.Sleep(30 * time.Millisecond)

	// Node 1 drops out, removing prov-xg-1 from the waiter's candidates. The only
	// remaining path is prov-xg-2, and the run holding it belongs to group B.
	// Generation 0 fences the node unconditionally (whole-node release).
	svc.HandleNodeDisconnect("xg-node-1", 0, "disconnected")

	// Terminating the group B run must wake the group A waiter.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatchB.RunID, Type: "complete"})

	wg.Wait()
	if errQ != nil {
		t.Fatalf("queued group A run error: %v", errQ)
	}
	if queued == nil {
		t.Fatal("queued group A run: expected non-nil result")
	}
	defer queued.Close()

	if got := queued.Dispatch().NodeID; got != "xg-node-2" {
		t.Errorf("queued run dispatched to %q, want xg-node-2 (the surviving node freed by another model group)", got)
	}
}

// TestSubmitRunModelQueueSurvivingProviderFallbackAfterDisconnect verifies that
// when a queued run's preferred provider disconnects, the authoritative
// resolve-and-pump pass rebuilds the candidate universe through the live
// resolver, discovers a surviving provider on a different node, and dispatches
// the waiter to that surviving provider — all without waiting for the queue
// timeout. This is the integration path for the REFACTOR-1 change.
func TestSubmitRunModelQueueSurvivingProviderFallbackAfterDisconnect(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "fb-node-1", Client: edgeClient1})

	catalog := []config.ModelCatalogEntry{
		{
			ID: "fb-group",
			Providers: map[string]string{
				"prov-fb-1": "fb-model",
				"prov-fb-2": "fb-model",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "fb-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-fb-1", Adapter: "mock", Models: []string{"fb-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID: "fb-node-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-fb-2", Adapter: "mock", Models: []string{"fb-model"}, Health: "available", Capacity: 1},
		},
	})

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

	// Fill fb-node-1.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "fb-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()
	if dispatch := res1.Dispatch(); dispatch.NodeID != "fb-node-1" || dispatch.ProviderID != "prov-fb-1" {
		t.Fatalf("run1 dispatch = (%q, %q), want (fb-node-1, prov-fb-1)", dispatch.NodeID, dispatch.ProviderID)
	}

	// Start run2 while fb-node-1 is the only live provider. Its capacity is held
	// by run1, so run2 must enter the provider-pool queue.
	type submitOutcome struct {
		result edgeservice.RunResult
		err    error
	}
	run2Ctx, cancelRun2 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancelRun2()
	run2Done := make(chan submitOutcome, 1)
	go func() {
		result, submitErr := svc.SubmitRun(run2Ctx, edgeservice.SubmitRunRequest{
			ModelGroupKey: "fb-group",
			ProviderPool:  true,
			Background:    true,
		})
		run2Done <- submitOutcome{result: result, err: submitErr}
	}()

	// Synchronize on observable queue state instead of assuming goroutine
	// scheduling with a fixed sleep. The waiter must be pressure on the only live
	// provider and must not have returned before any survivor is connected.
	waitForProviderSnapshotCounts(t, svc, "fb-node-1", "prov-fb-1", 1, 1)
	select {
	case outcome := <-run2Done:
		if outcome.result != nil {
			outcome.result.Close()
		}
		t.Fatalf("run2 completed before survivor registration: %v", outcome.err)
	default:
	}

	// Expose the configured survivor to the live resolver. Registry registration
	// alone has no queue-pump side effect; observing its candidate pressure proves
	// the waiter is still queued with the new live universe before disconnect.
	reg.Register(&edgenode.NodeEntry{NodeID: "fb-node-2", Client: edgeClient2})
	waitForProviderSnapshotCounts(t, svc, "fb-node-2", "prov-fb-2", 0, 1)
	select {
	case outcome := <-run2Done:
		if outcome.result != nil {
			outcome.result.Close()
		}
		t.Fatalf("run2 completed from survivor registration without disconnect: %v", outcome.err)
	default:
	}

	// Mirror the transport's authoritative ordering: current-client ownership is
	// checked and removed first, then that exact generation drives settlement.
	disconnectedGeneration, ok := reg.UnregisterIfClient("fb-node-1", edgeClient1)
	if !ok || disconnectedGeneration == 0 {
		t.Fatalf("authoritative unregister = (%d, %t), want non-zero generation and true", disconnectedGeneration, ok)
	}
	svc.HandleNodeDisconnect("fb-node-1", disconnectedGeneration, "disconnected")

	// The disconnect transition is the only queue wake-up. Completion is bounded
	// and the single submit goroutine must publish exactly one outcome.
	var outcome submitOutcome
	select {
	case outcome = <-run2Done:
	case <-time.After(time.Second):
		t.Fatal("run2 did not complete within 1s of authoritative disconnect")
	}
	select {
	case duplicate := <-run2Done:
		if duplicate.result != nil {
			duplicate.result.Close()
		}
		t.Fatalf("run2 completed more than once: %v", duplicate.err)
	default:
	}
	if outcome.err != nil {
		t.Fatalf("run2 error: %v", outcome.err)
	}
	if outcome.result == nil {
		t.Fatal("run2: expected non-nil result")
	}
	res2 := outcome.result
	defer res2.Close()

	dispatch2 := res2.Dispatch()
	if dispatch2.NodeID != "fb-node-2" || dispatch2.ProviderID != "prov-fb-2" {
		t.Errorf("run2 dispatch = (%q, %q), want (fb-node-2, prov-fb-2)", dispatch2.NodeID, dispatch2.ProviderID)
	}
	if dispatch2.QueueReason != "capacity_full" {
		t.Errorf("run2 queue reason = %q, want capacity_full", dispatch2.QueueReason)
	}

	// The disconnected holder lease and queue item are already gone, while the
	// fallback lease is the sole remaining in-flight reservation.
	provider1 := waitForProviderSnapshotCounts(t, svc, "fb-node-1", "prov-fb-1", 0, 0)
	provider2 := waitForProviderSnapshotCounts(t, svc, "fb-node-2", "prov-fb-2", 1, 0)
	if provider1.GetLongInFlight() != 0 || provider1.GetLongQueued() != 0 {
		t.Errorf("fb-node-1 long counters after disconnect = (%d, %d), want (0, 0)", provider1.GetLongInFlight(), provider1.GetLongQueued())
	}
	if provider2.GetLongInFlight() != 0 || provider2.GetLongQueued() != 0 {
		t.Errorf("fb-node-2 long counters after fallback = (%d, %d), want (0, 0)", provider2.GetLongInFlight(), provider2.GetLongQueued())
	}

	// A terminal event returns the surviving provider lease. Repeating the same
	// event is an idempotence check: all provider and queue counters stay at zero.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch2.RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch2.RunID, Type: "complete"})
	provider1 = waitForProviderSnapshotCounts(t, svc, "fb-node-1", "prov-fb-1", 0, 0)
	provider2 = waitForProviderSnapshotCounts(t, svc, "fb-node-2", "prov-fb-2", 0, 0)
	if provider1.GetLongInFlight() != 0 || provider1.GetLongQueued() != 0 ||
		provider2.GetLongInFlight() != 0 || provider2.GetLongQueued() != 0 {
		t.Fatalf("long counters after cleanup: provider1=(%d, %d), provider2=(%d, %d), want all zero",
			provider1.GetLongInFlight(), provider1.GetLongQueued(), provider2.GetLongInFlight(), provider2.GetLongQueued())
	}
}

func waitForProviderSnapshotCounts(
	t *testing.T,
	svc *edgeservice.Service,
	nodeID string,
	providerID string,
	wantInFlight int32,
	wantQueued int32,
) *iop.ProviderSnapshot {
	t.Helper()

	deadline := time.Now().Add(time.Second)
	var (
		lastInFlight int32
		lastQueued   int32
		found        bool
	)
	for {
		for _, snapshot := range svc.ListNodeSnapshots() {
			if snapshot.NodeID != nodeID {
				continue
			}
			for _, provider := range snapshot.ProviderSnapshots {
				if provider.GetId() != providerID {
					continue
				}
				found = true
				lastInFlight = provider.GetInFlight()
				lastQueued = provider.GetQueued()
				if lastInFlight == wantInFlight && lastQueued == wantQueued {
					return provider
				}
			}
		}
		if !time.Now().Before(deadline) {
			t.Fatalf("provider %s/%s snapshot counts = (%d, %d), want (%d, %d), found=%t",
				nodeID, providerID, lastInFlight, lastQueued, wantInFlight, wantQueued, found)
		}
		time.Sleep(time.Millisecond)
	}
}

// TestSubmitRunModelQueueLastProviderDisconnectReturnsUnavailable verifies that
// when ALL live candidates for a queued run disappear (last provider
// disconnect), the waiter receives a terminal unavailable error immediately
// instead of waiting for the queue timeout. This is the critical
// correctness condition for the REFACTOR-1 authoritative settlement pass.
func TestSubmitRunModelQueueLastProviderDisconnectReturnsUnavailable(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

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

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "lp-group",
			Providers: map[string]string{"prov-lp-1": "lp-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "lp-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-lp-1", Adapter: "mock", Models: []string{"lp-model"}, Health: "available", Capacity: 1},
		},
	})

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

	// Fill the only provider.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "lp-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()

	// Start run2 — will queue because capacity is full.
	var (
		res2 edgeservice.RunResult
		err2 error
		wg   sync.WaitGroup
	)
	wg.Add(1)
	go func() {
		defer wg.Done()
		res2, err2 = svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
			ModelGroupKey: "lp-group",
			ProviderPool:  true,
			Background:    true,
		})
	}()

	time.Sleep(30 * time.Millisecond)

	// Disconnect the last provider — all candidates disappear.
	svc.HandleNodeDisconnect("lp-node-1", 0, "last provider disconnect")

	wg.Wait()

	// The waiter must receive a terminal error (not a successful dispatch).
	if err2 == nil {
		t.Fatal("expected terminal unavailable error for last-provider disconnect, got nil")
	}
	if res2 != nil {
		res2.Close()
		t.Fatalf("expected nil result for last-provider disconnect, got %+v", res2.Dispatch())
	}

	// Verify all accounting is zero via snapshot.
	snaps := svc.ListNodeSnapshots()
	for _, s := range snaps {
		for _, ps := range s.ProviderSnapshots {
			if ps.GetInFlight() != 0 {
				t.Errorf("provider %s on %s: inflight after last disconnect = %d, want 0",
					ps.GetId(), s.NodeID, ps.GetInFlight())
			}
		}
	}
}

// TestReconnectRebuildsDisconnectedCandidateUniverse verifies the S15
// reconnect-candidate-recovery contract through the full SubmitRun path: when a
// waiter's target provider disconnected (dropping it from the live candidate
// universe) and the only alternate is full, the accepted reconnect alone rebuilds
// the candidate universe through the live resolver and dispatches the waiter back
// to the reconnected provider — without any new request, config refresh, or lease
// release. The transport's authoritative ordering (UnregisterIfClient →
// HandleNodeDisconnect on close; RegisterIfAbsent generation → HandleNodeConnect
// on reconnect) is mirrored directly.
func TestReconnectRebuildsDisconnectedCandidateUniverse(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	// node-alt is the full alternate; node-recon disconnects and reconnects. A
	// second pipe stands in for the reconnect's fresh connection.
	edgeAlt, nodeAlt := net.Pipe()
	edgeRecon1, nodeRecon1 := net.Pipe()
	edgeRecon2, nodeRecon2 := net.Pipe()
	defer edgeAlt.Close()
	defer nodeAlt.Close()
	defer edgeRecon1.Close()
	defer nodeRecon1.Close()
	defer edgeRecon2.Close()
	defer nodeRecon2.Close()

	edgeAltClient := toki.NewTcpClient(edgeAlt, 0, 0, parserMap)
	edgeReconClient1 := toki.NewTcpClient(edgeRecon1, 0, 0, parserMap)
	edgeReconClient2 := toki.NewTcpClient(edgeRecon2, 0, 0, parserMap)
	nodeAltClient := toki.NewTcpClient(nodeAlt, 0, 0, parserMap)
	nodeReconClient1 := toki.NewTcpClient(nodeRecon1, 0, 0, parserMap)
	nodeReconClient2 := toki.NewTcpClient(nodeRecon2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeAltClient.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeReconClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeReconClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-alt", Client: edgeAltClient})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-recon", Client: edgeReconClient1})

	catalog := []config.ModelCatalogEntry{{
		ID: "rc-group",
		Providers: map[string]string{
			"prov-alt":   "rc-model",
			"prov-recon": "rc-model",
		},
	}}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-alt",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"}},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-alt", Adapter: "mock", Models: []string{"rc-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID: "node-recon",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"}},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-recon", Adapter: "mock", Models: []string{"rc-model"}, Health: "available", Capacity: 1},
		},
	})

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

	// run1 fills the alternate provider (prov-alt sorts before prov-recon, so the
	// scheduler picks it deterministically for the first admission).
	run1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "rc-group", ProviderPool: true, Background: true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer run1.Close()
	if dispatch := run1.Dispatch(); dispatch.NodeID != "node-alt" || dispatch.ProviderID != "prov-alt" {
		t.Fatalf("run1 dispatch = (%q, %q), want (node-alt, prov-alt)", dispatch.NodeID, dispatch.ProviderID)
	}

	// Disconnect node-recon (idle) using the transport's authoritative ordering.
	// This drops prov-recon from the live candidate universe and orphans it.
	disconnectedGeneration, ok := reg.UnregisterIfClient("node-recon", edgeReconClient1)
	if !ok || disconnectedGeneration == 0 {
		t.Fatalf("authoritative unregister = (%d, %t), want non-zero generation and true", disconnectedGeneration, ok)
	}
	svc.HandleNodeDisconnect("node-recon", disconnectedGeneration, "disconnected")

	// run2 now has only prov-alt as a live candidate, and it is full — the request
	// must enter the queue rather than dispatch.
	type submitOutcome struct {
		result edgeservice.RunResult
		err    error
	}
	run2Ctx, cancelRun2 := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancelRun2()
	run2Done := make(chan submitOutcome, 1)
	go func() {
		result, submitErr := svc.SubmitRun(run2Ctx, edgeservice.SubmitRunRequest{
			ModelGroupKey: "rc-group", ProviderPool: true, Background: true,
		})
		run2Done <- submitOutcome{result: result, err: submitErr}
	}()

	// Confirm the waiter is queued against the only live provider before reconnect.
	waitForProviderSnapshotCounts(t, svc, "node-alt", "prov-alt", 1, 1)
	select {
	case outcome := <-run2Done:
		if outcome.result != nil {
			outcome.result.Close()
		}
		t.Fatalf("run2 resolved before reconnect: %v", outcome.err)
	default:
	}

	// Reconnect node-recon on a fresh connection: a strictly higher generation is
	// minted, and the accepted-connect hook alone rebuilds the candidate universe
	// and pumps the waiter. No config refresh, new request, or lease release.
	reconnectEntry := &edgenode.NodeEntry{NodeID: "node-recon", Client: edgeReconClient2}
	reg.Register(reconnectEntry)
	if reconnectEntry.ConnectionGeneration <= disconnectedGeneration {
		t.Fatalf("reconnect generation=%d, want > disconnected generation=%d", reconnectEntry.ConnectionGeneration, disconnectedGeneration)
	}
	svc.HandleNodeConnect("node-recon", reconnectEntry.ConnectionGeneration)

	var outcome submitOutcome
	select {
	case outcome = <-run2Done:
	case <-time.After(time.Second):
		t.Fatal("run2 did not dispatch within 1s of the accepted reconnect")
	}
	if outcome.err != nil {
		t.Fatalf("run2 error: %v", outcome.err)
	}
	if outcome.result == nil {
		t.Fatal("run2: expected non-nil result")
	}
	run2 := outcome.result
	defer run2.Close()

	dispatch2 := run2.Dispatch()
	if dispatch2.NodeID != "node-recon" || dispatch2.ProviderID != "prov-recon" {
		t.Fatalf("run2 dispatch = (%q, %q), want (node-recon, prov-recon)", dispatch2.NodeID, dispatch2.ProviderID)
	}
	if dispatch2.QueueReason != "capacity_full" {
		t.Errorf("run2 queue reason = %q, want capacity_full", dispatch2.QueueReason)
	}

	// The reconnected provider now carries the waiter's lease; the alternate still
	// holds run1 (its lease was never released), proving the reconnect drove the pump.
	waitForProviderSnapshotCounts(t, svc, "node-recon", "prov-recon", 1, 0)
	waitForProviderSnapshotCounts(t, svc, "node-alt", "prov-alt", 1, 0)

	// Terminal events drain both providers; repeating is an idempotence check.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: run1.Dispatch().RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch2.RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch2.RunID, Type: "complete"})
	waitForProviderSnapshotCounts(t, svc, "node-alt", "prov-alt", 0, 0)
	waitForProviderSnapshotCounts(t, svc, "node-recon", "prov-recon", 0, 0)
}