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

package openai

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"go.uber.org/zap"

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

type runService interface {
	SubmitRun(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error)
	SubmitProviderTunnel(context.Context, edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, 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,
		Adapter:   dispatch.Adapter,
		Target:    dispatch.Target,
		SessionID: dispatch.SessionID,
	}); 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
	modelCatalog                  []config.ModelCatalogEntry
	longContextThresholdTokens    int
	longContextThresholdTokensSet bool
	service                       runService
	logger                        *zap.Logger
	server                        *http.Server
}

func NewServer(cfg config.EdgeOpenAIConf, svc runService, logger *zap.Logger) *Server {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Server{cfg: cfg, service: svc, logger: logger}
}

// 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
			}
		}
	}
	return out
}

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

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