CrossRelay/internal/gateway/server.go

// Package gateway serves the Anthropic Messages API to Claude Code and
// the Responses API to Codex, plus the local management API to the
// CrossRelay UI.
package gateway

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"path/filepath"
	"sync"
	"time"

	"gitea.rutua.cn/kate522/CrossRelay/internal/anthropic"
	"gitea.rutua.cn/kate522/CrossRelay/internal/claudecfg"
	"gitea.rutua.cn/kate522/CrossRelay/internal/codexcfg"
	"gitea.rutua.cn/kate522/CrossRelay/internal/config"
	"gitea.rutua.cn/kate522/CrossRelay/internal/responses"
)

// Version is reported by /api/status.
const Version = "0.1.0"

// Stats aggregates gateway counters.
type Stats struct {
	RequestsTotal  int64 `json:"requests_total"`
	RequestsFailed int64 `json:"requests_failed"`
	InputTokens    int64 `json:"input_tokens"`
	OutputTokens   int64 `json:"output_tokens"`
}

// LogEntry is one recorded request.
type LogEntry struct {
	Time         time.Time `json:"time"`
	Method       string    `json:"method"`
	Path         string    `json:"path"`
	App          string    `json:"app,omitempty"`
	Provider     string    `json:"provider,omitempty"`
	Model        string    `json:"model,omitempty"`
	Status       int       `json:"status"`
	DurationMS   int64     `json:"duration_ms"`
	InputTokens  int       `json:"input_tokens,omitempty"`
	OutputTokens int       `json:"output_tokens,omitempty"`
	Error        string    `json:"error,omitempty"`
}

// Server wires the config store to the HTTP handlers.
type Server struct {
	store   *config.Store
	port    int
	started time.Time

	logs *logRing

	statsMu sync.Mutex
	stats   Stats
}

// New builds a server bound to the given port.
func New(store *config.Store, port int) *Server {
	return &Server{
		store:   store,
		port:    port,
		started: time.Now(),
		logs:    newLogRing(100),
	}
}

// Port returns the port the gateway listens on.
func (s *Server) Port() int { return s.port }

// Handler returns the full HTTP handler.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("POST /v1/messages", s.withLogging(s.withAuth(s.handleMessages)))
	mux.HandleFunc("POST /v1/messages/count_tokens", s.withAuth(s.handleCountTokens))
	mux.Handle("POST /v1/responses", s.withLogging(s.withAuth(s.handleResponses)))
	mux.HandleFunc("GET /api/status", s.handleStatus)
	mux.HandleFunc("POST /api/fetch-models", s.handleFetchModels)
	mux.HandleFunc("GET /api/providers", s.handleListProviders)
	mux.HandleFunc("POST /api/providers", s.handleCreateProvider)
	mux.HandleFunc("PUT /api/providers/{id}", s.handleUpdateProvider)
	mux.HandleFunc("DELETE /api/providers/{id}", s.handleDeleteProvider)
	mux.HandleFunc("POST /api/providers/{id}/apply", s.handleApply)
	mux.HandleFunc("POST /api/restore", s.handleRestore)
	mux.HandleFunc("GET /api/backups", s.handleBackups)
	mux.HandleFunc("GET /api/settings", s.handleGetSettings)
	mux.HandleFunc("PUT /api/settings", s.handleUpdateSettings)
	mux.HandleFunc("GET /api/logs", s.handleLogs)
	return mux
}

// withAuth protects the Anthropic-compatible endpoints with the local
// relay token.
func (s *Server) withAuth(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		token := s.store.Config().RelayToken
		if token != "" {
			got := bearerToken(r)
			if got == "" {
				got = r.Header.Get("x-api-key")
			}
			if got != token {
				writeAnthropicError(w, http.StatusUnauthorized, "authentication_error", "invalid relay token")
				return
			}
		}
		next(w, r)
	}
}

func bearerToken(r *http.Request) string {
	auth := r.Header.Get("Authorization")
	if len(auth) > 7 && auth[:7] == "Bearer " {
		return auth[7:]
	}
	return ""
}

// record carries per-request fields from handlers to the logging
// middleware.
type record struct {
	mu    sync.Mutex
	entry LogEntry
}

type recordKey struct{}

func recordFrom(r *http.Request) *record {
	rec, _ := r.Context().Value(recordKey{}).(*record)
	return rec
}

func (rec *record) update(fn func(*LogEntry)) {
	if rec == nil {
		return
	}
	rec.mu.Lock()
	defer rec.mu.Unlock()
	fn(&rec.entry)
}

func requestApp(path string) string {
	switch path {
	case "/v1/messages":
		return "Claude Code"
	case "/v1/responses":
		return "Codex"
	default:
		return ""
	}
}

func (s *Server) withLogging(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rec := &record{entry: LogEntry{
			Time:   start,
			Method: r.Method,
			Path:   r.URL.Path,
			App:    requestApp(r.URL.Path),
		}}
		ww := &statusWriter{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(ww, r.WithContext(context.WithValue(r.Context(), recordKey{}, rec)))

		rec.mu.Lock()
		rec.entry.Status = ww.status
		rec.entry.DurationMS = time.Since(start).Milliseconds()
		entry := rec.entry
		rec.mu.Unlock()

		s.logs.add(entry)
		s.statsMu.Lock()
		s.stats.RequestsTotal++
		if entry.Status >= 400 {
			s.stats.RequestsFailed++
		}
		s.stats.InputTokens += int64(entry.InputTokens)
		s.stats.OutputTokens += int64(entry.OutputTokens)
		s.statsMu.Unlock()
	})
}

type statusWriter struct {
	http.ResponseWriter
	status      int
	wroteHeader bool
}

func (w *statusWriter) WriteHeader(code int) {
	if !w.wroteHeader {
		w.status = code
		w.wroteHeader = true
	}
	w.ResponseWriter.WriteHeader(code)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	return w.ResponseWriter.Write(b)
}

// Flush keeps streaming responses unbuffered.
func (w *statusWriter) Flush() {
	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (s *Server) claudeManager() (*claudecfg.Manager, error) {
	path, err := claudecfg.DefaultPath()
	if err != nil {
		return nil, err
	}
	return claudecfg.New(path, filepath.Join(config.Dir(), "backups")), nil
}

func (s *Server) codexManager() (*codexcfg.Manager, error) {
	path, err := codexcfg.DefaultPath()
	if err != nil {
		return nil, err
	}
	return codexcfg.New(path, filepath.Join(config.Dir(), "backups")), nil
}

func (s *Server) snapshotStats() Stats {
	s.statsMu.Lock()
	defer s.statsMu.Unlock()
	return s.stats
}

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

func writeAnthropicError(w http.ResponseWriter, status int, errType, message string) {
	writeJSON(w, status, anthropic.ErrorResponse{
		Type:  "error",
		Error: anthropic.Detail{Type: errType, Message: message},
	})
}

func writeAPIError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]any{"error": message})
}

func writeResponsesError(w http.ResponseWriter, status int, errType, message string) {
	writeJSON(w, status, responses.ErrorResponse{
		Error: responses.Detail{Type: errType, Message: message},
	})
}

func decodeJSON(r *http.Request, v any, limit int64) error {
	dec := json.NewDecoder(io.LimitReader(r.Body, limit))
	return dec.Decode(v)
}