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

package service

import (
	"context"
	"errors"
	"fmt"
	"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, groupPolicy{})
	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, groupPolicy{})
	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, groupPolicy{})
	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, groupPolicy{})
	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")
	policyA := NewGroupPolicy(3, 500*time.Millisecond)
	storeB := newStore("prov-b", "served-b", 4)
	catalogB := newCatalog("prov-b", "served-b")
	policyB := NewGroupPolicy(7, 2*time.Second)
	svc.SetRuntimeConfig(storeA, catalogA, policyA)

	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, policy := svc.runtimeConfigSnapshot()
				// Snapshot must return a self-consistent (store, catalog, policy) tuple.
				// A reader must never observe storeA with policyB, or storeB with policyA.
				snapProviders := store.All()
				if len(snapProviders) != 1 {
					errCh <- fmt.Errorf("snapshot store has %d providers, want 1", len(snapProviders))
					return
				}
				providerID := snapProviders[0].Providers[0].ID
				// Each snapshot must pair the correct policy with the correct provider.
				if providerID == "prov-a" && (policy.maxQueue != policyA.maxQueue || policy.queueTimeout != policyA.queueTimeout) {
					errCh <- fmt.Errorf("A snapshot has B policy: store=%s policy=(max=%d,timeout=%v) want (max=%d,timeout=%v)",
						providerID, policy.maxQueue, policy.queueTimeout, policyA.maxQueue, policyA.queueTimeout)
					return
				}
				if providerID == "prov-b" && (policy.maxQueue != policyB.maxQueue || policy.queueTimeout != policyB.queueTimeout) {
					errCh <- fmt.Errorf("B snapshot has A policy: store=%s policy=(max=%d,timeout=%v) want (max=%d,timeout=%v)",
						providerID, policy.maxQueue, policy.queueTimeout, policyB.maxQueue, policyB.queueTimeout)
					return
				}
				// Policy must always be one of the two known policies (never a stale or mixed value).
				if policy.maxQueue != policyA.maxQueue && policy.maxQueue != policyB.maxQueue {
					errCh <- fmt.Errorf("snapshot has unknown policy maxQueue=%d (want %d or %d)",
						policy.maxQueue, policyA.maxQueue, policyB.maxQueue)
					return
				}
				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, policyA)
			svc.SetRuntimeConfig(storeB, catalogB, policyB)
		}
	}()

	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, groupPolicy{})

	// 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, groupPolicy{})

	// 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, groupPolicy{})
	}

	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, false, false); 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, false, false)
	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, groupPolicy{})
	}

	// 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, groupPolicy{})
	}

	// 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)
	}
}

// --- PROVIDER_POLICY_REFRESH-2: resource attribute refresh tests ---

// helper for capacity/priority/disable rebase tests.
// Returns a service with a mock client pair so SubmitRun can dispatch.
func buildCapacityRebaseService(t *testing.T, nodeID, providerID, servedModel string, initialCap, initialLongCap int, enabled *bool) (*Service, func(capacity, longCapacity int, enabled bool)) {
	t.Helper()
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		edgeConn.Close()
		nodeConn.Close()
	})
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
	svc := New(reg, edgeevents.NewBus())
	catalog := []config.ModelCatalogEntry{{
		ID:        "group-cap-rebase",
		Providers: map[string]string{providerID: servedModel},
	}}
	applyFn := func(capacity, longCapacity 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,
				LongContextCapacity: longCapacity,
				Enabled:             &enabledFlag,
			}},
		})
		svc.SetRuntimeConfig(store, catalog, groupPolicy{maxQueue: 8, queueTimeout: 5 * time.Second})
	}
	applyFn(initialCap, initialLongCap, *enabled)
	return svc, applyFn
}

// buildPriorityService sets up two providers with separate mock clients.
func buildPriorityService(t *testing.T) (*Service, func(capA, capB, prioA, prioB int, enabled bool)) {
	t.Helper()
	edgeConnA, nodeConnA := net.Pipe()
	edgeConnB, nodeConnB := net.Pipe()
	t.Cleanup(func() {
		edgeConnA.Close()
		nodeConnA.Close()
		edgeConnB.Close()
		nodeConnB.Close()
	})
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeClientA := toki.NewTcpClient(edgeConnA, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConnA, 0, 0, parserMap)
	edgeClientB := toki.NewTcpClient(edgeConnB, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConnB, 0, 0, parserMap)

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-prio-a", Client: edgeClientA})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-prio-b", Client: edgeClientB})
	svc := New(reg, edgeevents.NewBus())
	catalog := []config.ModelCatalogEntry{{
		ID:        "group-prio",
		Providers: map[string]string{"prov-a": "served-a", "prov-b": "served-b"},
	}}
	applyFn := func(capA, capB, prioA, prioB int, enabled bool) {
		enabledFlag := enabled
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-prio-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, Priority: prioA, Enabled: &enabledFlag,
			}},
		})
		store.Add(&edgenode.NodeRecord{
			ID: "node-prio-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, Priority: prioB, Enabled: &enabledFlag,
			}},
		})
		svc.SetRuntimeConfig(store, catalog, groupPolicy{maxQueue: 8, queueTimeout: 5 * time.Second})
	}
	applyFn(1, 0, 0, 0, true)
	return svc, applyFn
}

// TestRefreshCapacityShrinkPreservesLeaseAndBlocksUntilRecovery verifies that
// shrinking capacity while in-flight leases exist preserves those leases and
// blocks new admission, and that raising capacity after a release immediately
// pumps the queued waiter.
func TestRefreshCapacityShrinkPreservesLeaseAndBlocksUntilRecovery(t *testing.T) {
	enabled := true
	svc, applyFn := buildCapacityRebaseService(t, "node-cap-shrink", "prov-shrink", "served-shrink", 2, 0, &enabled)

	admitted1, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-cap-rebase", ProviderPool: true, Background: true,
	})
	if err != nil || admitted1 == nil {
		t.Fatalf("admit 1 failed: %v", err)
	}
	admitted2, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-cap-rebase", ProviderPool: true, Background: true,
	})
	if err != nil || admitted2 == nil {
		admitted1.Close()
		t.Fatalf("admit 2 failed: %v", err)
	}

	svc.queue.mu.Lock()
	res := svc.queue.resources[providerResourceKey{nodeID: "node-cap-shrink", providerID: "prov-shrink"}]
	if res == nil || res.inFlight != 2 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected inFlight=2 after admits, got %d", res.inFlight)
	}
	svc.queue.mu.Unlock()

	applyFn(1, 0, true)

	svc.queue.mu.Lock()
	if res.capacity != 1 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected capacity=1 after shrink, got %d", res.capacity)
	}
	if res.inFlight != 2 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected inFlight preserved at 2 after shrink, got %d", res.inFlight)
	}
	svc.queue.mu.Unlock()

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

	waitForQueueLen(t, svc.queue, "group-cap-rebase", 1)
	if inFlight, _ := providerResourceCounts(svc.queue, "node-cap-shrink", "prov-shrink"); inFlight != 2 {
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("blocked admission increased over-capacity in-flight to %d", inFlight)
	}

	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted1.Dispatch().RunID, Type: "complete"})
	applyFn(2, 0, true)

	select {
	case got := <-blockedCh:
		if got.err != nil {
			admitted2.Close()
			t.Fatalf("blocked waiter expected dispatch after raise, got error: %v", got.err)
		}
		if got.run == nil {
			admitted2.Close()
			t.Fatal("blocked waiter expected non-nil result")
		}
		got.run.Close()
		// Release the blocked waiter's lease.
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: got.run.Dispatch().RunID, Type: "complete"})
		got.run.Close()
		// Release admitted2's lease.
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted2.Dispatch().RunID, Type: "complete"})
		admitted2.Close()
	case <-time.After(2 * time.Second):
		admitted2.Close()
		t.Fatal("raised capacity did not pump the blocked waiter")
	}

	aInFlight, aLong := providerResourceCounts(svc.queue, "node-cap-shrink", "prov-shrink")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected full drain, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0, got %d", lc)
	}
}

// TestRefreshLongCapacityShrinkPreservesLeaseAndBlocksLongAdmission verifies
// that shrinking long context capacity preserves existing long leases and
// blocks new long admission until recovery.
func TestRefreshLongCapacityShrinkPreservesLeaseAndBlocksLongAdmission(t *testing.T) {
	enabled := true
	svc, applyFn := buildCapacityRebaseService(t, "node-lc-shrink", "prov-lc-shrink", "served-lc-shrink", 3, 2, &enabled)

	admitted1, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-cap-rebase", ProviderPool: true, ContextClass: "long",
	})
	if err != nil || admitted1 == nil {
		t.Fatalf("admit long 1 failed: %v", err)
	}
	admitted2, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-cap-rebase", ProviderPool: true, ContextClass: "long",
	})
	if err != nil || admitted2 == nil {
		admitted1.Close()
		t.Fatalf("admit long 2 failed: %v", err)
	}

	svc.queue.mu.Lock()
	res := svc.queue.resources[providerResourceKey{nodeID: "node-lc-shrink", providerID: "prov-lc-shrink"}]
	if res == nil || res.longInFlight != 2 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected longInFlight=2, got %d", res.longInFlight)
	}
	svc.queue.mu.Unlock()

	applyFn(3, 1, true)

	svc.queue.mu.Lock()
	if res.longCapacity != 1 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected longCapacity=1, got %d", res.longCapacity)
	}
	if res.longInFlight != 2 {
		svc.queue.mu.Unlock()
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("expected longInFlight preserved at 2, got %d", res.longInFlight)
	}
	svc.queue.mu.Unlock()

	blockedCh := make(chan struct {
		run RunResult
		err error
	}, 1)
	go func() {
		run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
			ModelGroupKey: "group-cap-rebase", ProviderPool: true, ContextClass: "long",
		})
		blockedCh <- struct {
			run RunResult
			err error
		}{run: run, err: err}
	}()

	waitForQueueLen(t, svc.queue, "group-cap-rebase", 1)
	inFlight, longInFlight := providerResourceCounts(svc.queue, "node-lc-shrink", "prov-lc-shrink")
	if inFlight != 2 || longInFlight != 2 {
		admitted1.Close()
		admitted2.Close()
		t.Fatalf("blocked long admission changed counts: inFlight=%d longInFlight=%d", inFlight, longInFlight)
	}

	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted1.Dispatch().RunID, Type: "complete"})
	applyFn(3, 2, true)

	select {
	case got := <-blockedCh:
		if got.err != nil {
			admitted2.Close()
			t.Fatalf("blocked long waiter expected dispatch, got: %v", got.err)
		}
		got.run.Close()
		// Release the blocked waiter's lease.
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: got.run.Dispatch().RunID, Type: "complete"})
		got.run.Close()
		// Release admitted2's lease.
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: admitted2.Dispatch().RunID, Type: "complete"})
		admitted2.Close()
	case <-time.After(2 * time.Second):
		admitted2.Close()
		t.Fatal("raised long capacity did not pump the blocked waiter")
	}

	aInFlight, aLong := providerResourceCounts(svc.queue, "node-lc-shrink", "prov-lc-shrink")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected full drain, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0, got %d", lc)
	}
}

// TestRefreshPriorityChangesQueuedProviderSelection verifies that a priority
// refresh via SetRuntimeConfig is picked up by the live candidate resolver
// so that a queued waiter selects the higher-priority provider. The priority
// update itself leaves capacity and enabled state unchanged.
func TestRefreshPriorityChangesQueuedProviderSelection(t *testing.T) {
	svc, applyConfig := buildPriorityService(t)
	applyConfig(1, 1, 0, 10, true)

	first, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-prio", ProviderPool: true, Background: true,
	})
	if err != nil {
		t.Fatalf("first SubmitRun: %v", err)
	}
	second, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "group-prio", ProviderPool: true, Background: true,
	})
	if err != nil {
		t.Fatalf("second SubmitRun: %v", err)
	}
	if first.Dispatch().ProviderID != "prov-a" || second.Dispatch().ProviderID != "prov-b" {
		t.Fatalf("unexpected initial provider order: first=%q second=%q", first.Dispatch().ProviderID, second.Dispatch().ProviderID)
	}

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

	// Change priority while both held slots are still occupied so the pump
	// cannot dispatch the waiter. Then release both held slots under the
	// manager lock so the waiter dispatches with the new priority ordering.
	applyConfig(1, 1, 10, 0, true)
	if pending := pendingItemCount(svc.queue); pending != 1 {
		t.Fatalf("priority-only refresh changed pending count to %d, want 1", pending)
	}

	// Release both runs while holding the queue lock so the pump does not
	// dispatch the waiter to the first-available provider on each release.
	svc.queue.mu.Lock()
	m := svc.queue
	m.releaseLeaseLocked(m.leaseByRun[first.Dispatch().RunID])
	m.releaseLeaseLocked(m.leaseByRun[second.Dispatch().RunID])
	m.pumpAllLocked()
	svc.queue.mu.Unlock()

	var third RunResult
	select {
	case got := <-blockedCh:
		if got.err != nil {
			t.Fatalf("queued SubmitRun: %v", got.err)
		}
		third = got.run
		if third.Dispatch().ProviderID != "prov-b" {
			t.Fatalf("priority refresh selected %q, want prov-b", third.Dispatch().ProviderID)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("priority refresh did not pump the queued waiter")
	}

	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: third.Dispatch().RunID, Type: "complete"})
	third.Close()
	if pending := pendingItemCount(svc.queue); pending != 0 {
		t.Fatalf("expected no pending items, got %d", pending)
	}
	if leases := leaseCount(svc.queue); leases != 0 {
		t.Fatalf("expected no live leases, got %d", leases)
	}
}

// TestRefreshDisableReenableReevaluatesAllWaiters verifies that a provider
// disable followed by re-enable re-evaluates ALL queued waiters through the
// live candidate resolver, not just the first.
func TestRefreshDisableReenableReevaluatesAllWaiters(t *testing.T) {
	enabled := true
	svc, applyConfig := buildCapacityRebaseService(t, "node-dre", "prov-dre", "served-dre", 2, 0, &enabled)

	held := make([]RunResult, 0, 2)
	for i := 0; i < 2; i++ {
		run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
			ModelGroupKey: "group-cap-rebase", ProviderPool: true, Background: true,
		})
		if err != nil {
			t.Fatalf("fill capacity %d: %v", i, err)
		}
		held = append(held, run)
	}

	type submitResult struct {
		run RunResult
		err error
	}
	waiters := make(chan submitResult, 2)
	for i := 0; i < 2; i++ {
		go func() {
			run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
				ModelGroupKey: "group-cap-rebase", ProviderPool: true, Background: true,
			})
			waiters <- submitResult{run: run, err: err}
		}()
	}
	waitForQueueLen(t, svc.queue, "group-cap-rebase", 2)

	// Disable while the waiters are queued, then free both held slots. A stale
	// candidate would dispatch here; the live provider snapshot must keep both queued.
	applyConfig(2, 0, false)
	for _, run := range held {
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: run.Dispatch().RunID, Type: "complete"})
		run.Close()
	}
	if pending := pendingItemCount(svc.queue); pending != 2 {
		t.Fatalf("expected 2 pending while provider is disabled, got %d", pending)
	}
	if inFlight, _ := providerResourceCounts(svc.queue, "node-dre", "prov-dre"); inFlight != 0 {
		t.Fatalf("expected held slots to drain while disabled, got inFlight=%d", inFlight)
	}

	// Re-enable through SetRuntimeConfig. Both real SubmitRun waiters must be
	// re-evaluated and dispatched, not merely a synthetic queue item.
	applyConfig(2, 0, true)
	dispatched := make([]RunResult, 0, 2)
	for i := 0; i < 2; i++ {
		select {
		case got := <-waiters:
			if got.err != nil {
				t.Fatalf("waiter %d after re-enable: %v", i, got.err)
			}
			dispatched = append(dispatched, got.run)
		case <-time.After(2 * time.Second):
			t.Fatalf("waiter %d not dispatched after re-enable", i)
		}
	}
	for _, run := range dispatched {
		svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: run.Dispatch().RunID, Type: "complete"})
		run.Close()
	}

	aInFlight, aLong := providerResourceCounts(svc.queue, "node-dre", "prov-dre")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected full drain, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0, got %d", lc)
	}
}

// --- PROVIDER_POLICY_REFRESH-3: concurrent refresh/admission/release race ---

// TestRefreshConcurrentAdmissionReleasePreservesProviderPolicyInvariants verifies
// refresh/admit/release/cancel races across normal and long-context requests.
func TestRefreshConcurrentAdmissionReleasePreservesProviderPolicyInvariants(t *testing.T) {
	const (
		nodeID   = "node-concurrent"
		provider = "prov-concurrent"
	)

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

	catalog := []config.ModelCatalogEntry{{
		ID:        "group-concurrent",
		Providers: map[string]string{provider: "served-concurrent"},
	}}

	cands := []candidateNode{{entry: entry, capacity: 4, providerID: provider}}

	newStore := func(capacity, longCapacity int) *edgenode.NodeStore {
		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:                  provider,
				Type:                "vllm",
				Adapter:             "vllm-gpu",
				Models:              []string{"served-concurrent"},
				Health:              "available",
				Capacity:            capacity,
				LongContextCapacity: longCapacity,
			}},
		})
		return store
	}
	highPolicy := groupPolicy{maxQueue: 64, queueTimeout: 15 * time.Millisecond}
	svc.SetRuntimeConfig(newStore(4, 2), catalog, highPolicy)

	var (
		seenMu     sync.Mutex
		seenLeases = map[uint64]struct{}{}
		peak       int
		peakLong   int
		peakMu     sync.Mutex
	)
	observe := func() {
		inFlight, longInFlight := providerResourceCounts(svc.queue, nodeID, provider)
		peakMu.Lock()
		defer peakMu.Unlock()
		if inFlight > peak {
			peak = inFlight
		}
		if longInFlight > peakLong {
			peakLong = longInFlight
		}
	}

	errs := make(chan error, 2048)
	report := func(err error) {
		if err != nil {
			errs <- err
		}
	}
	start := make(chan struct{})
	var wg sync.WaitGroup
	var cancelWG sync.WaitGroup

	// Repeatedly shrink and expand while workers are admitting and releasing.
	// During each shrink interval occupancy may remain above the new limit due
	// to preserved leases, but it must never grow beyond the post-shrink value.
	wg.Add(1)
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < 80; i++ {
			svc.SetRuntimeConfig(newStore(2, 1), catalog, groupPolicy{
				maxQueue: 64, queueTimeout: time.Duration(i%2) * 15 * time.Millisecond,
			})
			shrinkInFlight, shrinkLong := providerResourceCounts(svc.queue, nodeID, provider)
			normalLimit := max(shrinkInFlight, 2)
			longLimit := max(shrinkLong, 1)
			for sample := 0; sample < 4; sample++ {
				runtime.Gosched()
				inFlight, longInFlight := providerResourceCounts(svc.queue, nodeID, provider)
				if inFlight > normalLimit {
					report(fmt.Errorf("post-shrink inFlight grew from %d to %d (limit %d)", shrinkInFlight, inFlight, normalLimit))
				}
				if longInFlight > longLimit {
					report(fmt.Errorf("post-shrink longInFlight grew from %d to %d (limit %d)", shrinkLong, longInFlight, longLimit))
				}
				if inFlight > 4 || longInFlight > 2 {
					report(fmt.Errorf("global capacity bound exceeded: inFlight=%d longInFlight=%d", inFlight, longInFlight))
				}
			}
			svc.SetRuntimeConfig(newStore(4, 2), catalog, highPolicy)
		}
	}()

	for worker := 0; worker < 8; worker++ {
		worker := worker
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			for attempt := 0; attempt < 24; attempt++ {
				ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
				if attempt%4 == 0 {
					cancelWG.Add(1)
					go func() {
						defer cancelWG.Done()
						runtime.Gosched()
						cancel()
					}()
				}
				candidate, err := svc.queue.admit(ctx, "group-concurrent", "", "", cands,
					highPolicy, nil, worker%2 == 0, true)
				cancel()
				if err != nil {
					if !errors.Is(err, context.DeadlineExceeded) &&
						!errors.Is(err, context.Canceled) &&
						!errors.Is(err, errQueueTimeout) {
						report(fmt.Errorf("worker %d admit %d: %w", worker, attempt, err))
					}
					continue
				}
				seenMu.Lock()
				if _, duplicate := seenLeases[candidate.leaseID]; duplicate {
					report(fmt.Errorf("lease %d handed out more than once", candidate.leaseID))
				}
				seenLeases[candidate.leaseID] = struct{}{}
				seenMu.Unlock()
				observe()
				runtime.Gosched()
				svc.queue.releaseLease(candidate.leaseID, "complete")
				svc.queue.releaseLease(candidate.leaseID, "duplicate_complete")
			}
		}()
	}

	close(start)
	wg.Wait()
	cancelWG.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Error(err)
		}
	}

	svc.SetRuntimeConfig(newStore(4, 2), catalog, highPolicy)
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		inFlight, longInFlight := providerResourceCounts(svc.queue, nodeID, provider)
		if inFlight == 0 && longInFlight == 0 && pendingItemCount(svc.queue) == 0 && leaseCount(svc.queue) == 0 {
			break
		}
		runtime.Gosched()
	}

	peakMu.Lock()
	defer peakMu.Unlock()
	if len(seenLeases) == 0 {
		t.Error("expected at least one successful concurrent admission")
	}
	if peak > 4 {
		t.Errorf("provider in-flight peaked at %d, above configured capacity 4", peak)
	}
	if peakLong > 2 {
		t.Errorf("long in-flight peaked at %d, above configured long capacity 2", peakLong)
	}

	inFlight, longInFlight := providerResourceCounts(svc.queue, nodeID, provider)
	pending := pendingItemCount(svc.queue)
	lc := leaseCount(svc.queue)
	if inFlight != 0 {
		t.Errorf("expected inFlight=0 at end, got %d", inFlight)
	}
	if longInFlight != 0 {
		t.Errorf("expected longInFlight=0 at end, got %d", longInFlight)
	}
	if pending != 0 {
		t.Errorf("expected pending=0 at end, got %d", pending)
	}
	if lc != 0 {
		t.Errorf("expected leaseCount=0 at end, got %d", lc)
	}
}

// resourceStateForTest returns the (present, orphan, generation, inFlight) view
// of a provider resource under the manager lock. Reconnect-activation tests use
// it to assert the orphan marker and generation transitions directly.
func resourceStateForTest(m *modelQueueManager, nodeID, providerID string) (present, orphan bool, generation uint64, inFlight int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	res, ok := m.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
	if !ok {
		return false, false, 0, 0
	}
	return true, res.orphan, res.generation, res.inFlight
}

// TestAcceptedReconnectActivatesCandidateAndPumpsWaiter verifies the S15
// reconnect-candidate-recovery contract at the service/queue boundary: with the
// only alternate provider full and the target node disconnected (its resource
// orphaned), a waiter is stranded; the accepted reconnect alone — via
// HandleNodeConnect — re-activates the node's resource for the new generation and
// pumps the waiter to it, with no new request, config refresh, or lease release.
// It also pins that a stale/absent-generation connect callback is a no-op.
func TestAcceptedReconnectActivatesCandidateAndPumpsWaiter(t *testing.T) {
	const (
		nodeA      = "node-recon-a"
		nodeB      = "node-recon-b"
		provA      = "prov-recon-a"
		provB      = "prov-recon-b"
		modelGroup = "group-recon"
		servedA    = "served-recon-a"
		servedB    = "served-recon-b"
	)

	reg := edgenode.NewRegistry()
	entryA := &edgenode.NodeEntry{NodeID: nodeA}
	entryB := &edgenode.NodeEntry{NodeID: nodeB}
	reg.Register(entryA)
	reg.Register(entryB)
	genA := entryA.ConnectionGeneration
	if genA == 0 {
		t.Fatal("node A must carry a non-zero connection generation")
	}

	svc := New(reg, edgeevents.NewBus())
	store := edgenode.NewNodeStore()
	newRecord := func(nodeID, providerID, served string) *edgenode.NodeRecord {
		return &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{served}, Health: "available", Capacity: 1,
			}},
		}
	}
	store.Add(newRecord(nodeA, provA, servedA))
	store.Add(newRecord(nodeB, provB, servedB))
	catalog := []config.ModelCatalogEntry{{
		ID:        modelGroup,
		Providers: map[string]string{provA: servedA, provB: servedB},
	}}
	svc.SetRuntimeConfig(store, catalog, NewGroupPolicy(16, 3*time.Second))

	// Fill the alternate provider B so it cannot absorb the waiter. Hand-build the
	// candidate so the fill is forced onto B regardless of scheduler ordering.
	fillCand := []candidateNode{{
		entry: entryB, capacity: 1, providerID: provB,
		servedTarget: servedB, generation: entryB.ConnectionGeneration,
	}}
	filled, err := svc.queue.admit(t.Context(), modelGroup, "", "", fillCand, groupPolicy{}, nil, false, true)
	if err != nil || filled == nil {
		t.Fatalf("fill alternate provider B: %v", err)
	}
	if inFlight, _ := providerResourceCounts(svc.queue, nodeB, provB); inFlight != 1 {
		t.Fatalf("expected provB in-flight=1 after fill, got %d", inFlight)
	}

	// Disconnect node A (idle, holding no lease). The authoritative disconnect
	// marks provA's resource orphan/offline.
	reg.Unregister(nodeA)
	svc.HandleNodeDisconnect(nodeA, genA, "disconnected")
	present, orphan, _, _ := resourceStateForTest(svc.queue, nodeA, provA)
	if !present || !orphan {
		t.Fatalf("after disconnect provA resource present=%t orphan=%t, want true/true", present, orphan)
	}

	// A stale/absent-generation connect callback must not re-activate: node A is
	// no longer the registry owner (it is absent), so IsCurrentOwnerGeneration
	// fails and the callback is a no-op.
	svc.HandleNodeConnect(nodeA, genA)
	if _, stillOrphan, _, _ := resourceStateForTest(svc.queue, nodeA, provA); !stillOrphan {
		t.Fatal("stale connect callback re-activated the orphaned resource")
	}

	// Enqueue the waiter. At enqueue time node A is disconnected, so the only live
	// candidate is B, which is full — the request queues.
	req := SubmitRunRequest{ModelGroupKey: modelGroup, ProviderPool: true}
	type admitOutcome struct {
		candidate *candidateNode
		err       error
	}
	waiterCh := make(chan admitOutcome, 1)
	waiterCtx, cancelWaiter := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancelWaiter()
	go func() {
		enqCands, _, resolveErr := svc.resolveQueueCandidates(req)
		if resolveErr != nil {
			waiterCh <- admitOutcome{err: resolveErr}
			return
		}
		candidate, admitErr := svc.queue.admit(
			waiterCtx, modelGroup, "", "", enqCands, groupPolicy{},
			svc.resolveQueueCandidatesClosure(req), false, true,
		)
		waiterCh <- admitOutcome{candidate: candidate, err: admitErr}
	}()

	waitForQueueLen(t, svc.queue, modelGroup, 1)
	select {
	case outcome := <-waiterCh:
		t.Fatalf("waiter resolved before reconnect: candidate=%+v err=%v", outcome.candidate, outcome.err)
	default:
	}

	// Reconnect the same identity: a strictly higher generation is minted. The
	// accepted-connect hook alone must re-activate provA and dispatch the waiter —
	// no config refresh, no new request, no lease release is issued here.
	reconnectA := &edgenode.NodeEntry{NodeID: nodeA}
	reg.Register(reconnectA)
	newGen := reconnectA.ConnectionGeneration
	if newGen <= genA {
		t.Fatalf("reconnect generation %d must exceed the first %d", newGen, genA)
	}
	svc.HandleNodeConnect(nodeA, newGen)

	select {
	case outcome := <-waiterCh:
		if outcome.err != nil {
			t.Fatalf("waiter expected dispatch after reconnect, got error: %v", outcome.err)
		}
		if outcome.candidate == nil {
			t.Fatal("waiter expected a candidate after reconnect")
		}
		if outcome.candidate.providerID != provA || outcome.candidate.entry.NodeID != nodeA {
			t.Fatalf("waiter dispatched to (%q,%q), want (%q,%q)",
				outcome.candidate.entry.NodeID, outcome.candidate.providerID, nodeA, provA)
		}
		if outcome.candidate.generation != newGen {
			t.Fatalf("waiter candidate generation=%d, want the reconnect generation %d", outcome.candidate.generation, newGen)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("accepted reconnect did not pump the stranded waiter")
	}

	// The activation re-enabled provA for the reconnect generation and the pump
	// reserved it. The alternate B's lease was never released, proving the pump was
	// triggered by the reconnect alone.
	present, orphan, gen, inFlightA := resourceStateForTest(svc.queue, nodeA, provA)
	if !present || orphan {
		t.Fatalf("after reconnect provA present=%t orphan=%t, want true/false", present, orphan)
	}
	if gen != newGen {
		t.Fatalf("provA resource generation=%d, want %d", gen, newGen)
	}
	if inFlightA != 1 {
		t.Fatalf("provA in-flight=%d after waiter dispatch, want 1", inFlightA)
	}
	if inFlightB, _ := providerResourceCounts(svc.queue, nodeB, provB); inFlightB != 1 {
		t.Fatalf("provB in-flight=%d, want 1 (its lease must not have been released)", inFlightB)
	}
	_ = filled
}