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

package service

import (
	"context"
	"errors"
	"net"
	"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"
)

// waitForQueueLen polls until the group queue reaches wantLen or times out.
func waitForQueueLen(t *testing.T, m *modelQueueManager, groupKey string, wantLen int) {
	t.Helper()
	deadline := time.Now().Add(200 * time.Millisecond)
	for time.Now().Before(deadline) {
		m.mu.Lock()
		g, ok := m.groups[groupKey]
		got := 0
		if ok {
			got = len(g.queue)
		}
		m.mu.Unlock()
		if got >= wantLen {
			return
		}
		time.Sleep(1 * time.Millisecond)
	}
	t.Fatalf("timeout: queue for %q did not reach length %d", groupKey, wantLen)
}

// TestModelQueueFIFOOrdering verifies that queued items are dispatched in FIFO
// order when a slot becomes available.
func TestModelQueueFIFOOrdering(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-q1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-q1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(store)

	// Fill the only capacity slot.
	first, err := m.admit(context.Background(), "g-fifo", "", "", cands, defPolicy)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Manually insert two items in known order so we can verify FIFO.
	item1 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	item2 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-fifo", groupPolicy{})
	g.queue = append(g.queue, item1, item2)
	m.mu.Unlock()

	// Release one slot — item1 (head) must be dispatched first.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "")
	m.mu.Unlock()

	select {
	case res := <-item1.waitCh:
		if res.err != nil {
			t.Fatalf("item1 expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.entry.NodeID != "node-q1" {
			t.Fatalf("item1: unexpected candidate %v", res.candidate)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item1 was not dispatched after slot release")
	}

	// item2 must still be waiting (no slot available yet).
	select {
	case <-item2.waitCh:
		t.Fatal("item2 should not have been dispatched yet")
	default:
	}

	// Release the slot item1 holds so item2 gets dispatched.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "")
	m.mu.Unlock()

	select {
	case res := <-item2.waitCh:
		if res.err != nil {
			t.Fatalf("item2 expected dispatch, got error: %v", res.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item2 was not dispatched after second slot release")
	}
}

// TestModelQueueOverflow verifies that admit rejects requests when the queue
// is at max_queue capacity.
func TestModelQueueOverflow(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ov1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-ov1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill the node's capacity and set max_queue=1.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-overflow", groupPolicy{})
	g.policy.maxQueue = 1
	g.inflight["node-ov1"] = 1
	m.mu.Unlock()

	// Queue one item — should succeed.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	m.groups["g-overflow"].queue = append(m.groups["g-overflow"].queue, item)
	m.mu.Unlock()

	// Second admit should fail immediately with errQueueFull.
	_, err := m.admit(context.Background(), "g-overflow", "", "", cands, groupPolicy{})
	if !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull, got: %v", err)
	}
}

// TestModelQueueTimeout verifies that queued items receive a timeout error
// when no slot becomes available before the deadline.
func TestModelQueueTimeout(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-to1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-to1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill capacity and set a very short queue timeout.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-timeout", groupPolicy{})
	g.policy.queueTimeout = 20 * time.Millisecond
	g.inflight["node-to1"] = 1
	m.mu.Unlock()

	start := time.Now()
	_, err := m.admit(context.Background(), "g-timeout", "", "", cands, groupPolicy{})
	elapsed := time.Since(start)

	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected errQueueTimeout, got: %v", err)
	}
	if elapsed < 15*time.Millisecond {
		t.Fatalf("timed out too fast: %v", elapsed)
	}
}

// TestModelQueueTerminalReleaseDispatchesNext verifies that a terminal run event
// releases the in-flight slot and dispatches the next queued item.
func TestModelQueueTerminalReleaseDispatchesNext(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-tr1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-tr1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill capacity and record inflight.
	selected, err := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	m.trackInflight("g-tr", "run-tr-001", selected.entry.NodeID, selected.providerID)

	// Queue a second item in a goroutine.
	resultCh := make(chan admitResult, 1)
	go func() {
		n, e := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
		resultCh <- admitResult{candidate: n, err: e}
	}()

	waitForQueueLen(t, m, "g-tr", 1)

	// Terminal event for the first run — should unblock the queued admit.
	for _, termType := range []string{"complete", "error", "cancelled"} {
		t.Run(termType, func(t *testing.T) {
			// Reset state for each sub-test.
			m.mu.Lock()
			g := m.getOrCreateGroupLocked("g-tr-"+termType, groupPolicy{})
			g.inflight["node-tr1"] = 1
			item := &queueItem{
				candidates: cands,
				waitCh:     make(chan admitResult, 1),
				deadline:   time.Now().Add(2 * time.Second),
			}
			g.queue = []*queueItem{item}
			runID := "run-tr-" + termType
			m.inflightByRun[runID] = inflightRec{groupKey: "g-tr-" + termType, nodeID: "node-tr1"}
			m.mu.Unlock()

			bus.PublishRun(&iop.RunEvent{RunId: runID, Type: termType})

			select {
			case res := <-item.waitCh:
				if res.err != nil {
					t.Fatalf("expected dispatch, got error: %v", res.err)
				}
				if res.candidate == nil {
					t.Fatal("expected non-nil candidate")
				}
			case <-time.After(500 * time.Millisecond):
				t.Fatalf("timeout: queued item not dispatched after %q terminal event", termType)
			}
		})
	}

	// Cleanup the goroutine from the outer admit.
	m.mu.Lock()
	if g, ok := m.groups["g-tr"]; ok && len(g.queue) > 0 {
		g.inflight["node-tr1"] = 1
		m.tryDispatchLocked(g)
	}
	m.mu.Unlock()
	select {
	case <-resultCh:
	case <-time.After(100 * time.Millisecond):
	}
}

// TestModelQueueNodeDisconnectReleasesInflight verifies that a node disconnect
// event removes the node from queued item candidate lists, and that subsequent
// dispatch goes to a remaining live candidate (not the disconnected node).
func TestModelQueueNodeDisconnectReleasesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd2",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry1 := &edgenode.NodeEntry{NodeID: "node-nd1"}
	entry2 := &edgenode.NodeEntry{NodeID: "node-nd2"}
	// Queued item that can use either node (both at capacity 1).
	bothCands := []candidateNode{
		{entry: entry1, capacity: 1},
		{entry: entry2, capacity: 1},
	}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill node-nd1 and node-nd2 capacities and record inflight.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-nd", groupPolicy{})
	g.inflight["node-nd1"] = 1
	g.inflight["node-nd2"] = 1
	m.inflightByRun["run-nd-x"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd1"}
	m.inflightByRun["run-nd-y"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd2"}
	// Queue an item that can use either node.
	item := &queueItem{
		candidates: bothCands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	g.queue = []*queueItem{item}
	m.mu.Unlock()

	// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
	// node-nd2 is still full, so no dispatch yet.
	bus.PublishNode(&iop.EdgeNodeEvent{
		NodeId: "node-nd1",
		Type:   "node.disconnected",
	})

	// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
	bus.PublishRun(&iop.RunEvent{RunId: "run-nd-y", Type: "complete"})

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after disconnect+terminal, got error: %v", res.err)
		}
		if res.candidate == nil {
			t.Fatal("expected non-nil candidate after disconnect release")
		}
		// Must be nd2 — nd1 was removed from candidates on disconnect.
		if res.candidate.entry.NodeID != "node-nd2" {
			t.Errorf("expected dispatch to node-nd2, got %q (nd1 was disconnected)", res.candidate.entry.NodeID)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: queued item not dispatched after node disconnect and terminal event")
	}

	// The inflight record for run-nd-x should be gone (removed at disconnect).
	m.mu.Lock()
	_, stillInFlight := m.inflightByRun["run-nd-x"]
	m.mu.Unlock()
	if stillInFlight {
		t.Error("run-nd-x should be removed from inflightByRun after node disconnect")
	}
}

// TestModelQueueUsesProviderCapacity verifies that the capacity supplied in the
// candidateNode (derived from adapter config) overrides the default of 1.
func TestModelQueueUsesProviderCapacity(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-pc1"}
	// Provider capacity = 2 (two concurrent slots on this node).
	cands := []candidateNode{{entry: entry, capacity: 2}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// First admit: inflight=0 < cap=2 → dispatched immediately.
	n1, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
	if err != nil || n1 == nil {
		t.Fatalf("first admit: %v", err)
	}

	// Second admit: inflight=1 < cap=2 → still dispatched (not queued).
	n2, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
	if err != nil || n2 == nil {
		t.Fatalf("second admit (capacity=2 should allow): %v", err)
	}

	// Third request must queue because inflight=2 == cap=2.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(200 * time.Millisecond),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-pc", groupPolicy{})
	g.queue = append(g.queue, item)
	m.mu.Unlock()

	// Releasing one slot should dispatch the queued item.
	m.releaseSlot("g-pc", "node-pc1")

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after slot release, got: %v", res.err)
		}
		if res.candidate == nil || res.candidate.entry.NodeID != "node-pc1" {
			t.Fatalf("unexpected candidate: %v", res.candidate)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: item not dispatched after slot release")
	}
}

// TestRefreshProviderQueuePolicyUpdatesExistingGroup verifies that a tightened
// queue policy (e.g. provider max_queue lowered via config refresh) propagates
// to an already-created group on the next admission, while in-flight slots are
// preserved.
func TestRefreshProviderQueuePolicyUpdatesExistingGroup(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-rp"}
	cands := []candidateNode{{entry: entry, capacity: 2}}
	m := newModelQueueManager(nil)

	// Create the group with a generous policy and take one slot.
	initial := groupPolicy{maxQueue: 8, queueTimeout: 5 * time.Second}
	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, initial); err != nil {
		t.Fatalf("first admit: %v", err)
	}
	m.mu.Lock()
	if got := m.groups["g-rp"].policy.maxQueue; got != 8 {
		m.mu.Unlock()
		t.Fatalf("group created with maxQueue=%d, want 8", got)
	}
	m.mu.Unlock()

	// Config refresh tightens the policy to maxQueue=1; the recomputed admit
	// carries it and must update the existing group's policy. capacity=2 means
	// this second admit is dispatched (not queued).
	refreshed := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}
	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); err != nil {
		t.Fatalf("second admit: %v", err)
	}

	m.mu.Lock()
	g := m.groups["g-rp"]
	if g.policy.maxQueue != 1 {
		m.mu.Unlock()
		t.Fatalf("existing group policy not refreshed: maxQueue=%d, want 1", g.policy.maxQueue)
	}
	// Both slots now in-flight (cap=2). Queue one item to reach the refreshed
	// maxQueue=1 so the next admission must overflow.
	g.queue = append(g.queue, &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(5 * time.Second),
	})
	m.mu.Unlock()

	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull under refreshed maxQueue=1, got: %v", err)
	}
}

// TestModelQueueUsesProviderQueuePolicy verifies that max_queue and
// queue_timeout from the policy parameter are respected.
func TestModelQueueUsesProviderQueuePolicy(t *testing.T) {
	t.Run("maxQueue enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq1"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: maxQueue=1.
		policy := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}

		m := newModelQueueManager(nil)

		// Fill capacity and queue one item manually.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-max", policy)
		g.inflight["node-pq1"] = 1
		g.queue = []*queueItem{{
			candidates: cands,
			waitCh:     make(chan admitResult, 1),
			deadline:   time.Now().Add(5 * time.Second),
		}}
		m.mu.Unlock()

		// Second admit should fail: queue already at maxQueue=1.
		_, err := m.admit(context.Background(), "g-pq-max", "", "", cands, policy)
		if !errors.Is(err, errQueueFull) {
			t.Fatalf("expected errQueueFull, got: %v", err)
		}
	})

	t.Run("queueTimeout enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq2"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: very short timeout.
		policy := groupPolicy{maxQueue: 16, queueTimeout: 20 * time.Millisecond}

		m := newModelQueueManager(nil)

		// Fill capacity.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-to", policy)
		g.inflight["node-pq2"] = 1
		m.mu.Unlock()

		start := time.Now()
		_, err := m.admit(context.Background(), "g-pq-to", "", "", cands, policy)
		elapsed := time.Since(start)

		if !errors.Is(err, errQueueTimeout) {
			t.Fatalf("expected errQueueTimeout, got: %v", err)
		}
		if elapsed < 15*time.Millisecond {
			t.Fatalf("timed out too fast: %v", elapsed)
		}
	})
}

// TestModelQueueProviderLoadRatioSelection verifies that admit selects the
// candidate with the lowest in_flight/capacity ratio, not simply the first one.
func TestModelQueueProviderLoadRatioSelection(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-lr-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-lr-b"}
	// A: capacity=4, B: capacity=2.
	cands := []candidateNode{
		{entry: entryA, capacity: 4, providerID: "prov-a"},
		{entry: entryB, capacity: 2, providerID: "prov-b"},
	}
	m := newModelQueueManager(nil)

	// Put A at inflight=2 (ratio=0.5) and B at inflight=0 (ratio=0.0).
	// Uses provider-aware slot keys per REVIEW_API-2.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-lr", groupPolicy{})
	g.inflight["node-lr-a:prov-a"] = 2
	m.mu.Unlock()

	// Admit should pick B (ratio=0.0 < 0.5).
	sel, err := m.admit(context.Background(), "g-lr", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil || sel.entry.NodeID != "node-lr-b" {
		t.Fatalf("expected node-lr-b (lower ratio), got %v", sel)
	}

	// B is now at inflight=1 (ratio=0.5). A is still at inflight=2 (ratio=0.5).
	// With equal ratios, tie-break by providerID: "prov-a" < "prov-b" → A wins.
	sel2, err := m.admit(context.Background(), "g-lr", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit2: %v", err)
	}
	if sel2 == nil || sel2.entry.NodeID != "node-lr-a" {
		t.Fatalf("expected node-lr-a (tie-break by providerID), got %v", sel2)
	}
}

// TestModelQueueProviderServedTargetRewrite verifies that the selected
// candidateNode carries its servedTarget so callers can rewrite req.Target.
func TestModelQueueProviderServedTargetRewrite(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-tr-rewrite"}
	cands := []candidateNode{{
		entry:        entry,
		capacity:     2,
		providerID:   "prov-vllm",
		servedTarget: "qwen3-72b-instruct",
	}}
	m := newModelQueueManager(nil)

	sel, err := m.admit(context.Background(), "g-tr-rewrite", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil {
		t.Fatal("expected non-nil candidate")
	}
	if sel.servedTarget != "qwen3-72b-instruct" {
		t.Errorf("servedTarget: got %q, want %q", sel.servedTarget, "qwen3-72b-instruct")
	}
	if sel.providerID != "prov-vllm" {
		t.Errorf("providerID: got %q, want %q", sel.providerID, "prov-vllm")
	}
}

// TestProviderStatusInflightTracking verifies that getStatsForProviderLocked
// correctly counts in-flight runs keyed by (nodeID, providerID).
func TestProviderStatusInflightTracking(t *testing.T) {
	m := newModelQueueManager(nil)

	m.mu.Lock()
	m.inflightByRun["run-p1"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
	m.inflightByRun["run-p2"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
	m.inflightByRun["run-p3"] = inflightRec{groupKey: "g-alias2", nodeID: "node-x", providerID: "prov-2"}
	m.mu.Unlock()

	m.mu.Lock()
	inf1, q1 := m.getStatsForProviderLocked("node-x", "prov-1")
	inf2, q2 := m.getStatsForProviderLocked("node-x", "prov-2")
	inf3, q3 := m.getStatsForProviderLocked("node-y", "prov-1")
	m.mu.Unlock()

	if inf1 != 2 || q1 != 0 {
		t.Errorf("prov-1 on node-x: inflight=%d queued=%d, want 2/0", inf1, q1)
	}
	if inf2 != 1 || q2 != 0 {
		t.Errorf("prov-2 on node-x: inflight=%d queued=%d, want 1/0", inf2, q2)
	}
	if inf3 != 0 || q3 != 0 {
		t.Errorf("prov-1 on node-y: inflight=%d queued=%d, want 0/0", inf3, q3)
	}
}

// TestModelQueueContextCancelRemovesQueuedItem verifies that cancelling the
// context of a queued admit removes the item and returns context.Canceled.
func TestModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-cc1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// Fill the only slot.
	n, err := m.admit(context.Background(), "g-cc", "", "", cands, defPolicy)
	if err != nil || n == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Queue a second admit with a cancellable context.
	ctx, cancel := context.WithCancel(context.Background())
	resultCh := make(chan error, 1)
	go func() {
		_, err := m.admit(ctx, "g-cc", "", "", cands, defPolicy)
		resultCh <- err
	}()

	waitForQueueLen(t, m, "g-cc", 1)
	cancel()

	select {
	case err := <-resultCh:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("expected context.Canceled, got: %v", err)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: context cancellation not handled")
	}

	// Queue must be empty after cancellation.
	m.mu.Lock()
	qLen := 0
	if g, ok := m.groups["g-cc"]; ok {
		qLen = len(g.queue)
	}
	m.mu.Unlock()
	if qLen != 0 {
		t.Errorf("queue should be empty after cancel, got %d items", qLen)
	}
}

// TestModelQueueProviderCapacityIsPerProviderSlot verifies that same-node
// multiple provider candidates each have independent capacity/in-flight
// accounting per REVIEW_API-2.
func TestModelQueueProviderCapacityIsPerProviderSlot(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-same"}
	// Two provider candidates on the same node, each capacity=1.
	cands := []candidateNode{
		{entry: entry, capacity: 1, providerID: "prov-a", adapter: "vllm"},
		{entry: entry, capacity: 1, providerID: "prov-b", adapter: "vllm"},
	}
	m := newModelQueueManager(nil)

	// First admit: pick prov-a (deterministic tie-break by providerID).
	sel1, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
	if err != nil || sel1 == nil {
		t.Fatalf("first admit: %v", err)
	}
	if sel1.providerID != "prov-a" {
		t.Fatalf("expected prov-a, got %s", sel1.providerID)
	}

	// Second admit: should pick prov-b independently (prov-a is full, prov-b has capacity).
	sel2, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
	if err != nil || sel2 == nil {
		t.Fatalf("second admit (prov-b should be available): %v", err)
	}
	if sel2.providerID != "prov-b" {
		t.Fatalf("expected prov-b, got %s", sel2.providerID)
	}

	// Third admit: both providers at capacity, should queue.
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err = m.admit(ctx, "g-same", "", "", cands, groupPolicy{})
	if !errors.Is(err, errQueueTimeout) && !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("expected queue timeout or deadline exceeded, got: %v", err)
	}
}

// TestResolveProviderPoolCandidatesFiltersInvalidProviders verifies that
// unavailable providers, served model mismatch, empty adapter providers,
// and capacity zero/unknown providers are excluded from dispatch candidates.
// This test calls the actual Service.resolveProviderPoolCandidates method.
func TestResolveProviderPoolCandidatesFiltersInvalidProviders(t *testing.T) {
	// Build model catalog entry for "qwen3.6:35b".
	catalog := []config.ModelCatalogEntry{
		{
			ID: "qwen3.6:35b",
			Providers: map[string]string{
				"prov-available":   "served-qwen",
				"prov-unavailable": "served-qwen",
				"prov-mismatch":    "served-qwen",
				"prov-no-adapter":  "served-qwen",
				"prov-cap-zero":    "served-qwen",
				"prov-cap-unknown": "served-qwen",
			},
		},
	}

	// Build NodeStore with multiple providers.
	store := edgenode.NewNodeStore()

	// Valid provider: available, has adapter, has capacity, served model matches provider's models.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-valid",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-available",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen", "served-llama"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Invalid: health = "unavailable".
	store.Add(&edgenode.NodeRecord{
		ID:      "node-bad-health",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-unavailable",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "unavailable",
				Capacity: 2,
			},
		},
	})

	// Invalid: served model not in provider's own models list.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-mismatch",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-mismatch",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-llama"}, // does NOT include "served-qwen"
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Invalid: empty adapter.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-no-adapter",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-no-adapter",
				Adapter:  "",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Invalid: capacity = 0.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-cap-zero",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-cap-zero",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 0,
			},
		},
	})

	// Invalid: capacity < 0 (negative/unknown).
	store.Add(&edgenode.NodeRecord{
		ID:      "node-cap-unknown",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-cap-unknown",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: -1,
			},
		},
	})

	// Build a fake registry with all nodes.
	reg := edgenode.NewRegistry()
	allRecs := store.All()
	for _, rec := range allRecs {
		entry := &edgenode.NodeEntry{
			NodeID:         rec.ID,
			LifecycleState: edgenode.LifecycleConnected,
		}
		reg.Register(entry)
	}

	// Create Service with catalog and node store.
	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// Call the actual resolveProviderPoolCandidates.
	req := SubmitRunRequest{
		ModelGroupKey: "qwen3.6:35b",
		ProviderPool:  true,
	}
	storeSnapshot, catalogSnapshot := svc.runtimeConfigSnapshot()
	candidates, policy, err := svc.resolveProviderPoolCandidates(req, storeSnapshot, catalogSnapshot)
	if err != nil {
		t.Fatalf("resolveProviderPoolCandidates: %v", err)
	}

	// Only prov-available should be in candidates.
	if len(candidates) != 1 {
		t.Fatalf("expected 1 candidate, got %d: %v", len(candidates), candidates)
	}

	c := candidates[0]
	if c.providerID != "prov-available" {
		t.Errorf("providerID: got %q, want %q", c.providerID, "prov-available")
	}
	if c.adapter != "vllm-gpu" {
		t.Errorf("adapter: got %q, want %q", c.adapter, "vllm-gpu")
	}
	if c.servedTarget != "served-qwen" {
		t.Errorf("servedTarget: got %q, want %q", c.servedTarget, "served-qwen")
	}
	if c.capacity != 2 {
		t.Errorf("capacity: got %d, want %d", c.capacity, 2)
	}
	if c.capacity <= 0 {
		t.Error("capacity must be > 0 for dispatchable provider")
	}

	// Policy should use defaults since none of the providers set policy.
	if policy.maxQueue <= 0 {
		t.Errorf("policy.maxQueue: got %d, expected > 0", policy.maxQueue)
	}
	if policy.queueTimeout <= 0 {
		t.Errorf("policy.queueTimeout: got %v, expected > 0", policy.queueTimeout)
	}
}

// TestSubmitRunProviderPoolRewritesAdapterAndTarget verifies that provider-pool
// SubmitRun rewrites both req.Adapter and req.Target from the selected candidate.
// This test uses a fake TCP client via net.Pipe to capture the actual RunRequest
// sent by SubmitRun(ProviderPool=true), confirming that the winning candidate's
// adapter and servedTarget propagate correctly through the full service path.
func TestSubmitRunProviderPoolRewritesAdapterAndTarget(t *testing.T) {
	// Use net.Pipe to create a fake node connection that captures the RunRequest.
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer 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 capturedReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedReq = req
		capturedMu.Unlock()
	})

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

	// Build NodeStore with a provider-pool provider.
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pool",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm-01",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

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

	// Create Service with queue and catalog.
	// events bus must be non-nil to activate the queue path for provider-pool.
	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// SubmitRun with ProviderPool=true.
	result, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
		RunID:         "run-pool-test-001",
		ModelGroupKey: "qwen3.6:35b",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("SubmitRun: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil RunResult")
	}

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

	capturedMu.Lock()
	defer capturedMu.Unlock()

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

	// Verify that the adapter and target were rewritten from the provider-pool candidate.
	if capturedReq.GetAdapter() != "vllm-gpu" {
		t.Errorf("adapter: got %q, want %q", capturedReq.GetAdapter(), "vllm-gpu")
	}
	if capturedReq.GetTarget() != "served-qwen" {
		t.Errorf("target: got %q, want %q", capturedReq.GetTarget(), "served-qwen")
	}
	if capturedReq.GetRunId() != "run-pool-test-001" {
		t.Errorf("runID: got %q, want %q", capturedReq.GetRunId(), "run-pool-test-001")
	}
}