iop/apps/edge/internal/bootstrap/runtime.go

package bootstrap

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

	"go.uber.org/zap"

	"iop/apps/edge/internal/configrefresh"
	edgecontrolplane "iop/apps/edge/internal/controlplane"
	edgeevents "iop/apps/edge/internal/events"
	edgeinput "iop/apps/edge/internal/input"
	edgenode "iop/apps/edge/internal/node"
	edgeservice "iop/apps/edge/internal/service"
	"iop/apps/edge/internal/transport"
	"iop/packages/go/config"
	"iop/packages/go/observability"
	"iop/packages/go/version"
	iop "iop/proto/gen/iop"
)

// Runtime is the shared edge runtime assembly used by both `serve` and
// `console` entrypoints. NewRuntime wires logger, registry, node store,
// event bus, service, and transport server; Start performs handler wiring
// and binds the server and metrics endpoint.
type Runtime struct {
	Cfg          *config.EdgeConfig
	Logger       *zap.Logger
	Registry     *edgenode.Registry
	NodeStore    *edgenode.NodeStore
	EventBus     *edgeevents.Bus
	Service      *edgeservice.Service
	Server       *transport.Server
	Input        *edgeinput.Manager
	Artifact     *ArtifactServer
	ControlPlane *edgecontrolplane.Connector
	RefreshAdmin *RefreshAdminServer

	refreshMu      sync.Mutex
	cfgMu          sync.RWMutex
	lifetimeMu     sync.Mutex
	lifetimeCancel context.CancelFunc
}

// convertProviderPoolConf converts the config-level EdgeProviderPoolConf into
// the service-layer groupPolicy that the queue manager uses for provider-pool
// admission. Zero max_queue normalizes to the canonical default; explicit
// queue_timeout_ms of 0 preserves the legacy no-timeout semantics.
func convertProviderPoolConf(conf config.EdgeProviderPoolConf) edgeservice.GroupPolicy {
	maxQueue := conf.MaxQueue
	if maxQueue <= 0 {
		maxQueue = config.DefaultProviderPoolMaxQueue
	}
	timeout := time.Duration(conf.QueueTimeoutMS) * time.Millisecond
	return edgeservice.NewGroupPolicy(maxQueue, timeout)
}

func NewRuntime(cfg *config.EdgeConfig) (*Runtime, error) {
	logger, err := observability.NewLoggerWithFile(cfg.Logging.Level, cfg.Logging.Pretty, cfg.Logging.Path)
	if err != nil {
		return nil, err
	}

	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig(cfg.Nodes)
	if err != nil {
		return nil, fmt.Errorf("edge: seed node store: %w", err)
	}

	bus := edgeevents.NewBus()
	svc := edgeservice.New(registry, bus)
	svc.SetRuntimeConfig(nodeStore, cfg.Models, convertProviderPoolConf(cfg.ProviderPool))
	inputManager := edgeinput.NewManager(*cfg, svc, logger.Named("input"))
	artifactServer := NewArtifactServer(cfg.Bootstrap.Listen, cfg.Bootstrap.ArtifactDir, logger.Named("bootstrap"))

	server, err := transport.NewServer(cfg.Server.Listen, registry, nodeStore, logger)
	if err != nil {
		return nil, err
	}

	connector := edgecontrolplane.NewConnector(
		cfg.Edge,
		cfg.ControlPlane,
		version.Version,
		logger.Named("controlplane"),
		edgecontrolplane.WithStatusProvider(svc),
		edgecontrolplane.WithNodeEventBus(bus),
	)

	rt := &Runtime{
		Cfg:          cfg,
		Logger:       logger,
		Registry:     registry,
		NodeStore:    nodeStore,
		EventBus:     bus,
		Service:      svc,
		Server:       server,
		Input:        inputManager,
		Artifact:     artifactServer,
		ControlPlane: connector,
	}
	rt.RefreshAdmin = newRefreshAdminServer(cfg.Refresh.Listen, rt, logger.Named("refresh-admin"))
	return rt, nil
}

func (r *Runtime) wireHandlers() {
	// Authoritative lifecycle first: the service settles run/node accounting
	// synchronously from the transport, and the event bus stays a pure
	// observability fanout that is free to drop into full subscribers.
	r.Server.SetRunLifecycleHandler(r.Service.HandleRunLifecycleEvent)
	r.Server.SetNodeConnectHandler(r.Service.HandleNodeConnect)
	r.Server.SetNodeDisconnectHandler(r.Service.HandleNodeDisconnect)
	r.Server.SetRunEventHandler(r.EventBus.PublishRun)
	r.Server.SetNodeEventHandler(r.EventBus.PublishNode)
	// Tunnel frames bypass the event bus: raw provider bytes go to the
	// request-bound tunnel stream owned by the service.
	r.Server.SetTunnelFrameHandler(r.Service.RouteProviderTunnelFrame)
}

func (r *Runtime) Start(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	r.wireHandlers()
	lifetimeCtx := r.newLifetimeContext()
	if err := r.Server.Start(lifetimeCtx); err != nil {
		r.cancelLifetime()
		return err
	}
	if err := ctx.Err(); err != nil {
		r.cancelLifetime()
		_ = r.Server.Stop()
		return err
	}
	if err := r.Input.Start(lifetimeCtx); err != nil {
		r.cancelLifetime()
		_ = r.Server.Stop()
		return err
	}
	if err := r.Artifact.Start(lifetimeCtx); err != nil {
		r.cancelLifetime()
		_ = r.Input.Stop(context.Background())
		_ = r.Server.Stop()
		return err
	}
	if err := r.ControlPlane.Start(lifetimeCtx); err != nil {
		r.cancelLifetime()
		_ = r.Artifact.Stop(context.Background())
		_ = r.Input.Stop(context.Background())
		_ = r.Server.Stop()
		return err
	}
	if r.Cfg.Refresh.Enabled {
		if err := r.RefreshAdmin.Start(lifetimeCtx); err != nil {
			r.cancelLifetime()
			r.ControlPlane.Stop()
			_ = r.Artifact.Stop(context.Background())
			_ = r.Input.Stop(context.Background())
			_ = r.Server.Stop()
			return err
		}
	}
	r.startMetrics()
	return nil
}

func (r *Runtime) Stop() error {
	r.cancelLifetime()
	r.ControlPlane.Stop()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	refreshErr := r.RefreshAdmin.Stop(ctx)
	inputErr := r.Input.Stop(ctx)
	artifactErr := r.Artifact.Stop(ctx)
	err := r.Server.Stop()
	_ = r.Logger.Sync()
	if err == nil {
		err = refreshErr
	}
	if err == nil {
		err = inputErr
	}
	if err == nil {
		err = artifactErr
	}
	return err
}

// RefreshConfig evaluates a config refresh request against the current runtime
// config. In apply mode, mutable-only changes are committed to the runtime.
func (r *Runtime) RefreshConfig(ctx context.Context, req configrefresh.Request) (configrefresh.Result, error) {
	r.refreshMu.Lock()
	defer r.refreshMu.Unlock()

	r.cfgMu.RLock()
	current := r.Cfg
	r.cfgMu.RUnlock()

	result, candidate, err := configrefresh.Evaluate(ctx, current, req)
	if err != nil {
		return result, err
	}
	if req.Mode == configrefresh.ModeApply && result.Status == configrefresh.StatusApplied && candidate != nil {
		nodeResults, err := r.applyMutableConfig(ctx, candidate, result.Changes, req.RequestID)
		if err != nil {
			// Apply failed after a clean classification: preserve the previous
			// runtime snapshot (applyMutableConfig commits nothing on error) and
			// surface a rejected result with stable ops-report slices.
			return configrefresh.RejectedResult(req, "apply candidate config: "+err.Error()), nil
		}
		result.NodeResults = nodeResults
	}
	return result, nil
}

// applyMutableConfig commits a validated mutable-only candidate via copy-on-write
// replacement. Readers observe either the previous runtime snapshot or the new
// one, and validation/build failures leave the previous snapshot untouched.
// After the local commit, it pushes the new config to all connected nodes and
// returns per-node results. Node push failures do not roll back the local commit.
// When len(changes)==0 (no-change apply), the local snapshot is still committed
// but connected nodes are NOT pushed, preventing a spurious restart_required
// response from identical configs.
func (r *Runtime) applyMutableConfig(ctx context.Context, candidate *config.EdgeConfig, changes []configrefresh.Change, requestID string) ([]configrefresh.NodeResult, error) {
	nextStore, err := edgenode.LoadFromConfig(candidate.Nodes)
	if err != nil {
		return nil, err
	}

	r.cfgMu.Lock()
	r.Cfg = candidate
	r.NodeStore = nextStore
	r.cfgMu.Unlock()

	r.Server.SetNodeStore(nextStore)
	poolPolicy := convertProviderPoolConf(candidate.ProviderPool)
	r.Service.SetRuntimeConfig(nextStore, candidate.Models, poolPolicy)
	r.Input.SetModelCatalog(candidate.Models)
	r.Input.OpenAI.SetLongContextThreshold(candidate.LongContextThresholdTokens)

	// No-change apply: commit the snapshot but skip node push to prevent
	// spurious restart_required from identical configs.
	if len(changes) == 0 {
		r.Logger.Debug("no-change apply refresh: skipping node push", zap.String("request_id", requestID))
		return nil, nil
	}

	changedPaths := make([]string, 0, len(changes))
	for _, ch := range changes {
		changedPaths = append(changedPaths, ch.Path)
	}

	nodeResults := r.Server.PushConfigRefresh(ctx, buildNodeConfigRefreshRequest(requestID, changedPaths))
	return nodeResults, nil
}

func buildNodeConfigRefreshRequest(requestID string, changedPaths []string) *iop.NodeConfigRefreshRequest {
	return &iop.NodeConfigRefreshRequest{
		RequestId:    requestID,
		ChangedPaths: changedPaths,
	}
}

func (r *Runtime) newLifetimeContext() context.Context {
	r.lifetimeMu.Lock()
	defer r.lifetimeMu.Unlock()
	if r.lifetimeCancel != nil {
		r.lifetimeCancel()
	}
	lifetimeCtx, cancel := context.WithCancel(context.Background())
	r.lifetimeCancel = cancel
	return lifetimeCtx
}

func (r *Runtime) cancelLifetime() {
	r.lifetimeMu.Lock()
	cancel := r.lifetimeCancel
	r.lifetimeCancel = nil
	r.lifetimeMu.Unlock()
	if cancel != nil {
		cancel()
	}
}

func (r *Runtime) startMetrics() {
	if r.Cfg.Metrics.Port <= 0 {
		return
	}
	go func() {
		if err := observability.ServeMetrics(r.Cfg.Metrics.Port); err != nil {
			r.Logger.Warn("metrics server exited", zap.Error(err))
		}
	}()
}