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.
	svc.HandleNodeDisconnect(nd1NodeID, "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.
	svc.HandleNodeDisconnect("xg-node-1", "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)
	}
}