llmgw/internal/cli/server.go

package cli

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"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
			}
			controlMode := cfg.ControlPlaneMode()
			if !controlMode && len(cfg.Endpoints) == 0 {
				return errors.New("nothing to serve: declare models with profiles, " +
					"or give --endpoint for a router-only run")
			}

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

			// The control plane replaces the static endpoint list: agents
			// register, the reconciler places replicas, and Ready ones are
			// published to the router (plan §6.4).
			var agentHandler http.Handler
			var agentPath string
			var admin *control.Controller
			if controlMode {
				cat, err := cfg.Catalog()
				if err != nil {
					return err
				}

				ctl, svc, err := startControlPlane(ctx, cfg, cat, rt, cfg.AgentTLS.Enabled(), log)
				if err != nil {
					return err
				}
				defer ctl.Close()
				admin = ctl
				agentPath, agentHandler = svc.Handler()

				log.Info("control plane running",
					"models", len(cfg.Models), "agent_listen", cfg.AgentListen,
					"agent_mtls", cfg.AgentTLS.Enabled())
			}

			adminMux, err := rt.AdminHandler()
			if err != nil {
				return err
			}
			if admin != nil {
				adminapi.New(admin).Register(adminMux)
			}

			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:           adminMux,
				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,
			)

			var agentSrv *http.Server
			var agentLn net.Listener
			if agentHandler != nil {
				mux := http.NewServeMux()
				mux.Handle(agentPath, agentHandler)
				agentSrv = &http.Server{
					Handler:           mux,
					ReadHeaderTimeout: 10 * time.Second,
				}
				if !cfg.AgentTLS.Enabled() {
					// Without TLS, agents speak plaintext HTTP/2, which Go will
					// not do unless asked. Nothing then stops a node from
					// claiming to be another one, so this belongs on a trusted
					// network only.
					agentSrv.Handler = h2c.NewHandler(mux, &http2.Server{})
				}
				if agentLn, err = net.Listen("tcp", cfg.AgentListen); err != nil {
					_ = apiLn.Close()
					_ = adminLn.Close()
					return fmt.Errorf("listening on %s: %w", cfg.AgentListen, err)
				}
				if cfg.AgentTLS.Enabled() {
					tc, terr := cfg.AgentTLS.ServerConfig()
					if terr != nil {
						_ = apiLn.Close()
						_ = adminLn.Close()
						_ = agentLn.Close()
						return terr
					}
					// Serve on an already-wrapped listener rather than
					// ServeTLS: agentSrv.TLSConfig stays nil, which is exactly
					// the case where net/http sets up HTTP/2 for ALPN itself.
					agentLn = tls.NewListener(agentLn, tc)
				}
			}

			g, gctx := errgroup.WithContext(ctx)
			g.Go(func() error { return serve(apiSrv, apiLn) })
			g.Go(func() error { return serve(adminSrv, adminLn) })
			if agentSrv != nil {
				g.Go(func() error { return serve(agentSrv, agentLn) })
			}
			g.Go(func() error {
				<-gctx.Done()
				log.Info("shutting down")
				sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
				defer cancel()
				if agentSrv != nil {
					_ = agentSrv.Shutdown(sctx)
				}
				_ = 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
}

// startControlPlane builds the control plane and its agent service.
//
// The service is both the AgentService handler and the reconciler's Dispatcher:
// it holds the open streams, so it is the only thing that can deliver a command
// to a node.
func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Catalog,
	sink control.EndpointSink, requireIdentity bool, log *slog.Logger,
) (*control.Controller, *agentsvc.Service, error) {
	ccfg := control.DefaultConfig()
	ccfg.Reconciler = reconciler.Config{
		Interval:         cfg.Control.Interval,
		BaseBackoff:      cfg.Control.BaseBackoff,
		MaxBackoff:       cfg.Control.MaxBackoff,
		UnreachableAfter: cfg.Control.UnreachableAfter,
		LostAfter:        cfg.Control.LostAfter,
		MgmtInterface:    cfg.Control.MgmtInterface,
	}
	if len(cfg.Control.PortRange) == 2 {
		ccfg.Reconciler.PortRange = [2]int{cfg.Control.PortRange[0], cfg.Control.PortRange[1]}
	}

	// The service needs the inventory and the status sink, both of which the
	// controller owns; the controller needs the service as its dispatcher. The
	// cycle is broken by giving the controller a dispatcher that forwards to
	// the service once it exists.
	fwd := &dispatcherRef{}
	ctl := control.New(ccfg, cat, fwd, sink, log)
	svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), requireIdentity, log)
	fwd.set(svc)

	ctl.Start(ctx)
	return ctl, svc, nil
}

// dispatcherRef lets the controller be built before the service exists.
type dispatcherRef struct {
	mu sync.RWMutex
	d  reconciler.Dispatcher
}

func (r *dispatcherRef) set(d reconciler.Dispatcher) {
	r.mu.Lock()
	r.d = d
	r.mu.Unlock()
}

func (r *dispatcherRef) get() (reconciler.Dispatcher, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if r.d == nil {
		return nil, errors.New("the agent service is not ready yet")
	}
	return r.d, nil
}

func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
	d, err := r.get()
	if err != nil {
		return err
	}
	return d.StartShard(ctx, cmd)
}

func (r *dispatcherRef) StopShard(ctx context.Context, nodeID, shardID string) error {
	d, err := r.get()
	if err != nil {
		return err
	}
	return d.StopShard(ctx, nodeID, shardID)
}

func (r *dispatcherRef) DrainShard(ctx context.Context, nodeID, shardID string) error {
	d, err := r.get()
	if err != nil {
		return err
	}
	return d.DrainShard(ctx, nodeID, shardID)
}

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
}