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

package service

import (
	"fmt"
	"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 fences the leases held by the disconnecting connection
// identified by (nodeID, generation). The transport calls it only after the
// registry confirms the disconnecting client still owned the entry, and passes
// that owner's connection generation so a stale callback — or one racing a
// reconnect that has already taken the id — fences only its own generation and
// never the live connection's leases.
func (s *Service) HandleNodeDisconnect(nodeID string, generation uint64, reason string) {
	if nodeID == "" || s.queue == nil {
		return
	}
	s.queue.releaseNode(nodeID, generation, reason)
}

// HandleNodeConnect activates the provider resources owned by an accepted node
// connection identified by (nodeID, generation) and immediately re-resolves and
// pumps every pending item. The transport calls it only for an accepted
// (non-duplicate) registration, ahead of the observability fanout, so a waiter
// stranded when the node went offline is re-dispatched by the reconnect alone —
// with no new request, config refresh, or lease release to trigger the pump.
//
// The generation is re-checked against the registry's current owner, but the
// check runs inside activateNode under the queue lock — the same critical section
// that mutates the resources and pumps, and the one releaseNode holds when it
// fences a disconnect. Checking here and then activating would reopen a TOCTOU: a
// disconnect settling in that gap could unregister the owner and release its
// leases, after which a stale activation would resurrect the orphaned resources.
// A callback that a newer reconnect has already superseded, or one racing a
// disconnect that already removed the entry, is therefore a no-op. This is the
// connect counterpart to HandleNodeDisconnect.
func (s *Service) HandleNodeConnect(nodeID string, generation uint64) {
	if nodeID == "" || s.queue == nil {
		return
	}
	// A generation of 0 is untracked/legacy and never activates; only a
	// registry-confirmed current owner does. The currency check is deferred into
	// activateNode so it is atomic with the activation under the queue lock.
	if generation == 0 || s.registry == nil {
		return
	}
	s.queue.activateNode(nodeID, generation, func() bool {
		return s.registry.IsCurrentOwnerGeneration(nodeID, generation)
	})
}

// candidateIsCurrentOwner reports whether the selected candidate's connection
// generation still matches the registry's current owner for its node. It is the
// pre-send/handoff fence: a lease minted for a connection that has since
// disconnected or been superseded by a reconnect is caught here, deterministically,
// instead of relying on a send to a dead client to fail. An untracked candidate
// (generation 0 — legacy/direct dispatch) is always treated as current.
func (s *Service) candidateIsCurrentOwner(selected *candidateNode) bool {
	if selected == nil || selected.entry == nil || selected.generation == 0 || s.registry == nil {
		return true
	}
	return s.registry.IsCurrentOwnerGeneration(selected.entry.NodeID, selected.generation)
}

// staleGenerationError describes a dispatch fenced because the selected
// candidate's connection was superseded between admission and send.
func staleGenerationError(selected *candidateNode) error {
	return fmt.Errorf("provider node %q connection changed before dispatch (fenced generation %d)", selected.entry.NodeID, selected.generation)
}

// 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()
}

// NodeStore returns the current node store snapshot held by the service.
// Exported for test observability only.
func (s *Service) NodeStore() *edgenode.NodeStore {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.nodeStore
}

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