central_ai_llm_master/src/main.go

// Command hub is the Central AI Hub service.
//
// It serves the three frozen inbound routes from docs/hub-api.md:
//
//	POST /ask      — hear, understand, check, route, answer
//	POST /approve  — decide one Hub-owned approval
//	GET  /health   — liveness only
//
// The pipeline lives in package hub. This file is the HTTP layer only.
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"hub/apierr"
	"hub/check"
	"hub/config"
	"hub/destinations"
	"hub/hear"
	"hub/hub"
	"hub/logbook"
	"hub/route"
	"hub/understand"
)

const (
	serviceName    = "hub"
	serviceVersion = "1"

	maxBodyBytes = 64 * 1024
)

// server is the HTTP surface.
type server struct {
	svc  *hub.Service
	cfg  *config.Config
	rate *rateLimiter
}

func main() {
	hc := flag.Bool("healthcheck", false, "Call this process's /health and exit.")
	flag.Parse()

	cfg := config.Load()
	if *hc {
		healthcheck(cfg)
		return
	}

	rules, err := understand.LoadRules(cfg.RoutingFile)
	if err != nil {
		log.Printf("routing: using built-in rules (%v)", err)
	}

	lb, err := logbook.New(cfg.LogPath, cfg.LogMaxBytes, cfg.LogKeepDays)
	if err != nil {
		log.Fatalf("cannot open log %s: %v", cfg.LogPath, err)
	}
	defer lb.Close()

	projects := map[string]route.ProjectQuerier{}
	for name, p := range cfg.Projects {
		projects[name] = destinations.NewProjectClient(name, p.URL, p.Token, cfg.ProjectTimeout)
	}
	speaker := destinations.NewSpeaker(cfg.JervisURL, cfg.TTSTimeout)
	paperclip := destinations.NewPaperclip(cfg.PaperclipURL, cfg.PaperclipKey, cfg.PaperclipCompanyID, cfg.ApproveTimeout)
	approvals := hub.NewApprovals(10 * time.Minute)

	dispatcher := &route.Dispatcher{
		Rules:     rules,
		Projects:  projects,
		Speaker:   speaker,
		Paperclip: paperclip,
		Approvals: approvals,
	}
	svc := &hub.Service{
		Config:      cfg,
		Rules:       rules,
		Router:      understand.NewOmniRouter(cfg.OmniURL, cfg.OmniKey, cfg.OmniModel, cfg.LLMTimeout),
		Dispatcher:  dispatcher,
		Approvals:   approvals,
		Log:         lb,
		Speaker:     speaker,
		Paperclip:   paperclip,
		Started:     time.Now(),
		HealthProbe: omniProbe(cfg),
	}

	s := &server{svc: svc, cfg: cfg, rate: newRateLimiter(cfg.RatePerMin)}

	srv := &http.Server{
		Addr:              ":" + cfg.Port,
		Handler:           s.routes(),
		ReadHeaderTimeout: 1 * time.Second,
		ReadTimeout:       cfg.AskTimeout,
		// WriteTimeout must be longer than the /ask deadline, or the server can
		// cut the reply off when the handler finishes at the deadline.
		WriteTimeout: cfg.AskTimeout + 3*time.Second,
		IdleTimeout:  60 * time.Second,
	}

	log.Printf("%s %s listening on :%s", serviceName, serviceVersion, cfg.Port)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("%s stopped: %v", serviceName, err)
	}
}

// routes builds the HTTP multiplexer.
func (s *server) routes() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("/ask", s.handleAsk)
	mux.HandleFunc("/approve", s.handleApprove)
	mux.HandleFunc("/health", s.handleHealth)
	mux.HandleFunc("/", notFound)
	return mux
}

// handleAsk serves POST /ask.
func (s *server) handleAsk(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "", "bad_request", "Use POST for this route.")
		return
	}
	if !requireJSON(w, r) {
		return
	}
	source, ok := s.auth(w, r)
	if !ok {
		return
	}
	if ok, retry := s.rate.allow(source); !ok {
		w.Header().Set("Retry-After", fmt.Sprintf("%d", retry))
		writeError(w, http.StatusTooManyRequests, "", "rate_limited", "Too many requests. Try again shortly.")
		return
	}

	body, ok := readBody(w, r)
	if !ok {
		return
	}
	req, apiErr := hear.Parse(body)
	if apiErr != nil {
		writeAPIError(w, apiErr, "")
		return
	}
	if apiErr := check.Ask(req, source); apiErr != nil {
		writeAPIError(w, apiErr, req.RequestID)
		return
	}

	start := time.Now()
	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.AskTimeout)
	defer cancel()

	resp := s.svc.Ask(ctx, req)
	s.svc.LogAsk(req, resp, time.Since(start))
	writeJSON(w, http.StatusOK, resp)
}

// handleApprove serves POST /approve. Only HUB_TOKEN_ADMIN works.
func (s *server) handleApprove(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "", "bad_request", "Use POST for this route.")
		return
	}
	if !requireJSON(w, r) {
		return
	}
	source, ok := s.auth(w, r)
	if !ok {
		return
	}
	if ok, retry := s.rate.allow(source); !ok {
		w.Header().Set("Retry-After", fmt.Sprintf("%d", retry))
		writeError(w, http.StatusTooManyRequests, "", "rate_limited", "Too many requests. Try again shortly.")
		return
	}
	if apiErr := check.Admin(source); apiErr != nil {
		writeAPIError(w, apiErr, "")
		return
	}

	body, ok := readBody(w, r)
	if !ok {
		return
	}
	req, apiErr := hub.ParseApprove(body)
	if apiErr != nil {
		writeAPIError(w, apiErr, "")
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.ApproveTimeout)
	defer cancel()

	resp, apiErr := s.svc.Approve(ctx, req)
	if apiErr != nil {
		writeAPIError(w, apiErr, "")
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

// handleHealth serves GET /health. No authentication.
func (s *server) handleHealth(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "", "bad_request", "Use GET for this route.")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.HealthTimeout)
	defer cancel()

	resp := s.svc.Health(ctx)
	code := http.StatusOK
	if !resp.OK {
		code = http.StatusServiceUnavailable
	}
	writeJSON(w, code, resp)
}

// notFound answers an unknown path with the frozen error envelope.
func notFound(w http.ResponseWriter, r *http.Request) {
	writeError(w, http.StatusNotFound, "", "not_found", "Unknown path.")
}

// auth reads the bearer token and returns the client source. It writes the 401
// reply itself when the token is missing or unknown.
func (s *server) auth(w http.ResponseWriter, r *http.Request) (string, bool) {
	token := bearer(r)
	if token == "" {
		writeError(w, http.StatusUnauthorized, "", "unauthorized", "A bearer token is required.")
		return "", false
	}
	source, ok := s.cfg.SourceForToken(token)
	if !ok {
		writeError(w, http.StatusUnauthorized, "", "unauthorized", "The bearer token is not valid.")
		return "", false
	}
	return source, true
}

// bearer reads the token from the Authorization header.
func bearer(r *http.Request) string {
	h := r.Header.Get("Authorization")
	const prefix = "Bearer "
	if len(h) <= len(prefix) || !strings.EqualFold(h[:len(prefix)], prefix) {
		return ""
	}
	return strings.TrimSpace(h[len(prefix):])
}

// requireJSON checks the content type and returns false after writing 415.
func requireJSON(w http.ResponseWriter, r *http.Request) bool {
	ct := r.Header.Get("Content-Type")
	if ct == "" {
		writeError(w, http.StatusUnsupportedMediaType, "", "unsupported_media_type", "Send the body as application/json.")
		return false
	}
	mediaType := strings.ToLower(strings.TrimSpace(strings.Split(ct, ";")[0]))
	if mediaType != "application/json" {
		writeError(w, http.StatusUnsupportedMediaType, "", "unsupported_media_type", "Send the body as application/json.")
		return false
	}
	return true
}

// readBody reads at most 64 KiB. A larger body gets 413.
func readBody(w http.ResponseWriter, r *http.Request) ([]byte, bool) {
	r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes)
	body, err := io.ReadAll(r.Body)
	if err != nil {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			writeError(w, http.StatusRequestEntityTooLarge, "", "payload_too_large", "The request body is over 64 KiB.")
			return nil, false
		}
		writeError(w, http.StatusBadRequest, "", "bad_request", "The request body is not valid JSON.")
		return nil, false
	}
	return body, true
}

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

type errorBody struct {
	Status string    `json:"status"`
	Error  errorInfo `json:"error"`
}

type errorInfo struct {
	Code      string `json:"code"`
	Message   string `json:"message"`
	RequestID string `json:"request_id,omitempty"`
}

func writeError(w http.ResponseWriter, code int, requestID, errCode, message string) {
	writeJSON(w, code, errorBody{
		Status: "error",
		Error:  errorInfo{Code: errCode, Message: message, RequestID: requestID},
	})
}

// writeAPIError writes an *apierr.Error with the frozen envelope.
func writeAPIError(w http.ResponseWriter, e *apierr.Error, requestID string) {
	writeError(w, e.HTTP, requestID, e.Code, e.Message)
}

// rateLimiter is a fixed-window cap per client token (docs/hub-api.md section
// 4). It protects the Hub and is not a security control.
type rateLimiter struct {
	limit int
	mu    sync.Mutex
	win   map[string]*rateWindow
}

type rateWindow struct {
	start time.Time
	count int
}

func newRateLimiter(limit int) *rateLimiter {
	if limit <= 0 {
		limit = 30
	}
	return &rateLimiter{limit: limit, win: map[string]*rateWindow{}}
}

// allow reports whether the key may proceed. It returns the seconds to wait
// when it may not.
func (rl *rateLimiter) allow(key string) (bool, int) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	w, ok := rl.win[key]
	if !ok || now.Sub(w.start) >= time.Minute {
		rl.win[key] = &rateWindow{start: now, count: 1}
		return true, 0
	}
	if w.count >= rl.limit {
		retry := int(time.Until(w.start.Add(time.Minute)).Seconds())
		if retry < 1 {
			retry = 1
		}
		return false, retry
	}
	w.count++
	return true, 0
}

// omniProbe checks OmniRoute for GET /health. It never blocks over 1 second.
func omniProbe(cfg *config.Config) func(ctx context.Context) map[string]string {
	return func(ctx context.Context) map[string]string {
		if cfg.OmniKey == "" {
			return map[string]string{"omniroute": "unknown"}
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, cfg.OmniURL+"/v1/models", nil)
		if err != nil {
			return map[string]string{"omniroute": "down"}
		}
		req.Header.Set("Authorization", "Bearer "+cfg.OmniKey)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return map[string]string{"omniroute": "down"}
		}
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK {
			return map[string]string{"omniroute": "ok"}
		}
		return map[string]string{"omniroute": "down"}
	}
}

// healthcheck calls this process's own /health. It is the container health
// probe, because the runtime image has no shell and no curl.
func healthcheck(cfg *config.Config) {
	url := fmt.Sprintf("http://127.0.0.1:%s/health", cfg.Port)
	client := &http.Client{Timeout: 1 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		log.Printf("healthcheck failed: %v", err)
		os.Exit(1)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		log.Printf("healthcheck status: %d", resp.StatusCode)
		os.Exit(1)
	}
	os.Exit(0)
}