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

package service

import (
	"sync"

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

const (
	DefaultSessionID  = "default"
	DefaultTimeoutSec = 30
)

// Service is the surface-neutral application service shared by the console,
// HTTP/A2A input surfaces, and the Control Plane connector. Its responsibilities
// are split across this package by concern: status_provider.go (node snapshots,
// capabilities, health status), run_dispatch.go (run lifecycle), node_command.go
// (node command transport), and control_command.go (Control Plane command
// execution). The provider-pool policy is stored here alongside the runtime
// catalog so that snapshot readers observe it under the same service lock that
// the runtime writer uses, eliminating the race where status readers read the
// queue's policy field concurrently with a runtime apply.
type Service struct {
	mu                 sync.RWMutex
	registry           *edgenode.Registry
	events             *edgeevents.Bus
	nodeStore          *edgenode.NodeStore
	queue              *modelQueueManager
	modelCatalog       []config.ModelCatalogEntry
	providerPoolPolicy groupPolicy
	tunnels            *providerTunnelRouter
}

func New(registry *edgenode.Registry, events *edgeevents.Bus) *Service {
	s := &Service{registry: registry, events: events, tunnels: newProviderTunnelRouter()}
	if events != nil {
		s.queue = newModelQueueManager(nil)
	}
	return s
}

// HandleRunLifecycleEvent releases the lease owning a terminated run. The
// transport calls it directly, ahead of the observability fanout, because the
// event bus drops into full subscriber channels: lease accounting must not
// depend on a delivery that is allowed to fail. Non-terminal events are ignored.
func (s *Service) HandleRunLifecycleEvent(event *iop.RunEvent) {
	if event == nil || s.queue == nil || !isTerminalRunEvent(event) {
		return
	}
	s.queue.releaseRun(event.GetRunId(), event.GetType())
}

// HandleNodeDisconnect releases every lease held on a node that just lost its
// connection. The transport calls it only after the registry confirms the
// disconnecting client still owned the entry, so a stale duplicate connection
// closing cannot free the live connection's leases.
func (s *Service) HandleNodeDisconnect(nodeID, reason string) {
	if nodeID == "" || s.queue == nil {
		return
	}
	s.queue.releaseNode(nodeID, reason)
}

// SetNodeStore swaps the node catalog. Handing the store to the queue manager is
// not bookkeeping: the manager reconciles provider resource state against the new
// catalog and re-evaluates every queued request, so capacity raised or a provider
// re-enabled by this call wakes its waiters immediately instead of at their next
// group-local event.
func (s *Service) SetNodeStore(store *edgenode.NodeStore) {
	s.mu.Lock()
	s.nodeStore = store
	s.mu.Unlock()
	if s.queue != nil {
		s.queue.setStore(store)
	}
}

// SetRuntimeConfig atomically replaces the runtime config snapshot used by
// dispatch and status readers. It first acquires the queue lock (which the
// live resolver holds while recomputing candidates) and then the service lock
// (which the status snapshot and this writer both use), matching the lock
// order every concurrent path already follows so the resolver cannot deadlock
// on s.mu while this method waits for m.mu.
//
// The policy field lives on the service alongside the runtime snapshot so the
// status reader reads it under s.mu instead of reaching into the queue's
// internal policy field concurrently with a runtime apply.
func (s *Service) SetRuntimeConfig(store *edgenode.NodeStore, catalog []config.ModelCatalogEntry, policy groupPolicy) {
	catalog = cloneModelCatalog(catalog)
	if s.queue == nil {
		s.mu.Lock()
		s.nodeStore = store
		s.modelCatalog = catalog
		s.providerPoolPolicy = policy
		s.mu.Unlock()
		return
	}
	s.queue.mu.Lock()
	defer s.queue.mu.Unlock()
	s.mu.Lock()
	s.nodeStore = store
	s.modelCatalog = catalog
	s.providerPoolPolicy = policy
	s.mu.Unlock()

	// Existing queue-locked reconcile helper; already holds m.mu.
	s.queue.setProviderPoolPolicyLocked(store, policy)
}

// SetModelCatalog provides the top-level provider-pool model catalog to the
// service. Must be called before the first provider-pool SubmitRun.
func (s *Service) SetModelCatalog(catalog []config.ModelCatalogEntry) {
	s.mu.Lock()
	s.modelCatalog = cloneModelCatalog(catalog)
	s.mu.Unlock()
}

func (s *Service) runtimeConfigSnapshot() (*edgenode.NodeStore, []config.ModelCatalogEntry, groupPolicy) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	// The policy is now a service-owned field; reading it under s.mu is
	// race-free against SetRuntimeConfig which writes it under s.mu.Lock.
	return s.nodeStore, cloneModelCatalog(s.modelCatalog), s.providerPoolPolicy
}

func cloneModelCatalog(catalog []config.ModelCatalogEntry) []config.ModelCatalogEntry {
	if len(catalog) == 0 {
		return nil
	}
	out := make([]config.ModelCatalogEntry, len(catalog))
	for i, entry := range catalog {
		out[i] = entry
		if entry.Providers != nil {
			out[i].Providers = make(map[string]string, len(entry.Providers))
			for k, v := range entry.Providers {
				out[i].Providers[k] = v
			}
		}
	}
	return out
}

// ProviderPoolPolicy returns an exported snapshot of the current root
// provider-pool queue policy as observed by the service. It is race-free
// against SetRuntimeConfig because the policy is stored on the service and
// written under s.mu. Exported for test observability only.
func (s *Service) ProviderPoolPolicy() GroupPolicySnapshot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.providerPoolPolicy.Snapshot()
}

func (s *Service) ListNodes() []*edgenode.NodeEntry {
	return s.registry.All()
}

// NormalizeSessionID maps an empty session id to the default session.
func NormalizeSessionID(id string) string {
	if id == "" {
		return DefaultSessionID
	}
	return id
}

func normalizeTimeoutSec(timeoutSec int) int {
	if timeoutSec <= 0 {
		return DefaultTimeoutSec
	}
	return timeoutSec
}

func nodeLabel(entry *edgenode.NodeEntry) string {
	return entry.DisplayLabel()
}

// resolveQueueCandidatesClosure returns a closure that recomputes the candidate
// universe for the given request from the current store/catalog/registry at
// pump time. The global pump calls this under the manager lock when dispatching
// a queued item, so a provider that was disabled or capacity-zero at enqueue
// time can still wake its waiters once it becomes eligible — the pump no
// longer trusts the enqueue-time snapshot alone. Each call re-snapshots the
// runtime config, so concurrent SetRuntimeConfig is observed without stale
// reads.
func (s *Service) resolveQueueCandidatesClosure(req SubmitRunRequest) func() ([]candidateNode, error) {
	return func() ([]candidateNode, error) {
		candidates, _, err := s.resolveQueueCandidates(req)
		return candidates, err
	}
}