CHORUS/pkg/seqthink/proxy/server.go

package proxy

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	"chorus/pkg/seqthink/mcpclient"
	"chorus/pkg/seqthink/observability"
	"chorus/pkg/seqthink/policy"
	"github.com/gorilla/mux"
	"github.com/rs/zerolog/log"
)

// ServerConfig holds the proxy server configuration
type ServerConfig struct {
	MCPClient      *mcpclient.Client
	Metrics        *observability.Metrics
	MaxBodyMB      int
	AgeIdentPath   string
	AgeRecipsPath  string
	KachingJWKSURL string
	RequiredScope  string
}

// Server is the proxy server handling requests
type Server struct {
	config         ServerConfig
	router         *mux.Router
	authMiddleware *policy.AuthMiddleware
}

// NewServer creates a new proxy server
func NewServer(cfg ServerConfig) (*Server, error) {
	s := &Server{
		config: cfg,
		router: mux.NewRouter(),
	}

	// Setup policy enforcement if configured
	if cfg.KachingJWKSURL != "" && cfg.RequiredScope != "" {
		log.Info().
			Str("jwks_url", cfg.KachingJWKSURL).
			Str("required_scope", cfg.RequiredScope).
			Msg("Policy enforcement enabled")

		validator := policy.NewValidator(cfg.KachingJWKSURL, cfg.RequiredScope)

		// Pre-fetch JWKS
		if err := validator.RefreshJWKS(); err != nil {
			log.Warn().Err(err).Msg("Failed to pre-fetch JWKS, will retry on first request")
		}

		s.authMiddleware = policy.NewAuthMiddleware(validator, cfg.Metrics.IncrementPolicyDenials)
	} else {
		log.Warn().Msg("Policy enforcement disabled - no JWKS URL or required scope configured")
		s.authMiddleware = policy.NewAuthMiddleware(nil, cfg.Metrics.IncrementPolicyDenials)
	}

	// Setup routes
	s.setupRoutes()

	return s, nil
}

// Handler returns the HTTP handler
func (s *Server) Handler() http.Handler {
	return s.router
}

// setupRoutes configures the HTTP routes
func (s *Server) setupRoutes() {
	// Health checks (no auth required)
	s.router.HandleFunc("/health", s.handleHealth).Methods("GET")
	s.router.HandleFunc("/ready", s.handleReady).Methods("GET")

	// MCP tool endpoint - route based on encryption config, with auth
	if s.isEncryptionEnabled() {
		log.Info().Msg("Encryption enabled - using encrypted endpoint")
		s.router.Handle("/mcp/tool",
			s.authMiddleware.Wrap(http.HandlerFunc(s.handleToolCallEncrypted))).Methods("POST")
	} else {
		log.Warn().Msg("Encryption disabled - using plaintext endpoint")
		s.router.Handle("/mcp/tool",
			s.authMiddleware.Wrap(http.HandlerFunc(s.handleToolCall))).Methods("POST")
	}

	// SSE endpoint - route based on encryption config, with auth
	if s.isEncryptionEnabled() {
		s.router.Handle("/mcp/sse",
			s.authMiddleware.Wrap(http.HandlerFunc(s.handleSSEEncrypted))).Methods("GET")
	} else {
		s.router.Handle("/mcp/sse",
			s.authMiddleware.Wrap(http.HandlerFunc(s.handleSSEPlaintext))).Methods("GET")
	}

	// Metrics endpoint (no auth required for internal monitoring)
	s.router.Handle("/metrics", s.config.Metrics.Handler())
}

// isEncryptionEnabled checks if encryption is configured
func (s *Server) isEncryptionEnabled() bool {
	return s.config.AgeIdentPath != "" && s.config.AgeRecipsPath != ""
}

// handleHealth returns 200 OK if wrapper is running
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

// handleReady checks if MCP server is ready
func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
	defer cancel()

	if err := s.config.MCPClient.Health(ctx); err != nil {
		log.Error().Err(err).Msg("MCP server not ready")
		http.Error(w, "MCP server not ready", http.StatusServiceUnavailable)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("READY"))
}

// handleToolCall proxies tool calls to MCP server (plaintext for Beat 1)
func (s *Server) handleToolCall(w http.ResponseWriter, r *http.Request) {
	s.config.Metrics.IncrementRequests()
	startTime := time.Now()

	// Limit request body size
	r.Body = http.MaxBytesReader(w, r.Body, int64(s.config.MaxBodyMB)*1024*1024)

	// Read request body
	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Error().Err(err).Msg("Failed to read request body")
		s.config.Metrics.IncrementErrors()
		http.Error(w, "Failed to read request", http.StatusBadRequest)
		return
	}

	// Parse tool request
	var toolReq mcpclient.ToolRequest
	if err := json.Unmarshal(body, &toolReq); err != nil {
		log.Error().Err(err).Msg("Failed to parse tool request")
		s.config.Metrics.IncrementErrors()
		http.Error(w, "Invalid request format", http.StatusBadRequest)
		return
	}

	log.Info().
		Str("tool", toolReq.Tool).
		Msg("Proxying tool call to MCP server")

	// Call MCP server
	ctx, cancel := context.WithTimeout(r.Context(), 120*time.Second)
	defer cancel()

	toolResp, err := s.config.MCPClient.CallTool(ctx, &toolReq)
	if err != nil {
		log.Error().Err(err).Msg("MCP tool call failed")
		s.config.Metrics.IncrementErrors()
		http.Error(w, fmt.Sprintf("Tool call failed: %v", err), http.StatusInternalServerError)
		return
	}

	// Return response
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(toolResp); err != nil {
		log.Error().Err(err).Msg("Failed to encode response")
		s.config.Metrics.IncrementErrors()
		return
	}

	duration := time.Since(startTime)
	log.Info().
		Str("tool", toolReq.Tool).
		Dur("duration", duration).
		Msg("Tool call completed")
}