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

package service

import (
	"fmt"
	"net"
	"runtime"
	"sync"
	"sync/atomic"
	"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"
)

func TestListNodeSnapshotsUsesEdgeQueueState(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-q-1", Alias: "node-q-1"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-q-1",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{
					Name:     "ollama-inst",
					Enabled:  true,
					Capacity: 3,
				},
			},
		},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:      "test-group",
		adapter:  "ollama-inst",
		inflight: map[string]int{"node-q-1": 2},
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: reg.All()[0], capacity: 3},
				},
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]
	if len(snap.ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
	}
	p := snap.ProviderSnapshots[0]
	if p.Adapter != "ollama-inst" {
		t.Errorf("expected adapter ollama-inst, got %s", p.Adapter)
	}
	if p.Capacity != 3 {
		t.Errorf("expected capacity 3, got %d", p.Capacity)
	}
	if p.InFlight != 2 {
		t.Errorf("expected inflight 2, got %d", p.InFlight)
	}
	if p.Queued != 1 {
		t.Errorf("expected queued 1, got %d", p.Queued)
	}
}

func TestListNodeSnapshotsTypeRouteSingleNamedInstance(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-q-2", Alias: "node-q-2"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-q-2",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{
					Name:     "ollama-local",
					Enabled:  true,
					Capacity: 5,
				},
			},
		},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:      "test-group",
		adapter:  "ollama",
		inflight: map[string]int{"node-q-2": 3},
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: reg.All()[0], capacity: 5},
				},
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]
	if len(snap.ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
	}
	p := snap.ProviderSnapshots[0]
	if p.Adapter != "ollama-local" {
		t.Errorf("expected adapter ollama-local, got %s", p.Adapter)
	}
	if p.Capacity != 5 {
		t.Errorf("expected capacity 5, got %d", p.Capacity)
	}
	if p.InFlight != 3 {
		t.Errorf("expected inflight 3, got %d", p.InFlight)
	}
	if p.Queued != 1 {
		t.Errorf("expected queued 1, got %d", p.Queued)
	}
}

func TestListNodeSnapshotsProviderEmptyAdapterFallbackToProviderID(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-e-1", Alias: "node-e-1"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:       "node-e-1",
		Adapters: config.AdaptersConf{},
		Providers: []config.NodeProviderConf{
			{
				ID:       "standalone-provider",
				Type:     "cli",
				Category: config.CategoryCLI,
				Adapter:  "",
				Models:   []string{"test-model"},
				Capacity: 2,
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 2},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: "node-e-1", providerID: "standalone-provider"}
	res := svc.queue.resources[key]
	if res == nil {
		res = &providerResourceState{
			nodeID:     "node-e-1",
			providerID: "standalone-provider",
			capacity:   2,
			enabled:    true,
		}
		svc.queue.resources[key] = res
	}
	res.inFlight = 1

	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:     "test-group",
		adapter: "standalone-provider",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: reg.All()[0], capacity: 2, providerID: "standalone-provider"},
				},
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]

	found := false
	for i, ps := range snap.ProviderSnapshots {
		if ps.Id == "standalone-provider" {
			found = true
			if ps.InFlight != 1 {
				t.Errorf("expected InFlight=1 for standalone-provider (adapter empty), got %d", ps.InFlight)
			}
			if ps.Queued != 1 {
				t.Errorf("expected Queued=1 for standalone-provider (adapter empty), got %d", ps.Queued)
			}
			if ps.Capacity != 2 {
				t.Errorf("expected Capacity=2 for standalone-provider, got %d", ps.Capacity)
			}
			expectedLoadRatio := float32(1) / float32(2)
			if ps.LoadRatio != expectedLoadRatio {
				t.Errorf("expected LoadRatio=%.2f for standalone-provider, got %.2f", expectedLoadRatio, ps.LoadRatio)
			}
			t.Logf("provider snapshot[%d]: id=%s, adapter=%s, InFlight=%d, Queued=%d, Capacity=%d, LoadRatio=%.2f",
				i, ps.Id, ps.Adapter, ps.InFlight, ps.Queued, ps.Capacity, ps.LoadRatio)
			break
		}
	}
	if !found {
		t.Fatalf("expected provider snapshot with id=standalone-provider, got %d provider snapshots", len(snap.ProviderSnapshots))
	}
}

func TestListNodeSnapshotsProviderCatalogSupersedesAdapterSnapshots(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cs-1", Alias: "node-cs-1"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-cs-1",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{Enabled: true},
			OpenAICompatInstances: []config.OpenAICompatInstanceConf{
				{Name: "vllm-gpu", Enabled: true},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-cat-1", Adapter: "vllm-gpu", Models: []string{"model-x"}, Health: "available", Capacity: 4},
			{ID: "prov-cat-2", Adapter: "vllm-gpu", Models: []string{"model-y"}, Health: "available", Capacity: 2},
		},
		Runtime: config.RuntimeConf{Concurrency: 1},
	}
	store.Add(rec)

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

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]

	if len(snap.ProviderSnapshots) != 2 {
		t.Fatalf("expected 2 catalog snapshots, got %d: %+v",
			len(snap.ProviderSnapshots), snap.ProviderSnapshots)
	}
	ids := map[string]bool{}
	for _, p := range snap.ProviderSnapshots {
		if p.Id == "" {
			t.Errorf("unexpected adapter-only snapshot (empty Id): %+v", p)
		}
		ids[p.Id] = true
	}
	if !ids["prov-cat-1"] || !ids["prov-cat-2"] {
		t.Errorf("expected catalog IDs prov-cat-1 and prov-cat-2, got: %v", ids)
	}
}

func TestListNodeSnapshotsCatalogProviderZeroCapacityDoesNotUseRuntimeFallback(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-zero-cap", Alias: "node-zero-cap"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-zero-cap",
		Adapters: config.AdaptersConf{
			OpenAICompatInstances: []config.OpenAICompatInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Capacity: 7},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-zero", Adapter: "vllm-gpu", Models: []string{"model-zero"}, Health: "available", Capacity: 0},
		},
		Runtime: config.RuntimeConf{Concurrency: 7},
	}
	store.Add(rec)

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

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	if len(snaps[0].ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snaps[0].ProviderSnapshots))
	}
	p := snaps[0].ProviderSnapshots[0]
	if p.Id != "prov-zero" {
		t.Fatalf("expected provider id prov-zero, got %q", p.Id)
	}
	if p.Capacity != 0 {
		t.Fatalf("provider catalog capacity must not fall back to runtime concurrency: got %d", p.Capacity)
	}
	if p.LoadRatio != 0 {
		t.Fatalf("expected zero load ratio for zero capacity provider, got %.2f", p.LoadRatio)
	}
}

func TestStaticProviderCatalogSnapshotsZeroCapacityDoesNotUseRuntimeFallback(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID: "node-static-zero-cap",
		Providers: []config.NodeProviderConf{
			{ID: "prov-zero", Adapter: "vllm-gpu", Models: []string{"model-zero"}, Health: "available", Capacity: 0},
		},
		Runtime: config.RuntimeConf{Concurrency: 7},
	}

	snaps := staticProviderCatalogSnapshots(rec)
	if len(snaps) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snaps))
	}
	if snaps[0].Capacity != 0 {
		t.Fatalf("static provider catalog capacity must not fall back to runtime concurrency: got %d", snaps[0].Capacity)
	}
}

func TestListNodeSnapshotsLegacyAdapterFallbackWhenNoProviders(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-legacy-snap", Alias: "node-legacy-snap"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-legacy-snap",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{Enabled: true},
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, Capacity: 3},
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 2},
	}
	store.Add(rec)

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

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]

	if len(snap.ProviderSnapshots) != 2 {
		t.Fatalf("expected 2 adapter snapshots (cli + ollama-local), got %d: %+v",
			len(snap.ProviderSnapshots), snap.ProviderSnapshots)
	}
	adapters := map[string]bool{}
	for _, p := range snap.ProviderSnapshots {
		adapters[p.Adapter] = true
	}
	if !adapters["cli"] || !adapters["ollama-local"] {
		t.Errorf("expected cli and ollama-local adapter snapshots, got: %v", adapters)
	}
}

func TestListNodeSnapshotsProviderIdDiffersFromAdapterKey(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-p-1", Alias: "node-p-1"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-p-1",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Capacity: 4},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "provider-vllm-primary",
				Type:     "vllm",
				Category: config.CategoryAPI,
				Adapter:  "vllm-gpu",
				Models:   []string{"nvidia/Qwen3.6-35B"},
				Capacity: 4,
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 2},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: "node-p-1", providerID: "provider-vllm-primary"}
	res := svc.queue.resources[key]
	if res == nil {
		res = &providerResourceState{
			nodeID:     "node-p-1",
			providerID: "provider-vllm-primary",
			capacity:   4,
			enabled:    true,
		}
		svc.queue.resources[key] = res
	}
	res.inFlight = 2

	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:     "test-group",
		adapter: "vllm-gpu",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: reg.All()[0], capacity: 4, providerID: "provider-vllm-primary"},
				},
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]

	var p *config.NodeProviderConf
	for i := range snap.ProviderSnapshots {
		ps := snap.ProviderSnapshots[i]
		if ps.Id == "provider-vllm-primary" {
			p = &config.NodeProviderConf{
				Capacity: int(ps.Capacity),
			}
			if ps.InFlight != 2 {
				t.Errorf("expected InFlight=2 for provider-vllm-primary, got %d", ps.InFlight)
			}
			if ps.Queued != 1 {
				t.Errorf("expected Queued=1 for provider-vllm-primary, got %d", ps.Queued)
			}
			if ps.Adapter != "vllm-gpu" {
				t.Errorf("expected Adapter=vllm-gpu for provider-vllm-primary, got %s", ps.Adapter)
			}
			break
		}
	}
	if p == nil {
		t.Fatalf("expected provider snapshot with id=provider-vllm-primary, got %d provider snapshots", len(snap.ProviderSnapshots))
	}
}

func TestListNodeSnapshotsLongContextFields(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-long-1", Alias: "node-long-1"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-long-1",
		Providers: []config.NodeProviderConf{
			{
				ID:                  "long-provider",
				Type:                "vllm",
				Category:            config.CategoryAPI,
				Adapter:             "vllm-gpu",
				Models:              []string{"long-model"},
				Capacity:            4,
				LongContextCapacity: 2,
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 2},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: "node-long-1", providerID: "long-provider"}
	res := svc.queue.resources[key]
	if res == nil {
		res = &providerResourceState{
			nodeID:       "node-long-1",
			providerID:   "long-provider",
			capacity:     4,
			longCapacity: 2,
			enabled:      true,
		}
		svc.queue.resources[key] = res
	}
	res.inFlight = 1
	res.longInFlight = 1

	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:     "test-group",
		adapter: "vllm-gpu",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: reg.All()[0], capacity: 4, providerID: "long-provider"},
				},
				long: true,
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]
	if len(snap.ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
	}
	p := snap.ProviderSnapshots[0]
	if p.Id != "long-provider" {
		t.Errorf("expected long-provider, got %s", p.Id)
	}
	if p.LongContextCapacity != 2 {
		t.Errorf("expected LongContextCapacity 2, got %d", p.LongContextCapacity)
	}
	if p.LongInFlight != 1 {
		t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
	}
	if p.LongQueued != 1 {
		t.Errorf("expected LongQueued 1, got %d", p.LongQueued)
	}
}

func TestProviderSnapshotQueuedCountsCandidatePressureAcrossModelGroups(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{NodeID: "node-pressure", Alias: "node-pressure"}
	reg.Register(entry)

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-pressure",
		Providers: []config.NodeProviderConf{
			{
				ID:                  "provider-shared",
				Type:                "vllm",
				Category:            config.CategoryAPI,
				Models:              []string{"model-a", "model-b"},
				Health:              "available",
				Capacity:            1,
				LongContextCapacity: 1,
			},
			{
				ID:                  "provider-fallback",
				Type:                "vllm",
				Category:            config.CategoryAPI,
				Models:              []string{"model-a"},
				Health:              "available",
				Capacity:            1,
				LongContextCapacity: 1,
			},
		},
	}
	store.Add(rec)

	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)

	svc.queue.mu.Lock()
	svc.queue.groups["model-a"] = &modelQueueGroup{
		key: "model-a",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: entry, providerID: "provider-shared"},
					{entry: entry, providerID: "provider-fallback"},
				},
			},
		},
	}
	svc.queue.groups["model-b"] = &modelQueueGroup{
		key: "model-b",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: entry, providerID: "provider-shared"},
				},
				long: true,
			},
		},
	}
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	if len(snaps[0].ProviderSnapshots) != 2 {
		t.Fatalf("expected 2 provider snapshots, got %d", len(snaps[0].ProviderSnapshots))
	}

	byID := make(map[string]*iop.ProviderSnapshot, len(snaps[0].ProviderSnapshots))
	for _, snapshot := range snaps[0].ProviderSnapshots {
		byID[snapshot.GetId()] = snapshot
	}

	shared := byID["provider-shared"]
	if shared == nil {
		t.Fatal("provider-shared snapshot not found")
	}
	if shared.GetQueued() != 2 || shared.GetLongQueued() != 1 {
		t.Fatalf("shared candidate pressure mismatch: queued=%d long_queued=%d", shared.GetQueued(), shared.GetLongQueued())
	}

	fallback := byID["provider-fallback"]
	if fallback == nil {
		t.Fatal("provider-fallback snapshot not found")
	}
	if fallback.GetQueued() != 1 || fallback.GetLongQueued() != 0 {
		t.Fatalf("fallback candidate pressure mismatch: queued=%d long_queued=%d", fallback.GetQueued(), fallback.GetLongQueued())
	}

	if total := shared.GetQueued() + fallback.GetQueued(); total != 3 {
		t.Fatalf("provider candidate pressure should be non-additive across 2 Edge pending requests: got total=%d", total)
	}
}

func TestProviderSnapshotReadsResourceState(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-snap-test", Alias: "node-snap-test"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:      "node-snap-test",
		Runtime: config.RuntimeConf{Concurrency: 5},
		Providers: []config.NodeProviderConf{
			{
				ID:                  "snap-prov",
				Adapter:             "vllm",
				Capacity:            3,
				LongContextCapacity: 2,
				Enabled:             nil,
			},
		},
	}
	store.Add(rec)

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

	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: "node-snap-test", providerID: "snap-prov"}
	res := svc.queue.resources[key]
	if res == nil {
		svc.queue.mu.Unlock()
		t.Fatal("resource state not found after SetNodeStore")
	}
	res.inFlight = 2
	res.longInFlight = 1
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
	}
	snap := snaps[0]
	if len(snap.ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
	}
	p := snap.ProviderSnapshots[0]
	if p.InFlight != 2 {
		t.Errorf("expected InFlight 2, got %d", p.InFlight)
	}
	if p.LongInFlight != 1 {
		t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
	}
}

func TestProviderSnapshotIgnoresLegacyFallbackState(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-fallback-test", Alias: "node-fallback-test"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:      "node-fallback-test",
		Runtime: config.RuntimeConf{Concurrency: 5},
		Providers: []config.NodeProviderConf{
			{
				ID:                  "fallback-prov",
				Adapter:             "vllm",
				Capacity:            3,
				LongContextCapacity: 2,
				Enabled:             nil,
			},
		},
	}
	store.Add(rec)

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

	trackTestLease(svc.queue, "test-group", "run-legacy", "node-fallback-test", "fallback-prov", true)

	svc.queue.mu.Lock()
	svc.queue.groups["test-group"] = &modelQueueGroup{
		key:          "test-group",
		adapter:      "vllm",
		inflight:     map[string]int{"node-fallback-test:fallback-prov": 1},
		longInflight: map[string]int{"node-fallback-test:fallback-prov": 1},
	}
	key := providerResourceKey{nodeID: "node-fallback-test", providerID: "fallback-prov"}
	res := svc.queue.resources[key]
	if res == nil {
		svc.queue.mu.Unlock()
		t.Fatal("resource state not found")
	}
	res.inFlight = 0
	res.longInFlight = 0
	svc.queue.mu.Unlock()

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot")
	}
	p := snaps[0].ProviderSnapshots[0]
	if p.InFlight != 0 {
		t.Errorf("expected InFlight 0, got %d", p.InFlight)
	}
	if p.LongInFlight != 0 {
		t.Errorf("expected LongInFlight 0, got %d", p.LongInFlight)
	}
}

// findProviderSnapshotOnNode finds a provider snapshot by nodeID and providerID
// from the production ListNodeSnapshots output. It is used by lifecycle tests
// that assert provider state through the actual snapshot surface rather than
// reading internal queue/resource maps directly.
func findProviderSnapshotOnNode(snaps []NodeSnapshot, nodeID, providerID string) *iop.ProviderSnapshot {
	for _, snap := range snaps {
		if snap.NodeID == nodeID {
			for _, ps := range snap.ProviderSnapshots {
				if ps.Id == providerID {
					return ps
				}
			}
		}
	}
	return nil
}

// assertProviderSnapshotState asserts individual provider snapshot fields
// read from production ListNodeSnapshots. All four counters are independent so
// that mixed-field regressions (e.g., inFlight set but queued not cleared) are
// caught. The production surface is the only assertion channel: the internal
// queue/resource maps are implementation details.
func assertProviderSnapshotState(t *testing.T, snaps []NodeSnapshot, nodeID, providerID, label string,
	inFlight, queued, longInFlight, longQueued int32) {
	t.Helper()
	ps := findProviderSnapshotOnNode(snaps, nodeID, providerID)
	if ps == nil {
		t.Fatalf("%s: provider %s snapshot not found on %s in ListNodeSnapshots", label, providerID, nodeID)
	}
	if ps.GetId() != providerID {
		t.Errorf("%s %s id mismatch: want %s, got %s", label, providerID, providerID, ps.GetId())
	}
	if ps.GetInFlight() != inFlight {
		t.Errorf("%s %s inFlight: want %d, got %d", label, providerID, inFlight, ps.GetInFlight())
	}
	if ps.GetQueued() != queued {
		t.Errorf("%s %s queued: want %d, got %d", label, providerID, queued, ps.GetQueued())
	}
	if ps.GetLongInFlight() != longInFlight {
		t.Errorf("%s %s longInFlight: want %d, got %d", label, providerID, longInFlight, ps.GetLongInFlight())
	}
	if ps.GetLongQueued() != longQueued {
		t.Errorf("%s %s longQueued: want %d, got %d", label, providerID, longQueued, ps.GetLongQueued())
	}
}

// snapshotTuple describes the exact provider snapshot tuple (status, configured
// capacity, in-flight, queued, long-context in-flight, long-context queued).
// Used by the pure validateSnapshotTuple so every concurrent reader compares
// the same canonical field set without touching *testing.T.
type snapshotTuple struct {
	status       string
	capacity     int32
	inFlight     int32
	queued       int32
	longInFlight int32
	longQueued   int32
}

// validateSnapshotTuple compares a snapshotTuple against an expected tuple and
// returns nil on exact match or an error listing the first mismatched field.
// Pure function: does not touch *testing.T so storm readers can collect
// failures through a bounded channel without blocking on t.Fatal.
func validateSnapshotTuple(got, want snapshotTuple) error {
	if got.status != want.status {
		return fmt.Errorf("status: want %q, got %q", want.status, got.status)
	}
	if got.capacity != want.capacity {
		return fmt.Errorf("capacity: want %d, got %d", want.capacity, got.capacity)
	}
	if got.inFlight != want.inFlight {
		return fmt.Errorf("inFlight: want %d, got %d", want.inFlight, got.inFlight)
	}
	if got.queued != want.queued {
		return fmt.Errorf("queued: want %d, got %d", want.queued, got.queued)
	}
	if got.longInFlight != want.longInFlight {
		return fmt.Errorf("longInFlight: want %d, got %d", want.longInFlight, got.longInFlight)
	}
	if got.longQueued != want.longQueued {
		return fmt.Errorf("longQueued: want %d, got %d", want.longQueued, got.longQueued)
	}
	return nil
}

// snapshotFailure carries one reader's first validation error for bounded
// delivery through a snapReaders-sized channel during the concurrent storm.
// The main goroutine owns all t.Fatal/Error calls.
type snapshotFailure struct {
	reader int
	err    error
}

// exactProviderTuples builds a provider-ID → snapshotTuple map from a node's
// snapshot, rejecting missing or duplicate providers with an error. Used by the
// concurrent storm reader so each expected provider is proven present and only
// once before any tuple comparison.
func exactProviderTuples(snaps []NodeSnapshot, nodeID string, expectedIDs ...string) (map[string]snapshotTuple, error) {
	var nodeSnap *NodeSnapshot
	for i := range snaps {
		if snaps[i].NodeID == nodeID {
			nodeSnap = &snaps[i]
			break
		}
	}
	if nodeSnap == nil {
		return nil, fmt.Errorf("node %s not found in ListNodeSnapshots", nodeID)
	}
	seen := make(map[string]snapshotTuple, len(expectedIDs))
	for _, ps := range nodeSnap.ProviderSnapshots {
		pid := ps.GetId()
		if pid == "" {
			continue
		}
		t := snapshotTuple{
			status:       ps.GetStatus(),
			capacity:     ps.GetCapacity(),
			inFlight:     ps.GetInFlight(),
			queued:       ps.GetQueued(),
			longInFlight: ps.GetLongInFlight(),
			longQueued:   ps.GetLongQueued(),
		}
		if _, dup := seen[pid]; dup {
			return nil, fmt.Errorf("duplicate provider %s on %s", pid, nodeID)
		}
		seen[pid] = t
	}
	for _, wantID := range expectedIDs {
		if _, ok := seen[wantID]; !ok {
			return nil, fmt.Errorf("expected provider %s missing on %s", wantID, nodeID)
		}
	}
	return seen, nil
}

// reportFirstSnapshotFailure sends err through errCh as a reader's first
// validation failure. Non-blocking: if the channel is full, the reader exits
// silently so a stuck channel never deadlocks the storm.
func reportFirstSnapshotFailure(errCh chan<- snapshotFailure, readerID int, err error) {
	select {
	case errCh <- snapshotFailure{reader: readerID, err: err}:
	default:
	}
}

// TestProviderSnapshotLiveCandidatePressureFollowsRefreshAcrossModelGroups
// verifies that a provider snapshot reflects live candidate pressure from actual
// SubmitRun waiter dispatches across model groups, covering both the normal and
// long-context lifecycle (enqueue, dispatch, terminal zero). The holder and
// waiter go through real admission → pump → terminal lifecycle, not direct
// queue injection.
func TestProviderSnapshotLiveCandidatePressureFollowsRefreshAcrossModelGroups(t *testing.T) {
	for _, tc := range []struct {
		name         string
		contextClass string
		wantLong     int
	}{
		{name: "normal"},
		{name: "long", contextClass: contextClassLong, wantLong: 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			parserMap := toki.ParserMap{
				toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.RunRequest{}
					return m, proto.Unmarshal(b, m)
				},
			}

			edgeConn, nodeConn := net.Pipe()
			defer edgeConn.Close()
			defer nodeConn.Close()
			edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
			_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)

			const nodeID = "node-lifecycle"
			reg := edgenode.NewRegistry()
			reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
			bus := edgeevents.NewBus()
			svc := New(reg, bus)

			catalog := []config.ModelCatalogEntry{
				{
					ID:        "model-a",
					Providers: map[string]string{"prov-shared": "served-a"},
				},
				{
					ID:        "model-b",
					Providers: map[string]string{"prov-fallback": "served-b"},
				},
				{
					ID: "model-wait",
					Providers: map[string]string{
						"prov-shared":   "served-shared",
						"prov-fallback": "served-fallback",
					},
				},
			}

			makeStore := func(longCap int) *edgenode.NodeStore {
				store := edgenode.NewNodeStore()
				store.Add(&edgenode.NodeRecord{
					ID: nodeID,
					Adapters: config.AdaptersConf{
						VllmInstances: []config.VllmInstanceConf{
							{Name: "vllm-gpu", Enabled: true},
						},
					},
					Providers: []config.NodeProviderConf{
						{
							ID: "prov-shared", Type: "vllm", Category: config.CategoryAPI,
							Adapter: "vllm-gpu", Models: []string{"served-a", "served-shared"},
							Health: "available", Capacity: 1, LongContextCapacity: longCap,
						},
						{
							ID: "prov-fallback", Type: "vllm", Category: config.CategoryAPI,
							Adapter: "vllm-gpu", Models: []string{"served-b", "served-fallback"},
							Health: "available", Capacity: 1, LongContextCapacity: longCap,
						},
					},
				})
				return store
			}

			svc.SetRuntimeConfig(makeStore(tc.wantLong), catalog, groupPolicy{})

			// 1. Holder A for model-a → dispatched to prov-shared immediately.
			holderA, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
				ModelGroupKey: "model-a", ProviderPool: true, Background: true,
				ContextClass: tc.contextClass,
			})
			if err != nil || holderA == nil {
				t.Fatalf("holder A submit: %v", err)
			}
			hA := holderA.Dispatch()
			if got := hA.ProviderID; got != "prov-shared" {
				t.Fatalf("holder A expected prov-shared, got %s", got)
			}

			// 2. Holder B for model-b → dispatched to prov-fallback immediately.
			holderB, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
				ModelGroupKey: "model-b", ProviderPool: true, Background: true,
				ContextClass: tc.contextClass,
			})
			if err != nil || holderB == nil {
				t.Fatalf("holder B submit: %v", err)
			}
			hB := holderB.Dispatch()
			if got := hB.ProviderID; got != "prov-fallback" {
				t.Fatalf("holder B expected prov-fallback, got %s", got)
			}

			// Snapshot while both holders are in-flight and no waiter: queued must be 0.
			snaps := svc.ListNodeSnapshots()
			assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step1-holders",
				1, 0, int32(tc.wantLong), 0)
			assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step1-holders",
				1, 0, int32(tc.wantLong), 0)

			// 3. Waiter for model-wait → queues because both providers are full.
			type submitResult struct {
				run RunResult
				err error
			}
			resultCh := make(chan submitResult, 1)
			go func() {
				run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
					ModelGroupKey: "model-wait", ProviderPool: true, Background: true,
					ContextClass: tc.contextClass,
				})
				resultCh <- submitResult{run: run, err: err}
			}()

			// Poll until the waiter is enqueued.
			deadline := time.After(2 * time.Second)
			waiterQueued := false
			for {
				svc.queue.mu.Lock()
				group := svc.queue.groups["model-wait"]
				count := 0
				if group != nil {
					count = len(group.queue)
				}
				svc.queue.mu.Unlock()
				if count == 1 {
					waiterQueued = true
					break
				}
				select {
				case <-deadline:
					t.Fatal("waiter was not enqueued within timeout")
				default:
					runtime.Gosched()
				}
			}
			if !waiterQueued {
				t.Fatal("waiter not enqueued")
			}

			// Snapshot while waiter is queued: each provider has in_flight=1
			// from its holder AND queued=1 from the waiter's candidate.
			snaps = svc.ListNodeSnapshots()
			assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step2-queued",
				1, 1, int32(tc.wantLong), int32(tc.wantLong))
			assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step2-queued",
				1, 1, int32(tc.wantLong), int32(tc.wantLong))

			// 4. Release holder B (fallback) → waiter dispatches to fallback.
			svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: hB.RunID, Type: "complete"})

			var waiterRes submitResult
			select {
			case waiterRes = <-resultCh:
			case <-time.After(2 * time.Second):
				t.Fatal("waiter did not dispatch after fallback release")
			}
			if waiterRes.err != nil {
				t.Fatalf("waiter submit error: %v", waiterRes.err)
			}
			if waiterRes.run == nil {
				t.Fatal("waiter expected non-nil result")
			}
			defer waiterRes.run.Close()
			wDispatch := waiterRes.run.Dispatch()
			if got := wDispatch.ProviderID; got != "prov-fallback" {
				t.Fatalf("waiter expected dispatch to prov-fallback, got %s", got)
			}

			// Snapshot after dispatch.
			snaps = svc.ListNodeSnapshots()
			assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step3-dispatched",
				1, 0, int32(tc.wantLong), 0)
			assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step3-dispatched",
				1, 0, int32(tc.wantLong), 0)

			// 5. Terminal for holder A (shared).
			svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: hA.RunID, Type: "complete"})
			snaps = svc.ListNodeSnapshots()
			assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step4-shared-terminal",
				0, 0, 0, 0)
			assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step4-shared-terminal",
				1, 0, int32(tc.wantLong), 0)

			// 6. Terminal for waiter → full drain.
			svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})
			snaps = svc.ListNodeSnapshots()
			assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step5-all-terminal",
				0, 0, 0, 0)
			assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step5-all-terminal",
				0, 0, 0, 0)

			// Final counters and lease must be fully zero.
			aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-shared")
			bInFlight, bLong := providerResourceCounts(svc.queue, nodeID, "prov-fallback")
			if aInFlight != 0 || aLong != 0 {
				t.Fatalf("expected prov-shared drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
			}
			if bInFlight != 0 || bLong != 0 {
				t.Fatalf("expected prov-fallback drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
			}
			if lc := leaseCount(svc.queue); lc != 0 {
				t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
			}
		})
	}
}

// TestProviderSnapshotConcurrentRefreshPressureUsesLiveCandidates drives
// concurrent runtime refresh and snapshot reads against a queued waiter.
// Every snapshot tuple is asserted to be a fully-consistent old-or-new state:
// enabled providers must have (status=available, configured capacity, in-flight
// from current lease, queued from live candidates); disabled providers must have
// every effective counter and capacity zero. Each reader records at most one
// invalid-tuple failure so the storm cannot block on the bounded channel.
func TestProviderSnapshotConcurrentRefreshPressureUsesLiveCandidates(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	const nodeID = "node-concurrent"
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
	bus := edgeevents.NewBus()
	svc := New(reg, bus)

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "model-a",
			Providers: map[string]string{"prov-shared": "served-a"},
		},
		{
			ID:        "model-b",
			Providers: map[string]string{"prov-fallback": "served-b"},
		},
		{
			ID: "model-conc",
			Providers: map[string]string{
				"prov-shared":   "served-shared",
				"prov-fallback": "served-fallback",
			},
		},
	}

	makeStore := func(fallbackEnabled bool) *edgenode.NodeStore {
		enabledFlag := fallbackEnabled
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: nodeID,
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID: "prov-shared", Type: "vllm", Category: config.CategoryAPI,
					Adapter: "vllm-gpu", Models: []string{"served-a", "served-shared"},
					Health: "available", Capacity: 1,
				},
				{
					ID: "prov-fallback", Type: "vllm", Category: config.CategoryAPI,
					Adapter: "vllm-gpu", Models: []string{"served-b", "served-fallback"},
					Health: "available", Capacity: 1, Enabled: &enabledFlag,
				},
			},
		})
		return store
	}

	enabledStore := makeStore(true)
	concDisabledStore := makeStore(false)

	svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})

	holderA, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "model-a", ProviderPool: true, Background: true,
	})
	if err != nil || holderA == nil {
		t.Fatalf("holder A submit: %v", err)
	}
	defer holderA.Close()

	holderB, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "model-b", ProviderPool: true, Background: true,
	})
	if err != nil || holderB == nil {
		t.Fatalf("holder B submit: %v", err)
	}
	defer holderB.Close()

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

	waitForQueueLen(t, svc.queue, "model-conc", 1)

	// Snapshot while queued with both enabled — using production ListNodeSnapshots
	// so the assertion exercises the actual snapshot surface.
	snaps := svc.ListNodeSnapshots()
	assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-pre-storm",
		1, 1, 0, 0)
	assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-pre-storm",
		1, 1, 0, 0)

	// Expected tuples for exact concurrent storm validation.
	// prov-shared is always enabled: capacity=1, inFlight=1 (holder A), queued=1 (waiter candidate).
	expectedShared := snapshotTuple{
		status: "available", capacity: 1, inFlight: 1, queued: 1,
		longInFlight: 0, longQueued: 0,
	}
	// prov-fallback alternates: enabled (same tuple as shared) or disabled (all zeros).
	expectedFallbackEnabled := snapshotTuple{
		status: "available", capacity: 1, inFlight: 1, queued: 1,
		longInFlight: 0, longQueued: 0,
	}
	expectedFallbackDisabled := snapshotTuple{
		status: "disabled", capacity: 0, inFlight: 0, queued: 0,
		longInFlight: 0, longQueued: 0,
	}

	// --- Concurrent storm ---
	// Bounded failure channel sized to reader count: each reader delivers at
	// most one first-error, so the channel never fills and the storm cannot
	// deadlock on send. Main goroutine owns all t.Fatal/Error.
	const snapReaders = 4
	errCh := make(chan snapshotFailure, snapReaders)

	start := make(chan struct{})
	stop := make(chan struct{})
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		<-start
		for {
			select {
			case <-stop:
				return
			default:
				svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
				svc.SetRuntimeConfig(concDisabledStore, catalog, groupPolicy{})
			}
		}
	}()

	var (
		fallbackSeenAvailable atomic.Bool
		fallbackSeenDisabled  atomic.Bool
	)

	for r := 0; r < snapReaders; r++ {
		wg.Add(1)
		go func(readerID int) {
			defer wg.Done()
			<-start
			for {
				select {
				case <-stop:
					return
				default:
				}
				snaps := svc.ListNodeSnapshots()
				if len(snaps) == 0 {
					runtime.Gosched()
					continue
				}

				byID, err := exactProviderTuples(snaps, nodeID, "prov-shared", "prov-fallback")
				if err != nil {
					reportFirstSnapshotFailure(errCh, readerID, err)
					return
				}
				if err := validateSnapshotTuple(byID["prov-shared"], expectedShared); err != nil {
					reportFirstSnapshotFailure(errCh, readerID, err)
					return
				}
				switch fb := byID["prov-fallback"]; {
				case validateSnapshotTuple(fb, expectedFallbackEnabled) == nil:
					fallbackSeenAvailable.Store(true)
				case validateSnapshotTuple(fb, expectedFallbackDisabled) == nil:
					fallbackSeenDisabled.Store(true)
				default:
					reportFirstSnapshotFailure(errCh, readerID, fmt.Errorf(
						"fallback tuple mismatch: got status=%q capacity=%d inflight=%d queued=%d",
						fb.status, fb.capacity, fb.inFlight, fb.queued))
					return
				}
			}
		}(r)
	}

	close(start)
	time.Sleep(300 * time.Millisecond)
	close(stop)
	wg.Wait()
	close(errCh)

	// Main goroutine owns all t.Fatal/Error: each reader delivered at most one
	// first-error through the bounded channel, so this loop never blocks.
	for f := range errCh {
		t.Errorf("reader %d: %v", f.reader, f.err)
	}

	// Fallback must have been observed in both enabled and disabled states during
	// the storm — a torn-tuple regression that silently passed every fallback read
	// as one state would leave one of the flags unset.
	if !fallbackSeenAvailable.Load() {
		t.Fatal("concurrent storm never sampled fallback in available state")
	}
	if !fallbackSeenDisabled.Load() {
		t.Fatal("concurrent storm never sampled fallback in disabled state")
	}

	// --- Post-storm ---
	svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
	snaps = svc.ListNodeSnapshots()
	assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-post-storm",
		1, 1, 0, 0)
	assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-post-storm",
		1, 1, 0, 0)

	// --- Dispatch the waiter via holder release ---
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holderB.Dispatch().RunID, Type: "complete"})

	var waiterRes submitResult
	select {
	case waiterRes = <-resultCh:
	case <-time.After(2 * time.Second):
		t.Fatal("waiter did not dispatch after holder release")
	}
	if waiterRes.err != nil {
		t.Fatalf("waiter submit error: %v", waiterRes.err)
	}
	if waiterRes.run == nil {
		t.Fatal("waiter expected non-nil result")
	}
	defer waiterRes.run.Close()
	wDispatch := waiterRes.run.Dispatch()
	if got := wDispatch.ProviderID; got != "prov-fallback" {
		t.Fatalf("waiter expected dispatch to prov-fallback, got %s", got)
	}

	snaps = svc.ListNodeSnapshots()
	assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-after-dispatch",
		1, 0, 0, 0)
	assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-after-dispatch",
		1, 0, 0, 0)

	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holderA.Dispatch().RunID, Type: "complete"})
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})

	aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-shared")
	bInFlight, bLong := providerResourceCounts(svc.queue, nodeID, "prov-fallback")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected prov-shared drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if bInFlight != 0 || bLong != 0 {
		t.Fatalf("expected prov-fallback drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
	}
}

// TestProviderSnapshotResolverErrorDoesNotUseStaleCandidate verifies that when
// a queued item's live resolver returns an error or an empty slice, the
// production ListNodeSnapshots surface does not fall back to the enqueue-time
// candidate list. The resolver path is exercised by setting resolveCandidates
// on an actual queued item, then asserting the snapshot reports zero candidate
// pressure (queued=0, longQueued=0) for the provider the item targets.
//
// This regression test exists because a previous loop removed it: with a
// regular expression test list, a zero-match run exits successfully, so the
// missing test was silently passing. Both error and empty cases are covered
// here so the resolver contract (no stale enqueue-time snapshot fallback) is
// enforced.
func TestProviderSnapshotResolverErrorDoesNotUseStaleCandidate(t *testing.T) {
	const nodeID = "node-resolver-stale"
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Alias: nodeID})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: nodeID,
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-shared",
				Type:     "vllm",
				Category: config.CategoryAPI,
				Adapter:  "vllm-gpu",
				Models:   []string{"served-a"},
				Health:   "available",
				Capacity: 1,
			},
		},
	})

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

	// Pre-fill provider resource state to simulate an in-flight holder already
	// consuming the slot. The queued waiter's pressure must not show up in the
	// snapshot if its resolver is broken.
	svc.queue.mu.Lock()
	key := providerResourceKey{nodeID: nodeID, providerID: "prov-shared"}
	res := svc.queue.resources[key]
	if res == nil {
		res = &providerResourceState{
			nodeID:     nodeID,
			providerID: "prov-shared",
			capacity:   1,
			enabled:    true,
		}
		svc.queue.resources[key] = res
	}
	res.inFlight = 1

	entry := reg.All()[0]
	svc.queue.groups["model-a"] = &modelQueueGroup{
		key: "model-a",
		queue: []*queueItem{
			{
				candidates: []candidateNode{
					{entry: entry, providerID: "prov-shared", capacity: 1},
				},
				waitCh: make(chan admitResult, 1),
			},
		},
	}
	svc.queue.mu.Unlock()

	assertZeroPressure := func(t *testing.T, label string) {
		t.Helper()
		snaps := svc.ListNodeSnapshots()
		ps := findProviderSnapshotOnNode(snaps, nodeID, "prov-shared")
		if ps == nil {
			t.Fatalf("%s: provider snapshot not found in ListNodeSnapshots", label)
		}
		if got := ps.GetQueued(); got != 0 {
			t.Errorf("%s: expected queued=0, got %d", label, got)
		}
		if got := ps.GetLongQueued(); got != 0 {
			t.Errorf("%s: expected longQueued=0, got %d", label, got)
		}
	}

	// Case 1: resolver returns error.
	t.Run("resolverError", func(t *testing.T) {
		svc.queue.mu.Lock()
		wg := svc.queue.groups["model-a"]
		if len(wg.queue) == 1 {
			wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
				return nil, fmt.Errorf("resolver error")
			}
		}
		svc.queue.mu.Unlock()
		assertZeroPressure(t, "resolverError")
	})

	// Case 2: resolver returns empty slice.
	t.Run("resolverEmpty", func(t *testing.T) {
		svc.queue.mu.Lock()
		wg := svc.queue.groups["model-a"]
		if len(wg.queue) == 1 {
			wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
				return []candidateNode{}, nil
			}
		}
		svc.queue.mu.Unlock()
		assertZeroPressure(t, "resolverEmpty")
	})

	// Case 3: resolver works normally — pressure should be reported.
	t.Run("resolverOkHasPressure", func(t *testing.T) {
		svc.queue.mu.Lock()
		wg := svc.queue.groups["model-a"]
		if len(wg.queue) == 1 {
			wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
				return []candidateNode{{entry: entry, providerID: "prov-shared", capacity: 1}}, nil
			}
		}
		svc.queue.mu.Unlock()
		snaps := svc.ListNodeSnapshots()
		ps := findProviderSnapshotOnNode(snaps, nodeID, "prov-shared")
		if ps == nil {
			t.Fatal("resolverOk: provider snapshot not found")
		}
		if got := ps.GetQueued(); got != 1 {
			t.Errorf("resolverOk: expected queued=1, got %d", got)
		}
	})
}

// TestProviderSnapshotRuntimeRefreshIsOldOrNew forces a queue-held writer to
// update the runtime config while a reader holds the queue lock, then asserts
// that the returned snapshot is either the complete old tuple or the complete
// new tuple — never a torn mix such as available/1/0 between an old available
// (capacity=1, queued=1) and a new disabled (capacity=0, queued=0) state. The
// ordering is enforced by having the writer block until the reader holds the
// queue lock, then releasing; the reader must capture a single consistent
// snapshot under the queue critical section established in
// SetRuntimeConfig → ListNodeSnapshots.
func TestProviderSnapshotRuntimeRefreshIsOldOrNew(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	const nodeID = "node-old-or-new"
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
	bus := edgeevents.NewBus()
	svc := New(reg, bus)

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "model-x",
			Providers: map[string]string{"prov-x": "served-x"},
		},
	}

	enabledFlag := true
	disabledFlag := false
	enabledStore := edgenode.NewNodeStore()
	enabledStore.Add(&edgenode.NodeRecord{
		ID: nodeID,
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID: "prov-x", Type: "vllm", Category: config.CategoryAPI,
				Adapter:  "vllm-gpu",
				Models:   []string{"served-x"},
				Health:   "available",
				Capacity: 1,
				Enabled:  &enabledFlag,
			},
		},
	})

	disabledStore := edgenode.NewNodeStore()
	disabledStore.Add(&edgenode.NodeRecord{
		ID: nodeID,
		Providers: []config.NodeProviderConf{
			{
				ID: "prov-x", Type: "vllm", Category: config.CategoryAPI,
				Health: "available", Capacity: 1,
				Enabled: &disabledFlag,
			},
		},
	})

	svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})

	// Submit a holder so the provider has inflight+queued pressure under the
	// enabled config. This makes the old tuple (available/capacity=1/inflight>0)
	// distinguishable from the new tuple (disabled/capacity=0/inflight=0).
	holder, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
		ModelGroupKey: "model-x", ProviderPool: true, Background: true,
	})
	if err != nil || holder == nil {
		t.Fatalf("holder submit: %v", err)
	}
	defer holder.Close()

	// Force the queue to hold the lease so inflight > 0 is observable.
	waitForInflight(t, svc.queue, nodeID, "prov-x", 1)

	// Submit a waiter so the model-x group has one queued item targeting
	// prov-x. This creates live candidate pressure (queued=1) that, combined
	// with the holder's in-flight lease (inFlight=1), makes the old tuple
	// (available/capacity=1/inFlight=1/queued=1) distinguishable from the new
	// disabled tuple (disabled/0/0/0/0/0). Without the waiter the reader would
	// always see queued=0 and the regression would be silent.
	type submitResult struct {
		run RunResult
		err error
	}
	waiterResult := make(chan submitResult, 1)
	go func() {
		run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
			ModelGroupKey: "model-x", ProviderPool: true, Background: true,
		})
		waiterResult <- submitResult{run: run, err: err}
	}()
	waitForQueueLen(t, svc.queue, "model-x", 1)

	// --- Public reader barrier: enforce queue→service lock ordering ---
	// Hold the service write lock so the reader goroutine (calling production
	// ListNodeSnapshots) blocks on the service read lock inside the queue-held
	// critical section. This proves the reader actually acquires and holds the
	// queue lock, and the writer must wait for it.
	svc.mu.Lock()
	serviceLocked := true
	defer func() {
		if serviceLocked {
			svc.mu.Unlock()
		}
	}()

	readerResult := make(chan []NodeSnapshot, 1)
	go func() {
		readerResult <- svc.ListNodeSnapshots()
	}()
	if err := waitForQueueLockHeld(svc.queue, time.Second); err != nil {
		svc.mu.Unlock()
		serviceLocked = false
		t.Fatal(err)
	}

	// Start writer while reader still holds queue lock and waits on service lock.
	writerDone := make(chan struct{})
	go func() {
		svc.SetRuntimeConfig(disabledStore, catalog, groupPolicy{})
		close(writerDone)
	}()
	svc.mu.Unlock()
	serviceLocked = false

	// Bounded waits for reader, writer completions.
	var oldSnaps []NodeSnapshot
	select {
	case oldSnaps = <-readerResult:
	case <-time.After(time.Second):
		t.Fatal("public snapshot reader did not complete")
	}
	select {
	case <-writerDone:
	case <-time.After(time.Second):
		t.Fatal("runtime config writer did not complete")
	}

	expectedOld := snapshotTuple{
		status: "available", capacity: 1, inFlight: 1, queued: 1,
		longInFlight: 0, longQueued: 0,
	}
	assertExactSnapshotTuple(t, oldSnaps, nodeID, "prov-x", expectedOld)

	// Public surface now returns the new tuple (disabled/capacity=0).
	newSnaps := svc.ListNodeSnapshots()
	expectedNew := snapshotTuple{
		status: "disabled", capacity: 0, inFlight: 0, queued: 0,
		longInFlight: 0, longQueued: 0,
	}
	assertExactSnapshotTuple(t, newSnaps, nodeID, "prov-x", expectedNew)

	svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
	// Drain holder
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holder.Dispatch().RunID, Type: "complete"})
	// Drain waiter (holder release unblocks it)
	var waiterRes submitResult
	select {
	case waiterRes = <-waiterResult:
	case <-time.After(2 * time.Second):
		t.Fatal("waiter did not dispatch after holder release")
	}
	if waiterRes.err != nil || waiterRes.run == nil {
		t.Fatalf("waiter result: run=%v err=%v", waiterRes.run, waiterRes.err)
	}
	// Dispatch and complete the waiter to fully drain the provider.
	wDispatch := waiterRes.run.Dispatch()
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})

	// Terminal drain verification: provider resource counts and lease must be zero.
	aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-x")
	if aInFlight != 0 || aLong != 0 {
		t.Fatalf("expected prov-x drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
	}
	if lc := leaseCount(svc.queue); lc != 0 {
		t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
	}
}

// assertExactSnapshotTuple asserts that the production ListNodeSnapshots
// surface returns an exact tuple for (nodeID, providerID). Used by the
// tuple regression test to confirm the reader captures a single consistent
// snapshot under the queue lock via the public surface.
func assertExactSnapshotTuple(t *testing.T, snaps []NodeSnapshot, nodeID, providerID string, want snapshotTuple) {
	t.Helper()
	ps := findProviderSnapshotOnNode(snaps, nodeID, providerID)
	if ps == nil {
		t.Fatalf("provider %s snapshot not found on %s in ListNodeSnapshots", providerID, nodeID)
	}
	got := snapshotTuple{
		status: ps.GetStatus(), capacity: ps.GetCapacity(),
		inFlight: ps.GetInFlight(), queued: ps.GetQueued(),
		longInFlight: ps.GetLongInFlight(), longQueued: ps.GetLongQueued(),
	}
	if err := validateSnapshotTuple(got, want); err != nil {
		t.Fatalf("tuple mismatch for %s on %s: %v", providerID, nodeID, err)
	}
}

// waitForInflight blocks until the queue for (nodeID, providerID) observes at
// least minInflight in-flight leases. Returns immediately if already satisfied.
func waitForInflight(t *testing.T, q *modelQueueManager, nodeID, providerID string, minInflight int) {
	t.Helper()
	for i := 0; i < 100; i++ {
		inFlight, _ := providerResourceCounts(q, nodeID, providerID)
		if inFlight >= minInflight {
			return
		}
		runtime.Gosched()
		time.Sleep(time.Millisecond)
	}
}

// waitForQueueLockHeld proves that another goroutine holds q.mu by attempting
// TryLock from the caller's goroutine. It returns nil if TryLock fails within
// the deadline (proving the lock is held exclusively) or a formatted error if
// TryLock succeeds before the deadline (lock was not held).
func waitForQueueLockHeld(q *modelQueueManager, deadline time.Duration) error {
	deadlineTimer := time.NewTimer(deadline)
	defer deadlineTimer.Stop()
	ticker := time.NewTicker(time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-deadlineTimer.C:
			return fmt.Errorf("queue lock not held within %v (TryLock succeeded: lock was not held exclusively)", deadline)
		case <-ticker.C:
			if !q.mu.TryLock() {
				return nil
			}
			q.mu.Unlock()
			runtime.Gosched()
		}
	}
}