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

package service

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"time"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

const (
	defaultNodeCapacity  = 1
	defaultGroupMaxQueue = 16
	defaultQueueTimeout  = 30 * time.Second
)

var (
	errQueueFull    = fmt.Errorf("queue is full")
	errQueueTimeout = fmt.Errorf("queue timeout")
)

// candidateNode pairs a registry entry with the per-request capacity derived
// from the node's adapter config (or Runtime.Concurrency as fallback).
// For provider-pool candidates, providerID, adapter, and servedTarget carry
// the globally-unique provider id, the dispatch adapter key, and the concrete
// model name to dispatch. Priority is the provider's dispatch priority (lower
// is preferred) and is used as a tie-break when in_flight counts are equal.
type candidateNode struct {
	entry        *edgenode.NodeEntry
	capacity     int
	priority     int
	providerID   string // non-empty for provider-pool candidates
	adapter      string // non-empty for provider-pool candidates; dispatch adapter key
	servedTarget string // concrete served model name; used for target rewrite
	// longContextCapacity is the provider's concurrent long-context slot limit.
	// Zero means the provider declares no dedicated long-slot limit and long
	// requests are not gated on it.
	longContextCapacity int
}

// slotKey returns a unique slot key for inflight accounting.
// For provider-pool candidates (providerID non-empty) it uses "nodeID:providerID"
// so that multiple providers on the same node are tracked independently.
// For legacy candidates it falls back to nodeID only.
func (c *candidateNode) slotKey() string {
	if c.providerID != "" {
		return c.entry.NodeID + ":" + c.providerID
	}
	return c.entry.NodeID
}

type groupPolicy struct {
	maxQueue     int
	queueTimeout time.Duration
}

type inflightRec struct {
	groupKey   string
	nodeID     string
	providerID string // non-empty for provider-pool dispatches
	long       bool   // true when a long-context slot was reserved for this run
}

type admitResult struct {
	candidate *candidateNode
	err       error
}

type queueItem struct {
	candidates []candidateNode
	waitCh     chan admitResult
	deadline   time.Time
	long       bool // true when the queued request is a long-context request
	reason     string
}

type modelQueueGroup struct {
	key              string
	policy           groupPolicy
	queue            []*queueItem
	inflight         map[string]int // slotKey → current in-flight count
	longInflight     map[string]int // slotKey → current long-context in-flight count
	lastSelectedSlot string
	adapter          string
	target           string
}

type modelQueueManager struct {
	mu            sync.Mutex
	groups        map[string]*modelQueueGroup
	inflightByRun map[string]inflightRec // runID → {groupKey, nodeID}
	store         *edgenode.NodeStore
}

func newModelQueueManager(store *edgenode.NodeStore) *modelQueueManager {
	return &modelQueueManager{
		groups:        make(map[string]*modelQueueGroup),
		inflightByRun: make(map[string]inflightRec),
		store:         store,
	}
}

func (m *modelQueueManager) setStore(store *edgenode.NodeStore) {
	m.mu.Lock()
	m.store = store
	m.mu.Unlock()
}

// startEventWatcher subscribes to all run and node events, releasing in-flight
// slots when runs terminate or nodes disconnect. Returns a stop function.
func (m *modelQueueManager) startEventWatcher(bus *edgeevents.Bus) func() {
	runCh, unsubRun := bus.SubscribeAllRuns(256)
	nodeCh, unsubNode := bus.SubscribeAllNodes(64)
	done := make(chan struct{})
	go func() {
		defer unsubRun()
		defer unsubNode()
		for {
			select {
			case e, ok := <-runCh:
				if !ok {
					return
				}
				if isTerminalRunEvent(e) {
					m.releaseRun(e.GetRunId(), e.GetType())
				}
			case e, ok := <-nodeCh:
				if !ok {
					return
				}
				if e.GetType() == eventpkg.TypeNodeDisconnected {
					m.releaseNode(e.GetNodeId(), "disconnected")
				}
			case <-done:
				return
			}
		}
	}()
	return func() { close(done) }
}

func isTerminalRunEvent(e *iop.RunEvent) bool {
	t := e.GetType()
	return t == "complete" || t == "error" || t == "cancelled"
}

// getOrCreateGroupLocked returns an existing group or creates one with the
// given policy (applying defaults for zero fields). Policy is only applied
// at creation; existing groups retain their current policy until an explicit
// updateGroupPolicyLocked call (used for config-refresh policy propagation).
func (m *modelQueueManager) getOrCreateGroupLocked(key string, policy groupPolicy) *modelQueueGroup {
	g, ok := m.groups[key]
	if !ok {
		if policy.maxQueue <= 0 {
			policy.maxQueue = defaultGroupMaxQueue
		}
		if policy.queueTimeout <= 0 {
			policy.queueTimeout = defaultQueueTimeout
		}
		g = &modelQueueGroup{
			key:          key,
			policy:       policy,
			inflight:     make(map[string]int),
			longInflight: make(map[string]int),
		}
		m.groups[key] = g
	}
	return g
}

// updateGroupPolicyLocked refreshes an existing group's queue policy so that
// provider max_queue / queue_timeout_ms changes applied via config refresh take
// effect on subsequent admissions. Only non-zero policy fields override the
// current values; in-flight counts and queued items are preserved untouched.
func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy groupPolicy) {
	if g == nil {
		return
	}
	if policy.maxQueue > 0 {
		g.policy.maxQueue = policy.maxQueue
	}
	if policy.queueTimeout > 0 {
		g.policy.queueTimeout = policy.queueTimeout
	}
}

// findAvailableNodeLocked returns the available candidate with the lowest
// absolute in_flight count. Full candidates are skipped, so a provider that is
// at capacity is deferred behind providers that still have room. When in_flight
// is equal, the candidate with the lower priority value wins. When both
// in_flight and priority are equal, all tied candidates remain in best[] so
// that the rotation logic in the caller picks the next slot after the last
// selected one.
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode, long bool) *candidateNode {
	var best []candidateNode
	bestInflight := 0
	bestPriority := 0
	for i := range candidates {
		c := candidates[i]
		if c.capacity <= 0 {
			continue
		}
		slot := c.slotKey()
		inflight := group.inflight[slot]
		if inflight >= c.capacity {
			continue
		}
		// Long-context gate: a long request is excluded from a provider whose
		// long slots are full. Providers with long_context_capacity == 0 declare
		// no dedicated long-slot limit and are not gated here. Normal requests
		// (long == false) never consult long slot state.
		if long && c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
			continue
		}
		switch {
		case len(best) == 0 || inflight < bestInflight:
			best = []candidateNode{c}
			bestInflight = inflight
			bestPriority = c.priority
		case inflight == bestInflight:
			if c.priority < bestPriority {
				best = []candidateNode{c}
				bestPriority = c.priority
			} else if c.priority == bestPriority {
				best = append(best, c)
			}
		}
	}
	if len(best) == 0 {
		return nil
	}
	sort.Slice(best, func(i, j int) bool {
		return candidateLess(&best[i], &best[j])
	})
	idx := 0
	if group.lastSelectedSlot != "" {
		for i := range best {
			if best[i].slotKey() == group.lastSelectedSlot {
				idx = (i + 1) % len(best)
				break
			}
		}
	}
	return &best[idx]
}

// candidateLess provides a deterministic ordering for equal-inflight/priority rotation:
// providerID first, then nodeID.
func candidateLess(a, b *candidateNode) bool {
	if b == nil {
		return true
	}
	if a.providerID != b.providerID {
		return a.providerID < b.providerID
	}
	return a.entry.NodeID < b.entry.NodeID
}

// admit selects an available candidate for the given model group, or queues
// the request until a slot opens. Blocks until a candidate is assigned, the
// queue timeout expires, or ctx is cancelled. Returns the selected candidateNode
// so callers can rewrite the dispatch target for provider-pool requests.
// The optional long flag marks the request as long-context so it is gated on
// provider long slots and reserves a long slot on dispatch.
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, error) {
	candidate, _, err := m.admitWithReason(ctx, groupKey, adapter, target, candidates, policy, longFlag...)
	return candidate, err
}

func (m *modelQueueManager) admitWithReason(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, string, error) {
	long := len(longFlag) > 0 && longFlag[0]
	m.mu.Lock()

	group := m.getOrCreateGroupLocked(groupKey, policy)
	// Refresh the policy on an already-existing group so capacity/queue changes
	// applied via config refresh (and recomputed per-request from the live store)
	// take effect on this and subsequent admissions.
	m.updateGroupPolicyLocked(group, policy)
	if group.adapter == "" {
		group.adapter = adapter
	}
	if group.target == "" {
		group.target = target
	}

	candidate := m.findAvailableNodeLocked(group, candidates, long)
	if candidate != nil {
		slot := candidate.slotKey()
		group.inflight[slot]++
		if long && candidate.longContextCapacity > 0 {
			group.longInflight[slot]++
		}
		group.lastSelectedSlot = slot
		m.mu.Unlock()
		return candidate, "dispatched", nil
	}

	if len(group.queue) >= group.policy.maxQueue {
		m.mu.Unlock()
		return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueFull)
	}

	// Determine queue reason.
	reason := "capacity_full"
	if long {
		longBlocked := false
		for _, c := range candidates {
			if c.capacity <= 0 {
				continue
			}
			slot := c.slotKey()
			inflight := group.inflight[slot]
			if inflight < c.capacity {
				if c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
					longBlocked = true
					break
				}
			}
		}
		if longBlocked {
			reason = "long_context_capacity_full"
		}
	}

	item := &queueItem{
		candidates: candidates,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(group.policy.queueTimeout),
		long:       long,
		reason:     reason,
	}
	group.queue = append(group.queue, item)
	timeout := group.policy.queueTimeout
	m.mu.Unlock()

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case res := <-item.waitCh:
		return res.candidate, item.reason, res.err
	case <-timer.C:
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before timeout fired.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Use public releaseSlot (which acquires m.mu) because we are
				// outside the lock here.  This avoids shared-state mutation
				// without mutex that the original releaseSlotByCandidate call
				// would cause.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				longSlot := item.long && res.candidate.longContextCapacity > 0
				m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
			}
		default:
		}
		return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueTimeout)
	case <-ctx.Done():
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before cancellation.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Same fix as timeout path: use public releaseSlot with
				// explicit mutex to avoid concurrent map/slice mutation.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				longSlot := item.long && res.candidate.longContextCapacity > 0
				m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
			}
		default:
		}
		return nil, "", ctx.Err()
	}
}

// trackInflight records a dispatched run so release events can find it.
// providerID is non-empty for provider-pool dispatches and is used to
// attribute in-flight counts to the correct provider snapshot. The optional
// longSlot flag records that a long-context slot was reserved so the terminal
// release also frees the long slot.
func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string, longSlot ...bool) {
	long := len(longSlot) > 0 && longSlot[0]
	m.mu.Lock()
	m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID, long: long}
	m.mu.Unlock()
}

// releaseRun releases the in-flight slot for a terminated run and dispatches
// the next queued item if one is waiting.
func (m *modelQueueManager) releaseRun(runID, reason string) {
	m.mu.Lock()
	rec, ok := m.inflightByRun[runID]
	if !ok {
		m.mu.Unlock()
		return
	}
	delete(m.inflightByRun, runID)
	m.releaseSlotLocked(rec.groupKey, rec.nodeID, rec.providerID, rec.long)
	m.mu.Unlock()
}

// releaseNode resets all in-flight slots on a disconnected node, removes it
// from all queued items' candidate lists, and tries to re-dispatch to any
// remaining live candidates.
func (m *modelQueueManager) releaseNode(nodeID, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	for runID, rec := range m.inflightByRun {
		if rec.nodeID == nodeID {
			delete(m.inflightByRun, runID)
		}
	}

	for _, group := range m.groups {
		// Remove the disconnected node from all queued items so a future
		// tryDispatch cannot pick it as a stale available candidate.
		for _, item := range group.queue {
			filtered := make([]candidateNode, 0, len(item.candidates))
			for _, c := range item.candidates {
				if c.entry.NodeID != nodeID {
					filtered = append(filtered, c)
				}
			}
			item.candidates = filtered
		}

		// Clear all slot keys that start with this nodeID.
		// For provider-pool: "nodeID:providerID". For legacy: "nodeID".
		// Long-context in-flight counters are keyed identically, so clear them
		// alongside the normal counters.
		for slot := range group.inflight {
			if colonIdx := findLastColon(slot); colonIdx > 0 {
				candidateNodeID := slot[:colonIdx]
				// Only clear if the node part matches.
				if candidateNodeID == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
					continue
				}
			} else {
				// Legacy slot: exact nodeID match.
				if slot == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
				}
			}
		}

		// tryDispatchLocked will handle remaining queued items, but we need
		// to call it if we cleared any inflight.
		// Since we already deleted entries, tryDispatchLocked will find available slots.
		// We need to re-check: if any slots were cleared and there are queued items.
		if len(group.queue) > 0 {
			m.tryDispatchLocked(group)
		}
	}
}

// releaseSlot decrements the in-flight count and tries to dispatch the next
// queued item. Used when admit-path I/O fails after the slot was reserved.
// For provider-pool dispatches (nodeID:providerID slot), pass providerID.
// For legacy dispatches, pass empty providerID so nodeID is used directly.
func (m *modelQueueManager) releaseSlot(groupKey, nodeID string, providerID ...string) {
	var pid string
	if len(providerID) > 0 {
		pid = providerID[0]
	}
	m.releaseSlotWithLong(groupKey, nodeID, pid, false)
}

// releaseSlotWithLong releases a reserved slot, additionally freeing a long-context
// slot when longSlot is true. Used by admit-path failure/race handling where the
// caller knows whether a long slot was reserved.
func (m *modelQueueManager) releaseSlotWithLong(groupKey, nodeID, providerID string, longSlot bool) {
	m.mu.Lock()
	m.releaseSlotLocked(groupKey, nodeID, providerID, longSlot)
	m.mu.Unlock()
}

func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string, longSlot bool) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	// Use provider-aware slot key for provider-pool dispatches.
	slot := nodeID
	if providerID != "" {
		slot = nodeID + ":" + providerID
	}
	if group.inflight[slot] > 0 {
		group.inflight[slot]--
	}
	if longSlot && group.longInflight[slot] > 0 {
		group.longInflight[slot]--
	}
	m.tryDispatchLocked(group)
}

// findLastColon returns the index of the last ':' in s, or -1 if not found.
func findLastColon(s string) int {
	for i := len(s) - 1; i >= 0; i-- {
		if s[i] == ':' {
			return i
		}
	}
	return -1
}

// tryDispatchLocked scans the queue from the front and dispatches every item
// that currently has an available candidate. Unlike a head-only pump, a blocked
// head does not stop dispatchable items behind it: when the head is a long
// request whose provider long slots are full, later normal requests that fit
// ordinary capacity are still dispatched, avoiding head-of-line blocking (SDD
// D05/D06). FIFO order is preserved among items with the same dispatchability
// because the scan always picks the earliest dispatchable item first. Expired
// items are timed out in place regardless of position. Must be called with
// m.mu held.
func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) {
	for {
		progressed := false
		for i := 0; i < len(group.queue); i++ {
			item := group.queue[i]

			if time.Now().After(item.deadline) {
				group.queue = append(group.queue[:i], group.queue[i+1:]...)
				select {
				case item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errQueueTimeout)}:
				default:
				}
				progressed = true
				break
			}

			candidate := m.findAvailableNodeLocked(group, item.candidates, item.long)
			if candidate == nil {
				// No slot for this item under current provider state; skip it so a
				// later dispatchable item is not blocked behind it.
				continue
			}

			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			slot := candidate.slotKey()
			reserveLong := item.long && candidate.longContextCapacity > 0
			group.inflight[slot]++
			if reserveLong {
				group.longInflight[slot]++
			}

			select {
			case item.waitCh <- admitResult{candidate: candidate}:
				group.lastSelectedSlot = slot
			default:
				// Caller already timed out or cancelled; release the reserved slot.
				group.inflight[slot]--
				if reserveLong {
					group.longInflight[slot]--
				}
			}
			progressed = true
			break
		}
		if !progressed {
			break
		}
	}
}

func (m *modelQueueManager) removeItemLocked(groupKey string, item *queueItem) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	for i, it := range group.queue {
		if it == item {
			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			return
		}
	}
}

func (m *modelQueueManager) getSnapshotForNode(nodeID string, rec *edgenode.NodeRecord) []*iop.ProviderSnapshot {
	m.mu.Lock()
	defer m.mu.Unlock()

	var snaps []*iop.ProviderSnapshot

	// Catalog-first: nodes with a providers[] catalog emit only catalog snapshots.
	// Adapter snapshots (sections 1-4) are omitted to prevent duplicate metrics
	// when a node has both adapter instances and provider-pool catalog entries.
	if len(rec.Providers) > 0 {
		for _, prov := range rec.Providers {
			if prov.ID == "" {
				continue
			}

			servedModels := make([]string, len(prov.Models))
			copy(servedModels, prov.Models)
			lifecycleCaps := make([]string, len(prov.LifecycleCapabilities))
			copy(lifecycleCaps, prov.LifecycleCapabilities)

			// Disabled providers appear in the snapshot with status=disabled and
			// effective capacity 0 so operators can see the switch state.
			if !config.ProviderEnabled(prov) {
				snaps = append(snaps, &iop.ProviderSnapshot{
					Adapter:               prov.Adapter,
					Status:                "disabled",
					Capacity:              0,
					InFlight:              0,
					Queued:                0,
					Id:                    prov.ID,
					Type:                  prov.Type,
					Category:              string(prov.Category),
					ServedModels:          servedModels,
					Health:                "disabled",
					LoadRatio:             0,
					LifecycleCapabilities: lifecycleCaps,
					LongContextCapacity:   0,
					LongInFlight:          0,
					LongQueued:            0,
				})
				continue
			}

			capVal := prov.Capacity
			inflight, queued := m.getStatsForProviderLocked(nodeID, prov.ID)
			longInflight, longQueued := m.getLongStatsForProviderLocked(nodeID, prov.ID)

			var loadRatio float32
			if capVal > 0 {
				loadRatio = float32(inflight) / float32(capVal)
			}

			snaps = append(snaps, &iop.ProviderSnapshot{
				Adapter:               prov.Adapter,
				Status:                "available",
				Capacity:              int32(capVal),
				InFlight:              int32(inflight),
				Queued:                int32(queued),
				Id:                    prov.ID,
				Type:                  prov.Type,
				Category:              string(prov.Category),
				ServedModels:          servedModels,
				Health:                prov.Health,
				LoadRatio:             loadRatio,
				LifecycleCapabilities: lifecycleCaps,
				LongContextCapacity:   int32(prov.LongContextCapacity),
				LongInFlight:          int32(longInflight),
				LongQueued:            int32(longQueued),
			})
		}
		return snaps
	}

	// Legacy adapter snapshots for nodes with no providers catalog.
	concurrencyFallback := 1
	if rec.Runtime.Concurrency > 0 {
		concurrencyFallback = rec.Runtime.Concurrency
	}

	// 1. CLI
	if rec.Adapters.CLI.Enabled {
		capVal := concurrencyFallback
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, "cli")
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  "cli",
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 2. Ollama
	for _, inst := range rec.Adapters.OllamaInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "ollama"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 3. vLLM
	for _, inst := range rec.Adapters.VllmInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "vllm"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 4. OpenAI Compat
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "openai_compat"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	return snaps
}

// getStatsForProviderLocked returns in-flight and queued counts for a
// provider-pool provider identified by (nodeID, providerID). Must be called
// with m.mu held.
func (m *modelQueueManager) getStatsForProviderLocked(nodeID, providerID string) (inFlight, queued int) {
	for _, rec := range m.inflightByRun {
		if rec.nodeID == nodeID && rec.providerID == providerID {
			inFlight++
		}
	}
	for _, group := range m.groups {
		for _, item := range group.queue {
			for _, c := range item.candidates {
				if c.entry.NodeID == nodeID && c.providerID == providerID {
					queued++
					break
				}
			}
		}
	}
	return
}

// getLongStatsForProviderLocked returns the long-context in-flight and queued count for a
// provider-pool provider identified by (nodeID, providerID), derived from tracked
// runs. Must be called with m.mu held.
func (m *modelQueueManager) getLongStatsForProviderLocked(nodeID, providerID string) (longInFlight, longQueued int) {
	for _, rec := range m.inflightByRun {
		if rec.nodeID == nodeID && rec.providerID == providerID && rec.long {
			longInFlight++
		}
	}
	for _, group := range m.groups {
		for _, item := range group.queue {
			if item.long {
				for _, c := range item.candidates {
					if c.entry.NodeID == nodeID && c.providerID == providerID {
						longQueued++
						break
					}
				}
			}
		}
	}
	return
}

// longInFlightForProvider returns the long-context in-flight count for a provider,
// acquiring the manager lock. Exposed for status/snapshot reporting and tests.
func (m *modelQueueManager) longInFlightForProvider(nodeID, providerID string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	longInFlight, _ := m.getLongStatsForProviderLocked(nodeID, providerID)
	return longInFlight
}

func (m *modelQueueManager) getStatsForAdapterLocked(nodeID string, rec *edgenode.NodeRecord, adapterName string) (inFlight, queued int) {
	for _, group := range m.groups {
		canonical, ok := resolveSnapshotAdapterName(rec, group.adapter, group.target)
		if ok && canonical == adapterName {
			if val, ok := group.inflight[nodeID]; ok {
				inFlight += val
			}
			for _, item := range group.queue {
				for _, c := range item.candidates {
					if c.entry.NodeID == nodeID {
						queued++
						break
					}
				}
			}
		}
	}
	return
}

// resolveSnapshotAdapterName returns the adapter/instance key used to match
// queue state for a provider snapshot.
// When adapterType is empty it falls back to the provider id (target) so that
// queue stats are still resolved instead of being dropped on an empty key.
func resolveSnapshotAdapterName(rec *edgenode.NodeRecord, adapterType, target string) (string, bool) {
	// Fallback to provider id when adapter is not set.
	if adapterType == "" {
		adapterType = target
	}
	if rec == nil {
		return adapterType, true
	}

	// 1. Exact instance Name match (highest priority).
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}

	// 2. Type-name route.
	switch adapterType {
	case "ollama":
		var enabled []string
		for _, inst := range rec.Adapters.OllamaInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "ollama"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "ollama", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "vllm":
		var enabled []string
		for _, inst := range rec.Adapters.VllmInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "vllm"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "vllm", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "openai_compat":
		var enabled []string
		for _, inst := range rec.Adapters.OpenAICompatInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "openai_compat"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "openai_compat", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "cli":
		return "cli", true

	default:
		return adapterType, true
	}
}