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

package service

import (
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

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