iop/apps/edge/internal/openai/server.go

package openai

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"go.uber.org/zap"

	"iop/apps/edge/internal/authprojection"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/auth"
	"iop/packages/go/config"
	"iop/packages/go/streamgate"
)

type runService interface {
	SubmitRun(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error)
	SubmitProviderTunnel(context.Context, edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, error)
	SubmitProviderPool(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error)
	OllamaAPI(context.Context, edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error)
	CancelRun(context.Context, edgeservice.CancelRunRequest) (edgeservice.CommandResult, error)
}

// cancelRunOnHTTPGiveUp sends CancelRun to Node when the HTTP caller gave up
// (request cancellation/timeout) before the run reached a terminal state.
// Terminal run outcomes are not cancel-worthy; see isCancelWorthyRunError.
func (s *Server) cancelRunOnHTTPGiveUp(dispatch edgeservice.RunDispatch, err error) {
	if !isCancelWorthyRunError(err) || dispatch.RunID == "" {
		return
	}
	s.sendCancelRun(dispatch)
}

// sendCancelRun propagates cancellation for a dispatched run/tunnel to the
// Node cancel path.
func (s *Server) sendCancelRun(dispatch edgeservice.RunDispatch) {
	if dispatch.RunID == "" {
		return
	}
	if _, cancelErr := s.service.CancelRun(context.Background(), edgeservice.CancelRunRequest{
		NodeRef: dispatch.NodeID,
		RunID:   dispatch.RunID,
	}); cancelErr != nil {
		s.logger.Warn("openai cancel run failed",
			zap.String("run_id", dispatch.RunID),
			zap.Error(cancelErr),
		)
	}
}

type Server struct {
	mu                            sync.RWMutex
	cfg                           config.EdgeOpenAIConf
	edgeID                        string
	modelCatalog                  []config.ModelCatalogEntry
	longContextThresholdTokens    int
	longContextThresholdTokensSet bool
	service                       runService
	logger                        *zap.Logger
	server                        *http.Server
	obsSink                       streamgate.ObservationSink
	obsSinkIsDefault              bool
	livenessCollectors            *livenessRecoveryCollectors
	principalProjection           authprojection.Reader
	credentialMode                credentialMode
	executionPresets              []config.ExecutionPreset
	requestCoordinator            *logicalRequestCoordinator
	artifactFrontiers             *artifactFrontierStore
	lightFlows                    *hotPathLightStore
	hotPathObserver               hotPathObserver
	hotPathObserverHook           hotPathObserverFailureHook
}

// SetCredentialPlaneManaged selects the request authentication and provider
// credential source from the single Edge credential_plane.enabled switch.
// Runtime assembly calls this before the server starts; it is not a live mode
// switch.
func (s *Server) SetCredentialPlaneManaged(enabled bool) {
	s.mu.Lock()
	if enabled {
		s.credentialMode = credentialModeManaged
	} else {
		s.credentialMode = credentialModeLegacy
	}
	s.mu.Unlock()
}

func (s *Server) managedCredentialPlane() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.credentialMode == credentialModeManaged
}

// CredentialPlaneManaged reports the immutable startup credential mode for
// runtime assembly tests and diagnostics.
func (s *Server) CredentialPlaneManaged() bool {
	return s.managedCredentialPlane()
}

func NewServer(cfg config.EdgeOpenAIConf, svc runService, logger *zap.Logger) *Server {
	if logger == nil {
		logger = zap.NewNop()
	}
	s := &Server{
		cfg:                cfg,
		service:            svc,
		logger:             logger,
		obsSink:            newZapFilterObservationSink(logger),
		obsSinkIsDefault:   true,
		livenessCollectors: defaultLivenessRecoveryCollectors,
		requestCoordinator: newLogicalRequestCoordinator(logicalRequestCoordinatorOptions{}),
		artifactFrontiers:  newArtifactFrontierStore(defaultArtifactFrontierCapacity),
		lightFlows:         newHotPathLightStore(defaultHotPathLightCapacity),
		hotPathObserver:    newZapHotPathObserver(logger),
	}
	if s.hotPathObserver == nil {
		s.hotPathObserver = hotPathNoopObserver{}
	}
	return s
}

// logicalRequests returns the Edge-local coordinator installed for this server.
// Preset-backed Chat and Messages ingress join this coordinator before dispatch.
func (s *Server) logicalRequests() *logicalRequestCoordinator {
	return s.requestCoordinator
}

// SetPrincipalProjection installs the shared, transport-neutral projection
// reader. A nil or unmanaged reader preserves the legacy static auth mode.
func (s *Server) SetPrincipalProjection(projection authprojection.Reader) {
	s.mu.Lock()
	s.principalProjection = projection
	s.mu.Unlock()
}

// PrincipalProjection returns the exact reader installed on this server.
// Manager uses this narrow accessor to verify that all ingress auth reads the
// shared cache instance.
func (s *Server) PrincipalProjection() authprojection.Reader {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.principalProjection
}

// SetModelCatalog provides the provider-pool model catalog to the OpenAI server.
// When set, /v1/models lists catalog IDs and requests matching catalog entries
// are dispatched via the provider pool instead of the legacy model_routes path.
func (s *Server) SetModelCatalog(catalog []config.ModelCatalogEntry) {
	s.mu.Lock()
	s.modelCatalog = cloneModelCatalog(catalog)
	s.mu.Unlock()
}

func (s *Server) modelCatalogSnapshot() []config.ModelCatalogEntry {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return cloneModelCatalog(s.modelCatalog)
}

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
			}
		}
		if entry.TokenCounter != nil {
			counter := *entry.TokenCounter
			out[i].TokenCounter = &counter
		}
	}
	return out
}

// SetExecutionPresets provides the execution preset catalog to the OpenAI server using a deep clone snapshot.
func (s *Server) SetExecutionPresets(presets []config.ExecutionPreset) {
	s.mu.Lock()
	s.executionPresets = config.CloneExecutionPresetCatalog(presets)
	s.mu.Unlock()
}

// ExecutionPresetsSnapshot returns a deep cloned snapshot of the current execution preset catalog.
func (s *Server) ExecutionPresetsSnapshot() []config.ExecutionPreset {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return config.CloneExecutionPresetCatalog(s.executionPresets)
}

// ExecutionPreset returns a deep copy of the execution preset matching id.
func (s *Server) ExecutionPreset(id string) (config.ExecutionPreset, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, p := range s.executionPresets {
		if p.ID == id {
			return p.Clone(), true
		}
	}
	return config.ExecutionPreset{}, false
}

func (s *Server) Enabled() bool {
	return s != nil && s.cfg.Enabled
}

// SetEdgeID provides this Edge instance's stable identity, used as the low-
// cardinality edge_id label on OpenAI-compatible usage metrics.
func (s *Server) SetEdgeID(id string) {
	s.mu.Lock()
	s.edgeID = id
	s.mu.Unlock()
}

func (s *Server) edgeIDValue() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.edgeID != "" {
		return s.edgeID
	}
	return "edge-local"
}

// SetHotPathObserver installs a distinct observer for Hot Path lifecycle
// events. It is separate from Server.obsSink (Stream Gate) so the two
// observability contracts never share ownership. A nil observer installs a
// noop observer so observer failures can never alter response behavior.
func (s *Server) SetHotPathObserver(observer hotPathObserver) {
	s.mu.Lock()
	if observer == nil {
		s.hotPathObserver = hotPathNoopObserver{}
	} else {
		s.hotPathObserver = observer
	}
	s.mu.Unlock()
}

// HotPathObserver returns the exact observer installed on this server. It is
// exported for tests and diagnostics only; production code routes through
// emitHotPathObservation, which wraps the observer with failure isolation.
func (s *Server) HotPathObserver() hotPathObserver {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.hotPathObserver == nil {
		return hotPathNoopObserver{}
	}
	return s.hotPathObserver
}

// SetHotPathObserverHook installs the optional failure hook for the Hot Path
// observer. It is called whenever the observer returns an error or panics. The
// hook is isolated from request results.
func (s *Server) SetHotPathObserverHook(hook hotPathObserverFailureHook) {
	s.mu.Lock()
	s.hotPathObserverHook = hook
	s.mu.Unlock()
}

func (s *Server) hotPathObservationSnapshot() (hotPathObserver, hotPathObserverFailureHook) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	observer := s.hotPathObserver
	if observer == nil {
		observer = hotPathNoopObserver{}
	}
	return observer, s.hotPathObserverHook
}

// emitHotPathObservation is the single production emission seam for Hot Path
// observations. It snapshots observer state under the server lock, then emits
// through bounded, failure-isolated wrappers.
func (s *Server) emitHotPathObservation(ctx context.Context, projection hotPathLogProjection) {
	if s == nil {
		return
	}
	observer, hook := s.hotPathObservationSnapshot()
	failureHook := func(failed hotPathLogProjection, err error) {
		initHotPathMetrics().recordObserverFailure(s.edgeIDValue())
		invokeHotPathObserverFailureHookSafely(hook, failed, err)
	}
	safe := hotPathSafeObserver{
		inner:     &hotPathBoundedObserver{inner: observer},
		onFailure: failureHook,
	}
	_ = safe.Emit(ctx, projection)
}

// SetObservationSink replaces the default observation sink used to emit
// streamgate_filter_observation entries for this server's request runtimes.
// A nil sink installs a NoopObservationSink so observation failures can never
// alter response behavior. Every call transfers ownership to the application:
// the constructor-owned-default flag is cleared so the request-local liveness
// projection never suppresses forwarding to an explicitly installed sink, even
// when that sink is another *zapFilterObservationSink of the built-in type.
func (s *Server) SetObservationSink(sink streamgate.ObservationSink) {
	s.mu.Lock()
	if sink == nil {
		s.obsSink = streamgate.NoopObservationSink{}
	} else {
		s.obsSink = sink
	}
	s.obsSinkIsDefault = false
	s.mu.Unlock()
}

// observationSink returns a fresh request-local liveness observation projection
// wrapping the configured downstream sink. The wrapper only suppresses the
// private-liveness/ExactReplay rows from the generic writer when the downstream
// is this server's constructor-owned default sink; every explicitly installed
// sink receives the original immutable observations.
func (s *Server) observationSink() streamgate.ObservationSink {
	s.mu.RLock()
	downstream := s.obsSink
	logger := s.logger
	suppressDefault := s.obsSinkIsDefault
	collectors := s.livenessCollectors
	s.mu.RUnlock()
	if downstream == nil {
		downstream = streamgate.NoopObservationSink{}
	}
	return newOpenAILivenessObservationSink(downstream, logger, suppressDefault, collectors)
}

// SetLongContextThreshold sets the input-token threshold at or above which a
// request is classified as long-context for admission policy.
func (s *Server) SetLongContextThreshold(val int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.longContextThresholdTokens = val
	s.longContextThresholdTokensSet = true
}

func (s *Server) longContextThreshold() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.longContextThresholdTokensSet {
		return s.longContextThresholdTokens
	}
	return defaultLongContextThreshold
}

const defaultLongContextThreshold = 100000

func (s *Server) Start(ctx context.Context) error {
	if !s.Enabled() {
		return nil
	}
	if s.cfg.Listen == "" {
		s.cfg.Listen = "0.0.0.0:18081"
	}

	mux := s.routes()

	s.server = &http.Server{
		Addr:              s.cfg.Listen,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}
	ln, err := net.Listen("tcp", s.cfg.Listen)
	if err != nil {
		return fmt.Errorf("openai server listen %s: %w", s.cfg.Listen, err)
	}
	if s.cfg.TLS.Enabled {
		tlsConfig, loadErr := auth.LoadHTTPServerTLS(s.cfg.TLS.Cert, s.cfg.TLS.Key)
		if loadErr != nil {
			_ = ln.Close()
			return fmt.Errorf("openai server TLS: %w", loadErr)
		}
		ln = tls.NewListener(ln, tlsConfig)
	}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = s.server.Shutdown(shutdownCtx)
	}()
	go func() {
		s.logger.Info("openai-compatible server listening", zap.String("addr", s.cfg.Listen))
		if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			s.logger.Warn("openai-compatible server exited", zap.Error(err))
		}
	}()
	return nil
}

func (s *Server) Stop(ctx context.Context) error {
	if s == nil || s.server == nil {
		return nil
	}
	return s.server.Shutdown(ctx)
}

func responseModel(requestModel, target string) string {
	if requestModel != "" {
		return requestModel
	}
	return target
}

func httpStatusForRunError(err error) int {
	if errors.Is(err, context.Canceled) {
		return http.StatusRequestTimeout
	}
	return http.StatusBadGateway
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, errorResponse{Error: errorBody{Type: code, Message: message}})
}