llmgw/internal/cli/server.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/sync/errgroup"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router"
)

func newServerCmd() *cobra.Command {
	var (
		configPath  string
		listen      string
		adminListen string
		endpoints   []string
	)

	cmd := &cobra.Command{
		Use:   "server",
		Short: "Run the router (phase 1; the control plane arrives in phase 3)",
		Long: "Serves the OpenAI API in front of a set of engine replicas.\n\n" +
			"In this phase the replica list is static: give it in the config file, or\n" +
			"pass --endpoint id=URL=model[,model] once per replica for a quick test.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			log := logger()

			cfg, err := config.Load(configPath)
			if err != nil {
				return err
			}
			if listen != "" {
				cfg.Listen = listen
			}
			if adminListen != "" {
				cfg.AdminListen = adminListen
			}
			for _, spec := range endpoints {
				ep, err := parseEndpointFlag(spec)
				if err != nil {
					return err
				}
				cfg.Endpoints = append(cfg.Endpoints, ep)
			}
			if err := cfg.Validate(); err != nil {
				return err
			}
			if len(cfg.Endpoints) == 0 {
				return errors.New("no endpoints configured; use --endpoint or a config file")
			}

			rt, err := router.New(cfg, log)
			if err != nil {
				return err
			}

			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			rt.Start(ctx)
			defer rt.Stop()

			adminHandler, err := rt.AdminHandler()
			if err != nil {
				return err
			}

			apiSrv := &http.Server{
				Handler:           rt.Handler(),
				ReadHeaderTimeout: 10 * time.Second,
				// No write timeout: responses stream for minutes and the proxy
				// enforces its own per-phase deadlines.
			}
			adminSrv := &http.Server{
				Handler:           adminHandler,
				ReadHeaderTimeout: 10 * time.Second,
			}

			apiLn, err := net.Listen("tcp", cfg.Listen)
			if err != nil {
				return fmt.Errorf("listening on %s: %w", cfg.Listen, err)
			}
			adminLn, err := net.Listen("tcp", cfg.AdminListen)
			if err != nil {
				_ = apiLn.Close()
				return fmt.Errorf("listening on %s: %w", cfg.AdminListen, err)
			}

			log.Info("router listening",
				"api", apiLn.Addr().String(),
				"admin", adminLn.Addr().String(),
				"endpoints", len(cfg.Endpoints),
				"default_profile", cfg.Router.DefaultProfile,
			)

			g, gctx := errgroup.WithContext(ctx)
			g.Go(func() error { return serve(apiSrv, apiLn) })
			g.Go(func() error { return serve(adminSrv, adminLn) })
			g.Go(func() error {
				<-gctx.Done()
				log.Info("shutting down")
				sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
				defer cancel()
				_ = adminSrv.Shutdown(sctx)
				return apiSrv.Shutdown(sctx)
			})
			return g.Wait()
		},
	}

	f := cmd.Flags()
	f.StringVar(&configPath, "config", "", "configuration file (optional)")
	f.StringVar(&listen, "listen", "", "override the client-facing address")
	f.StringVar(&adminListen, "admin-listen", "", "override the metrics/pprof address")
	f.StringArrayVar(&endpoints, "endpoint", nil, "static endpoint: id=URL=model[,model] (repeatable)")

	return cmd
}

func serve(s *http.Server, ln net.Listener) error {
	if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
		return err
	}
	return nil
}

// parseEndpointFlag decodes id=URL=model[,model].
//
// The URL contains '=' only in pathological cases, so splitting into exactly
// three parts keeps the flag readable while staying unambiguous.
func parseEndpointFlag(spec string) (config.Endpoint, error) {
	parts := strings.SplitN(spec, "=", 3)
	if len(parts) != 3 {
		return config.Endpoint{}, fmt.Errorf("--endpoint %q: expected id=URL=model[,model]", spec)
	}
	models := strings.Split(parts[2], ",")
	for i := range models {
		models[i] = strings.TrimSpace(models[i])
	}
	return config.Endpoint{
		ID:     strings.TrimSpace(parts[0]),
		URL:    strings.TrimRight(strings.TrimSpace(parts[1]), "/"),
		Models: models,
		Pool:   "static",
		Weight: 1,
	}, nil
}