MOPAC/internal/serve/server.go

// Package serve is the `harness serve` OpenAI-compatible front door for
// OpenWebUI (DESIGN "OWUI front door"; any OpenAI client works):
//
//	GET  /v1/models           the servable model catalog
//	POST /v1/chat/completions ONE bounded stateless conductor turn
//	GET  /healthz             liveness
//
// The class -> tier map in harness.toml is the catalog: every [models.classes]
// class is exposed as a model named mopac-<class> and routes through the same
// tier table `once` uses. v0 is STATELESS (the client sends the full
// conversation history each call; no session storage), NON-streaming, and
// tools are OFF — a pure chat path over the shared conductor turn machinery.
package serve

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sort"
	"strings"
	"time"

	"ukrrs.com/mopac/harness/internal/config"
	"ukrrs.com/mopac/harness/internal/llm"
	"ukrrs.com/mopac/harness/internal/loop"
	"ukrrs.com/mopac/harness/internal/models"
)

// maxBodyBytes bounds request payloads; OWUI sends the whole conversation
// history per call, so this is roomier than the webhook receiver's cap.
const maxBodyBytes = 4 << 20

// ModelPrefix is the servable-model namespace: class "code" is served as
// ModelPrefix + "code" = "mopac-code".
const ModelPrefix = "mopac-"

// Turner is the conductor-facing turn sink; *loop.Conductor implements it
// (same machinery `once` uses, tools off).
type Turner interface {
	ServeTurn(ctx context.Context, msgs []llm.Message, model string, opts loop.ServeTurnOpts) (*loop.TurnResult, error)
}

// Server is the `harness serve` front door. Every route requires a valid
// Bearer vkey (constant-time compare); anything else gets 401 with a single
// generic body. The vkey is resolved at startup and never logged.
type Server struct {
	cfg     config.ServeConfig
	catalog map[string]models.Decision // model name -> routing decision
	names   []string                   // sorted servable model names
	vkey    string
	turner  Turner
	logger  *log.Logger
	started int64
}

// NewServer resolves the vkey (fail-fast) and builds the servable catalog
// from the router's class map — each class becomes ModelPrefix+class and
// carries the tier/model it routes to — optionally narrowed to
// [serve] enabled_models.
func NewServer(cfg config.ServeConfig, router *models.Router, turner Turner, out io.Writer) (*Server, error) {
	if cfg.VKeyRef == "" {
		return nil, fmt.Errorf("[serve]: vkey_ref is required (the bearer key OWUI connections present)")
	}
	vkey, err := config.ResolveKeyRef(cfg.VKeyRef)
	if err != nil {
		return nil, fmt.Errorf("[serve]: resolve vkey: %w", err)
	}
	catalog := make(map[string]models.Decision)
	for _, class := range router.Classes() {
		d, err := router.Resolve(class)
		if err != nil {
			return nil, fmt.Errorf("serve catalog: %w", err)
		}
		catalog[ModelPrefix+class] = d
	}
	names := make([]string, 0, len(catalog))
	for name := range catalog {
		names = append(names, name)
	}
	sort.Strings(names)
	if len(cfg.EnabledModels) > 0 {
		keep := make(map[string]bool, len(cfg.EnabledModels))
		for _, m := range cfg.EnabledModels {
			if _, ok := catalog[m]; !ok {
				return nil, fmt.Errorf("[serve]: enabled_models entry %q is not servable (valid: %s)", m, strings.Join(names, ", "))
			}
			keep[m] = true
		}
		filtered := make([]string, 0, len(keep))
		for _, n := range names {
			if keep[n] {
				filtered = append(filtered, n)
			} else {
				delete(catalog, n)
			}
		}
		names = filtered
	}
	if len(names) == 0 {
		return nil, fmt.Errorf("[serve]: model catalog is empty (define [models.classes] in harness.toml)")
	}
	return &Server{
		cfg:     cfg,
		catalog: catalog,
		names:   names,
		vkey:    vkey,
		turner:  turner,
		logger:  log.New(out, "serve: ", log.LstdFlags|log.Lmsgprefix),
		started: time.Now().Unix(),
	}, nil
}

// Handler builds the front door's HTTP routes.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"status":"ok"}`)
	})
	mux.HandleFunc("/v1/models", s.handleModels)
	mux.HandleFunc("/v1/chat/completions", s.handleChat)
	return mux
}

// ModelNames returns the sorted servable model names (startup banner, docs).
func (s *Server) ModelNames() []string {
	return append([]string(nil), s.names...)
}

func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) {
	if !s.authorize(r) {
		s.logger.Printf("rejected models remote=%s reason=auth", r.RemoteAddr)
		s.rejectAuth(w)
		return
	}
	if r.Method != http.MethodGet {
		w.Header().Set("Allow", http.MethodGet)
		s.writeError(w, http.StatusMethodNotAllowed, "GET only")
		return
	}
	data := make([]modelEntry, 0, len(s.names))
	for _, name := range s.names {
		data = append(data, modelEntry{ID: name, Object: "model", Created: s.started, OwnedBy: "mopac"})
	}
	s.writeJSON(w, http.StatusOK, modelList{Object: "list", Data: data})
}

func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	if !s.authorize(r) {
		s.logger.Printf("rejected chat remote=%s reason=auth", r.RemoteAddr)
		s.rejectAuth(w)
		return
	}
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", http.MethodPost)
		s.writeError(w, http.StatusMethodNotAllowed, "POST only")
		return
	}
	body, ok := s.readBody(w, r)
	if !ok {
		return
	}
	var req inboundRequest
	if err := json.Unmarshal(body, &req); err != nil {
		s.writeError(w, http.StatusBadRequest, "malformed request body")
		return
	}
	if req.Stream {
		s.writeError(w, http.StatusBadRequest, "streaming is not supported; set stream to false")
		return
	}
	decision, ok := s.catalog[req.Model]
	if !ok {
		s.writeError(w, http.StatusBadRequest,
			fmt.Sprintf("unknown model %q; valid models: %s", req.Model, strings.Join(s.names, ", ")))
		return
	}
	if len(req.Messages) == 0 {
		s.writeError(w, http.StatusBadRequest, "messages must not be empty")
		return
	}
	if req.MaxTokens < 0 {
		s.writeError(w, http.StatusBadRequest, "max_tokens must not be negative")
		return
	}
	msgs := make([]llm.Message, 0, len(req.Messages))
	for i, m := range req.Messages {
		switch m.Role {
		case "system", "user", "assistant":
		default:
			s.writeError(w, http.StatusBadRequest, fmt.Sprintf("messages[%d]: unsupported role %q", i, m.Role))
			return
		}
		msgs = append(msgs, llm.Message{Role: m.Role, Content: m.Content})
	}

	turn, err := s.turner.ServeTurn(r.Context(), msgs, decision.Model, loop.ServeTurnOpts{
		MaxTokens:   req.MaxTokens,
		Temperature: req.Temperature,
	})
	if err != nil {
		// Log the failure class only; upstream bodies never reach the client.
		s.logger.Printf("chat model=%s remote=%s error=upstream duration=%s", req.Model, r.RemoteAddr, time.Since(start).Round(time.Millisecond))
		s.writeError(w, http.StatusBadGateway, "upstream model error")
		return
	}
	// One audit line per request: routing + counters, no message contents.
	s.logger.Printf("chat model=%s tier=%s concrete=%s rounds=%d tools_refused=%d tokens=%d duration=%s remote=%s",
		req.Model, decision.Tier, decision.Model, turn.Rounds, turn.Denied, turn.TotalTokens,
		time.Since(start).Round(time.Millisecond), r.RemoteAddr)
	s.writeJSON(w, http.StatusOK, chatCompletion{
		ID:      newCompletionID(),
		Object:  "chat.completion",
		Created: time.Now().Unix(),
		Model:   req.Model,
		Choices: []outboundChoice{{
			Index:        0,
			Message:      outboundMessage{Role: "assistant", Content: turn.Content},
			FinishReason: "stop",
		}},
		Usage: usageOut{
			PromptTokens:     turn.PromptTokens,
			CompletionTokens: turn.CompletionTokens,
			TotalTokens:      turn.TotalTokens,
		},
	})
}

// authorize checks the Bearer vkey in constant time. The 401 body is byte-
// identical for missing, malformed, and wrong keys (no oracle).
func (s *Server) authorize(r *http.Request) bool {
	const prefix = "Bearer "
	h := r.Header.Get("Authorization")
	if len(h) <= len(prefix) || !strings.EqualFold(h[:len(prefix)], prefix) {
		return false
	}
	token := h[len(prefix):]
	// Compare SHA-256 digests so the comparison time does not depend on the
	// presented key length.
	got := sha256.Sum256([]byte(token))
	want := sha256.Sum256([]byte(s.vkey))
	return subtle.ConstantTimeCompare(got[:], want[:]) == 1
}

func (s *Server) rejectAuth(w http.ResponseWriter) {
	s.writeError(w, http.StatusUnauthorized, "invalid api key")
}

func (s *Server) readBody(w http.ResponseWriter, r *http.Request) ([]byte, bool) {
	body, err := io.ReadAll(io.LimitReader(r.Body, maxBodyBytes+1))
	if err != nil {
		s.writeError(w, http.StatusBadRequest, "unreadable body")
		return nil, false
	}
	if len(body) > maxBodyBytes {
		s.writeError(w, http.StatusRequestEntityTooLarge, "payload too large")
		return nil, false
	}
	return body, true
}

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

// writeError emits an OpenAI-shaped error body.
func (s *Server) writeError(w http.ResponseWriter, code int, msg string) {
	s.writeJSON(w, code, errorBody{Error: errorDetail{Message: msg, Type: errorType(code)}})
}

func errorType(code int) string {
	switch {
	case code == http.StatusUnauthorized:
		return "authentication_error"
	case code >= 500:
		return "server_error"
	default:
		return "invalid_request_error"
	}
}

// inboundRequest is the accepted OpenAI chat-completions request subset:
// model, messages, temperature, max_tokens. Streaming is rejected; tools
// are ignored (v0 serve turns are pure chat).
type inboundRequest struct {
	Model       string           `json:"model"`
	Messages    []inboundMessage `json:"messages"`
	Temperature *float64         `json:"temperature"`
	MaxTokens   int              `json:"max_tokens"`
	Stream      bool             `json:"stream"`
}

type inboundMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type modelEntry struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

type modelList struct {
	Object string       `json:"object"`
	Data   []modelEntry `json:"data"`
}

type outboundMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type outboundChoice struct {
	Index        int             `json:"index"`
	Message      outboundMessage `json:"message"`
	FinishReason string          `json:"finish_reason"`
}

type usageOut struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type chatCompletion struct {
	ID      string           `json:"id"`
	Object  string           `json:"object"`
	Created int64            `json:"created"`
	Model   string           `json:"model"`
	Choices []outboundChoice `json:"choices"`
	Usage   usageOut         `json:"usage"`
}

type errorDetail struct {
	Message string `json:"message"`
	Type    string `json:"type"`
}

type errorBody struct {
	Error errorDetail `json:"error"`
}

// newCompletionID mints an OpenAI-shaped id (chatcmpl-<hex>).
func newCompletionID() string {
	var b [12]byte
	if _, err := rand.Read(b[:]); err != nil {
		return fmt.Sprintf("chatcmpl-%d", time.Now().UnixNano())
	}
	return "chatcmpl-" + hex.EncodeToString(b[:])
}