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

package service

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

// TestRefreshProviderCapacityAffectsNextDispatch verifies that a provider
// capacity change applied via SetRuntimeConfig is reflected in the candidate
// capacity used for the next dispatch admission, since candidates are rebuilt
// from the live runtime snapshot on each request.
func TestRefreshProviderCapacityAffectsNextDispatch(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"})
	svc := New(reg, nil)

	newStore := func(capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cap",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:       "prov-a",
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{"served-a"},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	catalog := []config.ModelCatalogEntry{{
		ID:        "qwen3.6:35b",
		Providers: map[string]string{"prov-a": "served-a"},
	}}

	req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true}

	svc.SetRuntimeConfig(newStore(2), catalog)
	store, cat := svc.runtimeConfigSnapshot()
	cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve before refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 2 {
		t.Fatalf("before refresh: got candidates=%+v, want one with capacity=2", cands)
	}

	// Refresh raises capacity to 8; the next candidate build must observe it.
	svc.SetRuntimeConfig(newStore(8), catalog)
	store, cat = svc.runtimeConfigSnapshot()
	cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve after refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 8 {
		t.Fatalf("after refresh: got candidates=%+v, want one with capacity=8", cands)
	}
}

// TestRefreshProviderLongCapacityAffectsCandidate verifies that a provider
// long_context_capacity change applied via SetRuntimeConfig is reflected in the
// candidate longContextCapacity used for the next long-context admission.
func TestRefreshProviderLongCapacityAffectsCandidate(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"})
	svc := New(reg, nil)

	newStore := func(longCapacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cap",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:                  "prov-a",
					Type:                "vllm",
					Category:            config.CategoryAPI,
					Adapter:             "vllm-gpu",
					Models:              []string{"served-a"},
					Health:              "available",
					Capacity:            4,
					TotalContextTokens:  524288,
					LongContextCapacity: longCapacity,
				},
			},
		})
		return store
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                  "qwen3.6:35b",
		ContextWindowTokens: 262144,
		Providers:           map[string]string{"prov-a": "served-a"},
	}}

	req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true, ContextClass: "long"}

	svc.SetRuntimeConfig(newStore(1), catalog)
	store, cat := svc.runtimeConfigSnapshot()
	cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve before refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].longContextCapacity != 1 {
		t.Fatalf("before refresh: got candidates=%+v, want one with longContextCapacity=1", cands)
	}

	// Refresh raises long capacity to 2; the next candidate build must observe it.
	svc.SetRuntimeConfig(newStore(2), catalog)
	store, cat = svc.runtimeConfigSnapshot()
	cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve after refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].longContextCapacity != 2 {
		t.Fatalf("after refresh: got candidates=%+v, want one with longContextCapacity=2", cands)
	}
}

func TestRuntimeConfigSnapshotConcurrentReplace(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-refresh"})
	svc := New(reg, nil)

	newStore := func(providerID, servedModel string, capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID:      "node-refresh",
			Runtime: config.RuntimeConf{Concurrency: capacity},
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:       providerID,
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{servedModel},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	newCatalog := func(providerID, servedModel string) []config.ModelCatalogEntry {
		return []config.ModelCatalogEntry{
			{
				ID: "qwen3.6:35b",
				Providers: map[string]string{
					providerID: servedModel,
				},
			},
		}
	}

	storeA := newStore("prov-a", "served-a", 2)
	catalogA := newCatalog("prov-a", "served-a")
	storeB := newStore("prov-b", "served-b", 4)
	catalogB := newCatalog("prov-b", "served-b")
	svc.SetRuntimeConfig(storeA, catalogA)

	const readers = 8
	const iterations = 200
	start := make(chan struct{})
	errCh := make(chan error, readers)
	var wg sync.WaitGroup

	for i := 0; i < readers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			for j := 0; j < iterations; j++ {
				store, catalog := svc.runtimeConfigSnapshot()
				candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{
					ModelGroupKey: "qwen3.6:35b",
					ProviderPool:  true,
				}, store, catalog)
				if err != nil {
					errCh <- err
					return
				}
				if len(candidates) != 1 {
					errCh <- errors.New("expected exactly one provider candidate")
					return
				}
				_ = svc.ListNodeSnapshots()
			}
		}()
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < iterations; i++ {
			svc.SetRuntimeConfig(storeA, catalogA)
			svc.SetRuntimeConfig(storeB, catalogB)
		}
	}()

	close(start)
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("runtime config snapshot reader observed inconsistent state: %v", err)
		}
	}
}

func TestServiceRuntimeConfigRefreshPreservesInflight(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-ref-pres"})
	bus := edgeevents.NewBus()
	svc := New(reg, bus)

	newStore := func(capacity, longCapacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-ref-pres",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:                  "prov-a",
					Type:                "vllm",
					Category:            config.CategoryAPI,
					Adapter:             "vllm-gpu",
					Models:              []string{"served-a"},
					Health:              "available",
					Capacity:            capacity,
					LongContextCapacity: longCapacity,
				},
			},
		})
		return store
	}

	catalog := []config.ModelCatalogEntry{{
		ID:        "qwen3.6:35b",
		Providers: map[string]string{"prov-a": "served-a"},
	}}

	svc.SetRuntimeConfig(newStore(4, 2), catalog)

	// Simulate admission by reserving slot
	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: "node-ref-pres", providerID: "prov-a"}
	res := svc.queue.resources[key]
	if res == nil {
		svc.queue.mu.Unlock()
		t.Fatal("expected provider resource state to exist")
	}
	res.inFlight = 3
	res.longInFlight = 2
	svc.queue.mu.Unlock()

	// Update store (reduce capacity and long capacity)
	svc.SetRuntimeConfig(newStore(2, 1), catalog)

	// Verify capacity changed but inflight counts preserved
	svc.queue.mu.Lock()
	if res.capacity != 2 {
		t.Errorf("expected capacity updated to 2, got %d", res.capacity)
	}
	if res.longCapacity != 1 {
		t.Errorf("expected longCapacity updated to 1, got %d", res.longCapacity)
	}
	if res.inFlight != 3 {
		t.Errorf("expected inFlight preserved at 3, got %d", res.inFlight)
	}
	if res.longInFlight != 2 {
		t.Errorf("expected longInFlight preserved at 2, got %d", res.longInFlight)
	}
	svc.queue.mu.Unlock()
}

// refreshWakeupFixture builds a service whose single provider can be reconfigured
// between calls, plus the candidate list a queued request would carry for it.
func refreshWakeupFixture(t *testing.T, nodeID, providerID, servedModel string) (*Service, func(capacity int, enabled bool), []candidateNode) {
	t.Helper()

	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{NodeID: nodeID}
	reg.Register(entry)
	svc := New(reg, edgeevents.NewBus())

	catalog := []config.ModelCatalogEntry{{
		ID:        "group-refresh",
		Providers: map[string]string{providerID: servedModel},
	}}

	applyConfig := func(capacity int, enabled bool) {
		enabledFlag := enabled
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: nodeID,
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{{
				ID:       providerID,
				Type:     "vllm",
				Adapter:  "vllm-gpu",
				Models:   []string{servedModel},
				Health:   "available",
				Capacity: capacity,
				Enabled:  &enabledFlag,
			}},
		})
		svc.SetRuntimeConfig(store, catalog)
	}

	applyConfig(1, true)
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: providerID, servedTarget: servedModel}}
	return svc, applyConfig, cands
}

// TestGlobalPumpCapacityRefreshWakesWaiter verifies that raising provider
// capacity through a runtime config refresh wakes an already-queued request.
// A capacity increase frees room that no lease release will ever announce, so
// without pumping on refresh the waiter would sleep until its queue timeout.
func TestGlobalPumpCapacityRefreshWakesWaiter(t *testing.T) {
	svc, applyConfig, cands := refreshWakeupFixture(t, "node-cap-wake", "prov-cap-wake", "served-cap-wake")

	// Fill the provider's only slot.
	if _, err := svc.queue.admit(t.Context(), "g-cap-wake-a", "", "", cands, groupPolicy{}, nil); err != nil {
		t.Fatalf("initial admit: %v", err)
	}

	// A request in a different group queues behind it.
	waiter := queueItemForTest(cands, false)
	enqueueForTest(svc.queue, "g-cap-wake-b", waiter, nil)

	// Raising capacity must dispatch the waiter immediately.
	applyConfig(2, true)

	select {
	case res := <-waiter.waitCh:
		if res.err != nil {
			t.Fatalf("waiter expected dispatch after capacity refresh, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.capacity != 2 {
			t.Fatalf("waiter must be dispatched with the refreshed capacity, got %+v", res.candidate)
		}
	default:
		t.Fatal("capacity refresh did not wake the queued waiter")
	}

	// Over-capacity must still be refused: both slots are now taken.
	inFlight, _ := providerResourceCounts(svc.queue, "node-cap-wake", "prov-cap-wake")
	if inFlight != 2 {
		t.Fatalf("expected 2 in-flight after the refresh dispatch, got %d", inFlight)
	}
	extra := queueItemForTest(cands, false)
	enqueueForTest(svc.queue, "g-cap-wake-b", extra, nil)
	svc.queue.mu.Lock()
	svc.queue.pumpAllLocked()
	svc.queue.mu.Unlock()
	select {
	case <-extra.waitCh:
		t.Fatal("pump dispatched past the refreshed capacity")
	default:
	}
}

// TestGlobalPumpDisabledProviderIsNotDispatched verifies that dispatch re-reads
// the enabled switch at hand-off time. A request queued while the provider was
// enabled must not be handed the provider after a refresh disabled it, even
// though its candidate snapshot still says the provider is usable — and it must
// resume once the provider comes back.
func TestGlobalPumpDisabledProviderIsNotDispatched(t *testing.T) {
	svc, applyConfig, cands := refreshWakeupFixture(t, "node-dis", "prov-dis", "served-dis")

	admitted, err := svc.queue.admit(t.Context(), "g-dis-a", "", "", cands, groupPolicy{}, nil)
	if err != nil || admitted == nil {
		t.Fatalf("initial admit: %v", err)
	}

	waiter := queueItemForTest(cands, false)
	enqueueForTest(svc.queue, "g-dis-b", waiter, nil)

	// Disable the provider, then free the in-flight slot. The freed capacity must
	// not be handed to the waiter: the provider is no longer dispatchable.
	applyConfig(1, false)
	svc.queue.releaseLease(admitted.leaseID, "complete")

	select {
	case res := <-waiter.waitCh:
		t.Fatalf("waiter was dispatched to a disabled provider: %+v", res)
	default:
	}
	if got := pendingItemCount(svc.queue); got != 1 {
		t.Fatalf("expected the waiter to stay queued, got %d pending items", got)
	}

	// Re-enabling must wake it through the same refresh pump.
	applyConfig(1, true)
	select {
	case res := <-waiter.waitCh:
		if res.err != nil {
			t.Fatalf("waiter expected dispatch after re-enable, got error: %v", res.err)
		}
	default:
		t.Fatal("re-enabling the provider did not wake the queued waiter")
	}
}

// TestGlobalPumpReenabledProviderAbsentAtEnqueueWakesWaiter verifies the
// required regression: a provider that was disabled at enqueue time is absent
// from the candidate snapshot, so the original snapshot-only pump could never
// discover it after re-enable. With the live resolver the waiter must wake
// and dispatch to the re-enabled provider without any lease release.
func TestGlobalPumpReenabledProviderAbsentAtEnqueueWakesWaiter(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)
	_ = toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConn2, 0, 0, parserMap)

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-re-a", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-re-b", Client: edgeClient2})
	svc := New(reg, edgeevents.NewBus())

	catalog := []config.ModelCatalogEntry{{
		ID:        "group-re",
		Providers: map[string]string{"prov-a": "served-a", "prov-b": "served-b"},
	}}

	applyConfig := func(capA, capB int, enabledB bool) {
		enabledFlag := true
		if !enabledB {
			enabledFlag = false
		}
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-re-a",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{{
				ID: "prov-a", Type: "vllm", Adapter: "vllm-gpu",
				Models: []string{"served-a"}, Health: "available",
				Capacity: capA,
			}},
		})
		store.Add(&edgenode.NodeRecord{
			ID: "node-re-b",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
				},
			},
			Providers: []config.NodeProviderConf{{
				ID: "prov-b", Type: "vllm", Adapter: "vllm-gpu",
				Models: []string{"served-b"}, Health: "available",
				Capacity: capB, Enabled: &enabledFlag,
			}},
		})
		svc.SetRuntimeConfig(store, catalog)
	}

	// Fill prov-a's only slot so prov-b is the only candidate for the waiter.
	applyConfig(1, 1, true)
	admitted, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-re", ProviderPool: true, Background: true,
	})
	if err != nil || admitted == nil {
		t.Fatalf("initial admit: %v", err)
	}
	defer admitted.Close()

	// Disable prov-b and enqueue a waiter — prov-b must be absent from the
	// resolver's candidate list at enqueue time.
	applyConfig(1, 1, false)

	type submitResult struct {
		run RunResult
		err error
	}
	resultCh := make(chan submitResult, 1)
	go func() {
		run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
			ModelGroupKey: "group-re", ProviderPool: true, Background: true,
		})
		resultCh <- submitResult{run: run, err: err}
	}()

	// Verify the waiter is queued and prov-b is NOT in its stored enqueue snapshot.
	requireQueued := func() {
		deadline := time.After(2 * time.Second)
		for {
			svc.queue.mu.Lock()
			group := svc.queue.groups["group-re"]
			count := 0
			if group != nil {
				count = len(group.queue)
			}
			svc.queue.mu.Unlock()
			if count == 1 {
				return
			}
			select {
			case <-deadline:
				t.Fatalf("waiter was not enqueued within timeout")
			default:
				runtime.Gosched()
			}
		}
	}
	requireQueued()
	svc.queue.mu.Lock()
	group := svc.queue.groups["group-re"]
	if group == nil || len(group.queue) != 1 {
		svc.queue.mu.Unlock()
		t.Fatalf("expected exactly one queued item")
	}
	item := group.queue[0]
	if item.resolveCandidates == nil {
		svc.queue.mu.Unlock()
		t.Fatal("expected live resolver on queued item")
	}
	for _, candidate := range item.candidates {
		if candidate.providerID == "prov-b" {
			svc.queue.mu.Unlock()
			t.Fatalf("prov-b must be absent from stored enqueue snapshot")
		}
	}
	svc.queue.mu.Unlock()

	// Re-enable prov-b — the live resolver must now include it and the pump
	// must dispatch the waiter to prov-b without any lease release.
	applyConfig(1, 1, true)

	var res3 RunResult
	var err3 error
	select {
	case got := <-resultCh:
		err3 = got.err
		res3 = got.run
	case <-time.After(500 * time.Millisecond):
		t.Fatal("re-enabling prov-b did not wake the queued waiter (provider absent-at-enqueue regression)")
	}

	if err3 != nil {
		t.Fatalf("waiter expected dispatch after re-enable, got error: %v", err3)
	}
	if res3 == nil {
		t.Fatal("waiter expected non-nil result")
	}
	defer res3.Close()
	if got := res3.Dispatch().ProviderID; got != "prov-b" {
		t.Fatalf("waiter dispatched to provider %q, want prov-b", got)
	}

	// Verify both providers have in-flight=1 after dispatch.
	preAInFlight, _ := providerResourceCounts(svc.queue, "node-re-a", "prov-a")
	preBInFlight, _ := providerResourceCounts(svc.queue, "node-re-b", "prov-b")
	if preAInFlight != 1 {
		t.Fatalf("expected prov-a in-flight=1 after dispatch, got %d", preAInFlight)
	}
	if preBInFlight != 1 {
		t.Fatalf("expected prov-b in-flight=1 after dispatch, got %d", preBInFlight)
	}

	// Release both runs through terminal lifecycle events and verify full drain.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted.Dispatch().RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: res3.Dispatch().RunID, Type: "complete"})
	aInFlight, aLong := providerResourceCounts(svc.queue, "node-re-a", "prov-a")
	bInFlight, bLong := providerResourceCounts(svc.queue, "node-re-b", "prov-b")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected prov-a drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if bInFlight != 0 || bLong != 0 {
		t.Fatalf("expected prov-b drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0 after terminal events, got %d", lc)
	}
}

// TestGlobalPumpCapacityZeroProviderBecomesEligible verifies the required
// regression: a provider with capacity 0 at enqueue time is absent from the
// candidate snapshot, so the original snapshot-only pump could never discover
// it after capacity becomes positive. With the live resolver the waiter must
// wake and dispatch to the provider once capacity is raised.
func TestGlobalPumpCapacityZeroProviderBecomesEligible(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)
	_ = toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConn2, 0, 0, parserMap)

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cz-a", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cz-b", Client: edgeClient2})
	svc := New(reg, edgeevents.NewBus())

	catalog := []config.ModelCatalogEntry{{
		ID:        "group-cz",
		Providers: map[string]string{"prov-a": "served-a", "prov-b": "served-b"},
	}}

	applyConfig := func(capA, capB int) {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cz-a",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{{
				ID: "prov-a", Type: "vllm", Adapter: "vllm-gpu",
				Models: []string{"served-a"}, Health: "available",
				Capacity: capA,
			}},
		})
		store.Add(&edgenode.NodeRecord{
			ID: "node-cz-b",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
				},
			},
			Providers: []config.NodeProviderConf{{
				ID: "prov-b", Type: "vllm", Adapter: "vllm-gpu",
				Models: []string{"served-b"}, Health: "available",
				Capacity: capB,
			}},
		})
		svc.SetRuntimeConfig(store, catalog)
	}

	// Fill prov-a's only slot so prov-b is the only candidate for the waiter.
	applyConfig(1, 1)
	admitted, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-cz", ProviderPool: true, Background: true,
	})
	if err != nil || admitted == nil {
		t.Fatalf("initial admit: %v", err)
	}
	defer admitted.Close()

	// Set prov-b to capacity 0 and enqueue a waiter — prov-b must be absent
	// from the resolver's candidate list at enqueue time.
	applyConfig(1, 0)

	type submitResult struct {
		run RunResult
		err error
	}
	resultCh := make(chan submitResult, 1)
	go func() {
		run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
			ModelGroupKey: "group-cz", ProviderPool: true, Background: true,
		})
		resultCh <- submitResult{run: run, err: err}
	}()

	// Verify the waiter is queued and prov-b is NOT in its stored enqueue snapshot.
	requireQueued := func() {
		deadline := time.After(2 * time.Second)
		for {
			svc.queue.mu.Lock()
			group := svc.queue.groups["group-cz"]
			count := 0
			if group != nil {
				count = len(group.queue)
			}
			svc.queue.mu.Unlock()
			if count == 1 {
				return
			}
			select {
			case <-deadline:
				t.Fatalf("waiter was not enqueued within timeout")
			default:
				runtime.Gosched()
			}
		}
	}
	requireQueued()
	svc.queue.mu.Lock()
	group := svc.queue.groups["group-cz"]
	if group == nil || len(group.queue) != 1 {
		svc.queue.mu.Unlock()
		t.Fatalf("expected exactly one queued item")
	}
	item := group.queue[0]
	if item.resolveCandidates == nil {
		svc.queue.mu.Unlock()
		t.Fatal("expected live resolver on queued item")
	}
	for _, candidate := range item.candidates {
		if candidate.providerID == "prov-b" {
			svc.queue.mu.Unlock()
			t.Fatalf("prov-b must be absent from stored enqueue snapshot")
		}
	}
	svc.queue.mu.Unlock()

	// Raise prov-b's capacity — the live resolver must now include it and the
	// pump must dispatch the waiter to prov-b without any lease release.
	applyConfig(1, 1)

	var res3 RunResult
	var err3 error
	select {
	case got := <-resultCh:
		err3 = got.err
		res3 = got.run
	case <-time.After(500 * time.Millisecond):
		t.Fatal("raising prov-b capacity did not wake the queued waiter (capacity-zero-at-enqueue regression)")
	}

	if err3 != nil {
		t.Fatalf("waiter expected dispatch after capacity raise, got error: %v", err3)
	}
	if res3 == nil {
		t.Fatal("waiter expected non-nil result")
	}
	defer res3.Close()
	if got := res3.Dispatch().ProviderID; got != "prov-b" {
		t.Fatalf("waiter dispatched to provider %q, want prov-b", got)
	}

	// Verify both providers have in-flight=1 after dispatch.
	preAInFlight, _ := providerResourceCounts(svc.queue, "node-cz-a", "prov-a")
	preBInFlight, _ := providerResourceCounts(svc.queue, "node-cz-b", "prov-b")
	if preAInFlight != 1 {
		t.Fatalf("expected prov-a in-flight=1 after dispatch, got %d", preAInFlight)
	}
	if preBInFlight != 1 {
		t.Fatalf("expected prov-b in-flight=1 after dispatch, got %d", preBInFlight)
	}

	// Release both runs through terminal lifecycle events and verify full drain.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted.Dispatch().RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: res3.Dispatch().RunID, Type: "complete"})
	aInFlight, aLong := providerResourceCounts(svc.queue, "node-cz-a", "prov-a")
	bInFlight, bLong := providerResourceCounts(svc.queue, "node-cz-b", "prov-b")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected prov-a drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if bInFlight != 0 || bLong != 0 {
		t.Fatalf("expected prov-b drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0 after terminal events, got %d", lc)
	}
}