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, true)
	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()
		}
	}
}

// TestListNodeSnapshotsIncludesConfiguredOfflineNodes verifies that configured
// nodes that are not currently connected still appear in the snapshot with
// Connected=false and preserve their configured identity and provider catalog.
// It also verifies that reconnecting with the same NodeID recovers Connected=true
// while keeping the same identity and stable position, and that the snapshot
// order is deterministic across no-queue and queue-backed paths.
func TestListNodeSnapshotsIncludesConfiguredOfflineNodes(t *testing.T) {
	t.Run("noQueue", func(t *testing.T) {
		t.Parallel()
		store := edgenode.NewNodeStore()
		// Add records in reverse index order to prove All() sorts by Index.
		store.Add(&edgenode.NodeRecord{
			ID:    "node-c",
			Alias: "gamma",
			Index: 2,
			Adapters: config.AdaptersConf{
				CLI: config.CLIConf{Enabled: true},
			},
		})
		store.Add(&edgenode.NodeRecord{
			ID:    "node-a",
			Alias: "alpha",
			Index: 0,
			Adapters: config.AdaptersConf{
				CLI: config.CLIConf{Enabled: true},
			},
		})
		store.Add(&edgenode.NodeRecord{
			ID:    "node-b",
			Alias: "beta",
			Index: 1,
			Adapters: config.AdaptersConf{
				CLI: config.CLIConf{Enabled: true},
			},
			Providers: []config.NodeProviderConf{
				{ID: "prov-b", Type: "cli", Category: config.CategoryCLI, Adapter: "cli", Models: []string{"m"}, Capacity: 2},
			},
		})

		reg := edgenode.NewRegistry()
		// Only node-a is connected.
		reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})

		svc := New(reg, nil)
		if svc.queue != nil {
			t.Fatal("expected no-queue service")
		}
		svc.SetNodeStore(store)

		snaps := svc.ListNodeSnapshots()
		if len(snaps) != 3 {
			t.Fatalf("expected 3 node snapshots, got %d", len(snaps))
		}

		// Stable order: index 0, 1, 2.
		expectedIDs := []string{"node-a", "node-b", "node-c"}
		expectedLabels := []string{"node0", "node1", "node2"}
		expectedConnected := []bool{true, false, false}
		for i, s := range snaps {
			if s.NodeID != expectedIDs[i] {
				t.Errorf("snap[%d].NodeID = %q, want %q", i, s.NodeID, expectedIDs[i])
			}
			if s.Label != expectedLabels[i] {
				t.Errorf("snap[%d].Label = %q, want %q", i, s.Label, expectedLabels[i])
			}
			if s.Connected != expectedConnected[i] {
				t.Errorf("snap[%d].Connected = %v, want %v", i, s.Connected, expectedConnected[i])
			}
		}

		// The no-queue path preserves the configured provider catalog. Legacy
		// adapter snapshot parity is outside this catalog regression.
		for _, s := range snaps {
			if s.NodeID == "node-b" {
				if len(s.ProviderSnapshots) != 1 {
					t.Errorf("node-b: expected 1 catalog snapshot (prov-b), got %d", len(s.ProviderSnapshots))
				}
				if len(s.ProviderSnapshots) > 0 {
					provider := s.ProviderSnapshots[0]
					// Offline node: catalog identity preserved (id/adapter/type/category/models),
					// but effective status/health/capacity drop to unavailable/offline/0.
					if provider.Id != "prov-b" || provider.Adapter != "cli" || provider.Status != "unavailable" || provider.Health != "offline" || provider.Capacity != 0 || provider.InFlight != 0 || provider.Queued != 0 || len(provider.ServedModels) != 1 || provider.ServedModels[0] != "m" {
						t.Errorf("node-b: unexpected configured provider snapshot (expected unavailable/offline/zero tuple with catalog identity): %+v", provider)
					}
				}
			}
		}
	})

	t.Run("withQueue", func(t *testing.T) {
		t.Parallel()
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{ID: "node-c", Alias: "gamma", Index: 2, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
		store.Add(&edgenode.NodeRecord{ID: "node-a", Alias: "alpha", Index: 0, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
		store.Add(&edgenode.NodeRecord{
			ID:       "node-b",
			Alias:    "beta",
			Index:    1,
			Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}},
			Providers: []config.NodeProviderConf{
				{ID: "prov-b", Type: "cli", Category: config.CategoryCLI, Adapter: "cli", Models: []string{"m"}, Capacity: 2},
			},
		})

		reg := edgenode.NewRegistry()
		reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})

		bus := edgeevents.NewBus()
		svc := New(reg, bus)
		if svc.queue == nil {
			t.Fatal("expected queue-backed service")
		}
		svc.SetNodeStore(store)

		snaps := svc.ListNodeSnapshots()
		if len(snaps) != 3 {
			t.Fatalf("expected 3 node snapshots, got %d", len(snaps))
		}

		expectedIDs := []string{"node-a", "node-b", "node-c"}
		for i, s := range snaps {
			if s.NodeID != expectedIDs[i] {
				t.Errorf("snap[%d].NodeID = %q, want %q", i, s.NodeID, expectedIDs[i])
			}
			if s.Connected != (s.NodeID == "node-a") {
				t.Errorf("snap[%d] (%s) Connected = %v, want %v", i, s.NodeID, s.Connected, s.NodeID == "node-a")
			}
			if s.NodeID == "node-b" {
				if len(s.ProviderSnapshots) != 1 {
					t.Errorf("node-b: expected 1 queue-backed catalog snapshot, got %d", len(s.ProviderSnapshots))
				} else if provider := s.ProviderSnapshots[0]; provider.Id != "prov-b" || provider.Adapter != "cli" || provider.Status != "unavailable" || provider.Health != "offline" || provider.Capacity != 0 || provider.InFlight != 0 || provider.Queued != 0 || provider.LongInFlight != 0 || provider.LongQueued != 0 {
					t.Errorf("node-b: unexpected queue-backed provider snapshot (expected unavailable/offline/zero tuple): %+v", provider)
				}
			}
		}
	})

	t.Run("reconnectPreservesIdentity", func(t *testing.T) {
		t.Parallel()
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{ID: "node-a", Alias: "alpha", Index: 0, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
		store.Add(&edgenode.NodeRecord{ID: "node-b", Alias: "beta", Index: 1, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})

		// Initial: both configured nodes are connected, with node-b on its first
		// accepted registry generation.
		reg := edgenode.NewRegistry()
		reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})
		firstNodeB := &edgenode.NodeEntry{NodeID: "node-b", Alias: "beta"}
		reg.Register(firstNodeB)
		if firstNodeB.ConnectionGeneration == 0 {
			t.Fatal("first node-b registration did not receive a generation")
		}

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

		assertSnapshot := func(stage string, wantNodeBConnected bool) {
			t.Helper()
			snaps := svc.ListNodeSnapshots()
			if len(snaps) != 2 {
				t.Fatalf("%s: expected 2 snapshots, got %d", stage, len(snaps))
			}
			if snaps[0].NodeID != "node-a" || snaps[1].NodeID != "node-b" {
				t.Fatalf("%s: stable identity/order changed: got [%s %s]", stage, snaps[0].NodeID, snaps[1].NodeID)
			}
			if snaps[1].Alias != "beta" || snaps[1].Label != "node1" {
				t.Errorf("%s: node-b identity changed: alias=%q label=%q", stage, snaps[1].Alias, snaps[1].Label)
			}
			if snaps[1].Connected != wantNodeBConnected {
				t.Errorf("%s: node-b Connected=%v, want %v", stage, snaps[1].Connected, wantNodeBConnected)
			}
		}
		assertSnapshot("connected", true)

		// Disconnect the current owner. The configured identity remains in the
		// same position but connectivity becomes false.
		reg.Unregister("node-b")
		if _, ok := reg.CurrentGeneration("node-b"); ok {
			t.Fatal("node-b generation remained live after unregister")
		}
		assertSnapshot("disconnected", false)

		// Reconnect the same Node ID through a new registry entry. The registry
		// must mint a later generation while the configured identity/order and
		// snapshot position stay stable.
		secondNodeB := &edgenode.NodeEntry{NodeID: "node-b", Alias: "beta"}
		reg.Register(secondNodeB)
		if secondNodeB.ConnectionGeneration <= firstNodeB.ConnectionGeneration {
			t.Fatalf("reconnect generation=%d, want > first generation=%d", secondNodeB.ConnectionGeneration, firstNodeB.ConnectionGeneration)
		}
		assertSnapshot("reconnected", true)
	})
}

// TestProviderSnapshotOfflineVsDisabled verifies that connectivity flipping
// only affects enabled providers: disabled providers keep status=disabled and
// health=disabled regardless of connectivity, while enabled disconnected
// providers show status=unavailable and health=offline. Catalog identity (id,
// adapter, type, category, models) is preserved in both cases.
func TestProviderSnapshotOfflineVsDisabled(t *testing.T) {
	disabled := false
	rec := &edgenode.NodeRecord{
		ID: "node-off-disabled",
		Providers: []config.NodeProviderConf{
			{ID: "prov-enabled", Type: "vllm", Category: config.CategoryAPI,
				Adapter: "vllm-gpu", Models: []string{"served-a"},
				Health: "available", Capacity: 3},
			{ID: "prov-disabled", Type: "vllm", Category: config.CategoryAPI,
				Adapter: "vllm-gpu", Models: []string{"served-b"},
				Health: "available", Capacity: 3, Enabled: &disabled},
		},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
		},
	}

	// wrapForFind converts []*iop.ProviderSnapshot into a single-element
	// []NodeSnapshot so we can reuse findProviderSnapshotOnNode.
	wrapForFind := func(ps []*iop.ProviderSnapshot) []NodeSnapshot {
		return []NodeSnapshot{{NodeID: "node-off-disabled", ProviderSnapshots: ps}}
	}

	assertDisabled := func(stage string, snaps []NodeSnapshot) {
		t.Helper()
		p := findProviderSnapshotOnNode(snaps, "node-off-disabled", "prov-disabled")
		if p == nil {
			t.Fatalf("%s: prov-disabled not found", stage)
		}
		if p.GetStatus() != "disabled" {
			t.Errorf("%s prov-disabled: status=%q want disabled", stage, p.GetStatus())
		}
		if p.GetHealth() != "disabled" {
			t.Errorf("%s prov-disabled: health=%q want disabled", stage, p.GetHealth())
		}
		if p.GetCapacity() != 0 || p.GetInFlight() != 0 || p.GetQueued() != 0 || p.GetLongInFlight() != 0 || p.GetLongQueued() != 0 {
			t.Errorf("%s prov-disabled: counters must be zero", stage)
		}
	}
	assertEnabled := func(stage string, snaps []NodeSnapshot, wantCap int32, wantStatus, wantHealth string) {
		t.Helper()
		p := findProviderSnapshotOnNode(snaps, "node-off-disabled", "prov-enabled")
		if p == nil {
			t.Fatalf("%s: prov-enabled not found", stage)
		}
		if p.GetStatus() != wantStatus {
			t.Errorf("%s prov-enabled: status=%q want %q", stage, p.GetStatus(), wantStatus)
		}
		if p.GetHealth() != wantHealth {
			t.Errorf("%s prov-enabled: health=%q want %q", stage, p.GetHealth(), wantHealth)
		}
		if p.GetCapacity() != wantCap {
			t.Errorf("%s prov-enabled: capacity=%d want %d", stage, p.GetCapacity(), wantCap)
		}
		if p.GetInFlight() != 0 || p.GetQueued() != 0 || p.GetLongInFlight() != 0 || p.GetLongQueued() != 0 {
			t.Errorf("%s prov-enabled: effective counters must be zero", stage)
		}
	}

	// --- Disabled: connected ---
	snaps := staticProviderCatalogSnapshots(rec, true)
	assertDisabled("connected", wrapForFind(snaps))
	assertEnabled("connected", wrapForFind(snaps), 3, "available", "available")

	// --- Disabled: disconnected (must still be disabled, not unavailable) ---
	snaps = staticProviderCatalogSnapshots(rec, false)
	assertDisabled("disconnected", wrapForFind(snaps))
	assertEnabled("disconnected", wrapForFind(snaps), 0, "unavailable", "offline")
}

// TestProviderSnapshotReconnectRestoresCapacityAndAdmission verifies the full
// generation-aware lease lifecycle across disconnect→offline→reconnect. The
// first connection acquires a real long-context lease, authoritative disconnect
// settlement releases it and orphans that generation, and a later connection
// restores configured capacity while exposing only its own newly admitted lease.
func TestProviderSnapshotReconnectRestoresCapacityAndAdmission(t *testing.T) {
	const (
		nodeID      = "node-reconnect-cycle"
		providerID  = "prov-rc"
		modelGroup  = "model-rc"
		servedModel = "served-x"
	)

	rec := &edgenode.NodeRecord{
		ID: nodeID,
		Providers: []config.NodeProviderConf{
			{
				ID:                  providerID,
				Type:                "vllm",
				Category:            config.CategoryAPI,
				Adapter:             "vllm-gpu",
				Models:              []string{servedModel},
				Health:              "available",
				Capacity:            5,
				LongContextCapacity: 2,
			},
		},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
		},
	}

	reg := edgenode.NewRegistry()
	firstEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "rc"}
	reg.Register(firstEntry)
	if firstEntry.ConnectionGeneration == 0 {
		t.Fatal("first connection generation must be non-zero")
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	policy := NewGroupPolicy(16, time.Second)
	svc.SetRuntimeConfig(store, []config.ModelCatalogEntry{
		{ID: modelGroup, Providers: map[string]string{providerID: servedModel}},
	}, policy)

	req := SubmitRunRequest{
		ModelGroupKey: modelGroup,
		ContextClass:  contextClassLong,
		ProviderPool:  true,
	}
	admitForGeneration := func(stage string, generation uint64) *candidateNode {
		t.Helper()
		candidates, _, err := svc.resolveQueueCandidates(req)
		if err != nil {
			t.Fatalf("%s: resolve provider-pool candidates: %v", stage, err)
		}
		selected, err := svc.queue.admit(
			t.Context(), modelGroup, "", "", candidates, policy,
			svc.resolveQueueCandidatesClosure(req), true, true,
		)
		if err != nil {
			t.Fatalf("%s: admit long-context lease: %v", stage, err)
		}
		if selected.providerID != providerID {
			t.Fatalf("%s: admitted provider=%q, want %q", stage, selected.providerID, providerID)
		}
		if selected.generation != generation {
			t.Fatalf("%s: admitted generation=%d, want %d", stage, selected.generation, generation)
		}
		if selected.leaseID == 0 {
			t.Fatalf("%s: admitted lease id must be non-zero", stage)
		}
		svc.queue.trackLease(selected.leaseID, "run-"+stage)
		return selected
	}

	assertSnapshot := func(stage, wantStatus, wantHealth string, wantCapacity, wantInFlight, wantLongCapacity, wantLongInFlight int32) {
		t.Helper()
		ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID)
		if ps == nil {
			t.Fatalf("%s: %s snapshot not found", stage, providerID)
		}
		if ps.GetId() != providerID || ps.GetAdapter() != "vllm-gpu" {
			t.Errorf("%s: catalog identity=(%q,%q), want (%q,%q)", stage, ps.GetId(), ps.GetAdapter(), providerID, "vllm-gpu")
		}
		if len(ps.GetServedModels()) != 1 || ps.GetServedModels()[0] != servedModel {
			t.Errorf("%s: served models=%v, want [%s]", stage, ps.GetServedModels(), servedModel)
		}
		if ps.GetStatus() != wantStatus || ps.GetHealth() != wantHealth {
			t.Errorf("%s: status/health=(%q,%q), want (%q,%q)", stage, ps.GetStatus(), ps.GetHealth(), wantStatus, wantHealth)
		}
		if ps.GetCapacity() != wantCapacity || ps.GetInFlight() != wantInFlight {
			t.Errorf("%s: capacity/in-flight=(%d,%d), want (%d,%d)", stage, ps.GetCapacity(), ps.GetInFlight(), wantCapacity, wantInFlight)
		}
		if ps.GetLongContextCapacity() != wantLongCapacity || ps.GetLongInFlight() != wantLongInFlight {
			t.Errorf("%s: long capacity/in-flight=(%d,%d), want (%d,%d)", stage, ps.GetLongContextCapacity(), ps.GetLongInFlight(), wantLongCapacity, wantLongInFlight)
		}
		if ps.GetQueued() != 0 || ps.GetLongQueued() != 0 {
			t.Errorf("%s: queued/long-queued=(%d,%d), want (0,0)", stage, ps.GetQueued(), ps.GetLongQueued())
		}
		t.Logf("%s: generation status=%s health=%s capacity=%d in_flight=%d long_capacity=%d long_in_flight=%d",
			stage, ps.GetStatus(), ps.GetHealth(), ps.GetCapacity(), ps.GetInFlight(), ps.GetLongContextCapacity(), ps.GetLongInFlight())
	}

	// Stage 1: the first accepted generation owns one real tracked long lease.
	firstLease := admitForGeneration("first-generation", firstEntry.ConnectionGeneration)
	assertSnapshot("first-generation-admitted", "available", "available", 5, 1, 2, 1)

	// Stage 2: remove the current registry owner, then settle that exact owner
	// generation through the same authoritative service hook used by transport.
	firstGeneration := firstEntry.ConnectionGeneration
	reg.Unregister(nodeID)
	if _, ok := reg.CurrentGeneration(nodeID); ok {
		t.Fatal("disconnected node retained a current registry generation")
	}
	svc.HandleNodeDisconnect(nodeID, firstGeneration, "test authoritative disconnect")

	svc.queue.mu.Lock()
	_, firstLeaseStillLive := svc.queue.leases[firstLease.leaseID]
	resource := svc.queue.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
	resourceOrphan := resource != nil && resource.orphan
	resourceGeneration := uint64(0)
	resourceInFlight, resourceLongInFlight := 0, 0
	if resource != nil {
		resourceGeneration = resource.generation
		resourceInFlight, resourceLongInFlight = resource.snapshotCounts()
	}
	svc.queue.mu.Unlock()
	if firstLeaseStillLive {
		t.Fatal("authoritative disconnect did not settle the first-generation lease")
	}
	if resource == nil || !resourceOrphan || resourceGeneration != firstGeneration {
		t.Fatalf("settled resource=(present=%t orphan=%t generation=%d), want (true,true,%d)", resource != nil, resourceOrphan, resourceGeneration, firstGeneration)
	}
	if resourceInFlight != 0 || resourceLongInFlight != 0 {
		t.Fatalf("settled resource counters=(%d,%d), want (0,0)", resourceInFlight, resourceLongInFlight)
	}
	assertSnapshot("offline-after-settlement", "unavailable", "offline", 0, 0, 0, 0)

	// Stage 3: reconnect the same identity. Configured capacity is visible, but
	// the old generation contributes no counter before a new admission occurs.
	secondEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "rc"}
	reg.Register(secondEntry)
	if secondEntry.ConnectionGeneration <= firstGeneration {
		t.Fatalf("reconnect generation=%d, want > disconnected generation=%d", secondEntry.ConnectionGeneration, firstGeneration)
	}
	assertSnapshot("reconnected-before-admission", "available", "available", 5, 0, 2, 0)

	// Stage 4: a real lease on the higher generation is admitted and is the only
	// counter exposed. Releasing that lease returns both counters to zero.
	secondLease := admitForGeneration("second-generation", secondEntry.ConnectionGeneration)
	assertSnapshot("reconnected-after-admission", "available", "available", 5, 1, 2, 1)
	svc.queue.releaseLease(secondLease.leaseID, "test release")
	assertSnapshot("reconnected-after-release", "available", "available", 5, 0, 2, 0)
}

// TestReconnectConnectHookActivatesResourceGenerationInSnapshot verifies the
// auxiliary S15 convergence: the accepted-connect hook (HandleNodeConnect) alone,
// with no new admission, clears the disconnect orphan marker, advances the
// provider resource generation to the reconnect's, and returns the status
// snapshot to available with the configured capacity and zero in-flight counters.
func TestReconnectConnectHookActivatesResourceGenerationInSnapshot(t *testing.T) {
	const (
		nodeID      = "node-hook-activate"
		providerID  = "prov-hook"
		modelGroup  = "model-hook"
		servedModel = "served-hook"
	)

	rec := &edgenode.NodeRecord{
		ID: nodeID,
		Providers: []config.NodeProviderConf{{
			ID:                  providerID,
			Type:                "vllm",
			Category:            config.CategoryAPI,
			Adapter:             "vllm-gpu",
			Models:              []string{servedModel},
			Health:              "available",
			Capacity:            4,
			LongContextCapacity: 2,
		}},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
		},
	}

	reg := edgenode.NewRegistry()
	firstEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "hook"}
	reg.Register(firstEntry)
	genOne := firstEntry.ConnectionGeneration
	if genOne == 0 {
		t.Fatal("first connection generation must be non-zero")
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)

	svc := New(reg, edgeevents.NewBus())
	policy := NewGroupPolicy(16, time.Second)
	svc.SetRuntimeConfig(store, []config.ModelCatalogEntry{
		{ID: modelGroup, Providers: map[string]string{providerID: servedModel}},
	}, policy)

	// Admit a real lease on the first generation so the resource generation is
	// genOne (not zero) before the disconnect.
	req := SubmitRunRequest{ModelGroupKey: modelGroup, ProviderPool: true}
	cands, _, err := svc.resolveQueueCandidates(req)
	if err != nil {
		t.Fatalf("resolve provider-pool candidates: %v", err)
	}
	lease, err := svc.queue.admit(
		t.Context(), modelGroup, "", "", cands, policy,
		svc.resolveQueueCandidatesClosure(req), false, true,
	)
	if err != nil || lease == nil {
		t.Fatalf("admit generation-one lease: %v", err)
	}
	if lease.generation != genOne {
		t.Fatalf("admitted lease generation=%d, want %d", lease.generation, genOne)
	}
	svc.queue.trackLease(lease.leaseID, "run-hook")

	if present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID); !present || orphan || gen != genOne || inFlight != 1 {
		t.Fatalf("generation-one resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,false,%d,1)", present, orphan, gen, inFlight, genOne)
	}

	// Disconnect the current owner: the lease settles and the resource is orphaned
	// at generation one.
	reg.Unregister(nodeID)
	svc.HandleNodeDisconnect(nodeID, genOne, "disconnected")
	if present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID); !present || !orphan || gen != genOne || inFlight != 0 {
		t.Fatalf("offline resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,true,%d,0)", present, orphan, gen, inFlight, genOne)
	}
	if ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID); ps == nil || ps.GetStatus() != "unavailable" || ps.GetHealth() != "offline" {
		t.Fatalf("offline snapshot = %+v, want status=unavailable health=offline", ps)
	}

	// Reconnect and drive the accepted-connect hook alone — no new admission,
	// refresh, or lease release.
	secondEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "hook"}
	reg.Register(secondEntry)
	genTwo := secondEntry.ConnectionGeneration
	if genTwo <= genOne {
		t.Fatalf("reconnect generation=%d, want > disconnected generation=%d", genTwo, genOne)
	}
	svc.HandleNodeConnect(nodeID, genTwo)

	present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID)
	if !present || orphan || gen != genTwo || inFlight != 0 {
		t.Fatalf("activated resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,false,%d,0)", present, orphan, gen, inFlight, genTwo)
	}
	ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID)
	if ps == nil {
		t.Fatal("provider snapshot missing after activation")
	}
	if ps.GetStatus() != "available" || ps.GetHealth() != "available" {
		t.Errorf("activated snapshot status/health=(%q,%q), want (available,available)", ps.GetStatus(), ps.GetHealth())
	}
	if ps.GetCapacity() != 4 || ps.GetInFlight() != 0 {
		t.Errorf("activated snapshot capacity/in-flight=(%d,%d), want (4,0)", ps.GetCapacity(), ps.GetInFlight())
	}
	if ps.GetLongContextCapacity() != 2 || ps.GetLongInFlight() != 0 {
		t.Errorf("activated snapshot long capacity/in-flight=(%d,%d), want (2,0)", ps.GetLongContextCapacity(), ps.GetLongInFlight())
	}
}