ultisuite-backend/internal/api/ai/mcp_proxy.go

package aiapi

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/jackc/pgx/v5"

	"github.com/ultisuite/ulti-backend/internal/ai"
	"github.com/ultisuite/ulti-backend/internal/api/apiresponse"
	"github.com/ultisuite/ulti-backend/internal/apitokens"
	"github.com/ultisuite/ulti-backend/internal/users"
)

func (h *Handler) MCPProxy(w http.ResponseWriter, r *http.Request) {
	if h.cfg == nil || strings.TrimSpace(h.cfg.UltimailMCPURL) == "" {
		apiresponse.WriteError(w, r, http.StatusServiceUnavailable, apiresponse.CodeInternal, "mcp not configured", nil)
		return
	}

	token, enabledTools, err := h.resolveMCPToken(r)
	if err != nil {
		apiresponse.WriteError(w, r, http.StatusUnauthorized, apiresponse.CodeAuthUnauthorized, err.Error(), nil)
		return
	}

	target, err := url.Parse(strings.TrimRight(strings.TrimSpace(h.cfg.UltimailMCPURL), "/"))
	if err != nil {
		apiresponse.WriteError(w, r, http.StatusInternalServerError, apiresponse.CodeInternal, err.Error(), nil)
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(target)
	origDirector := proxy.Director
	upstreamPath := mapMCPUpstreamPath(r.URL.Path)
	proxy.Director = func(req *http.Request) {
		origDirector(req)
		req.Host = target.Host
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = upstreamPath
		req.Header.Set("X-Ulti-Token", token)
		if len(enabledTools) > 0 {
			req.Header.Set("X-Ulti-Enabled-Tools", strings.Join(enabledTools, ","))
		}
	}
	proxy.ModifyResponse = func(resp *http.Response) error {
		resp.Header.Del("Access-Control-Allow-Origin")
		return nil
	}
	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		apiresponse.WriteError(w, r, http.StatusBadGateway, apiresponse.CodeInternal, err.Error(), nil)
	}

	proxy.ServeHTTP(w, r)
}

func (h *Handler) resolveMCPToken(r *http.Request) (token string, enabledTools []string, err error) {
	ctx := r.Context()
	enabledTools = h.loadEnabledTools(ctx)

	if headerToken := strings.TrimSpace(r.Header.Get("X-Ulti-Token")); headerToken != "" {
		if _, authErr := apitokens.Authenticate(ctx, h.db, headerToken); authErr == nil {
			return headerToken, enabledTools, nil
		}
	}

	if bearer := bearerToken(r); strings.HasPrefix(bearer, apitokens.TokenPrefix()) {
		if _, authErr := apitokens.Authenticate(ctx, h.db, bearer); authErr == nil {
			return bearer, enabledTools, nil
		}
	}

	if bearer := bearerToken(r); h.cfg != nil && h.cfg.AIGatewayAPIKey != "" && bearer == h.cfg.AIGatewayAPIKey {
		email := openWebUIUserEmail(r)
		if email == "" {
			return "", nil, fmt.Errorf("missing openwebui user email")
		}
		created, createErr := h.createMCPSessionForEmail(ctx, email, apitokens.ChatSessionStandalone)
		if createErr != nil {
			return "", nil, createErr
		}
		return created.TokenSecret, enabledTools, nil
	}

	if claims, ok := h.resolveClaims(r); ok && strings.TrimSpace(claims.Sub) != "" {
		created, createErr := apitokens.CreateChatSession(ctx, h.db, claims.Sub, claims.Email, apitokens.ChatSessionInput{
			Preset: apitokens.ChatSessionStandalone,
		})
		if createErr != nil {
			return "", nil, createErr
		}
		return created.TokenSecret, enabledTools, nil
	}

	return "", nil, fmt.Errorf("unauthorized")
}

func (h *Handler) createMCPSessionForEmail(ctx context.Context, email string, preset apitokens.ChatSessionPreset) (apitokens.CreatedToken, error) {
	externalID, storedEmail, err := users.LookupIdentityByEmail(ctx, h.db, email)
	if err != nil {
		if err == pgx.ErrNoRows {
			return apitokens.CreatedToken{}, fmt.Errorf("user not found")
		}
		return apitokens.CreatedToken{}, err
	}
	return apitokens.CreateChatSession(ctx, h.db, externalID, storedEmail, apitokens.ChatSessionInput{
		Preset: preset,
	})
}

func openWebUIUserEmail(r *http.Request) string {
	for _, key := range []string{
		"X-OpenWebUI-User-Email",
		"X-Ulti-User-Email",
	} {
		if email := strings.TrimSpace(r.Header.Get(key)); email != "" {
			return email
		}
	}
	return ""
}

func (h *Handler) loadEnabledTools(ctx context.Context) []string {
	deployEnabled := h.cfg != nil && h.cfg.AIAssistantEnabled
	policy, _ := ai.LoadAssistantPolicy(ctx, h.db)
	if !deployEnabled && !policy.Enabled {
		return policy.EnabledTools
	}
	if len(policy.EnabledTools) == 0 {
		return []string{"mail", "drive", "contacts", "agenda", "search", "web_search", "docs"}
	}
	return policy.EnabledTools
}

func mapMCPUpstreamPath(path string) string {
	for _, prefix := range []string{"/api/v1/ai/mcp", "/mcp"} {
		if !strings.HasPrefix(path, prefix) {
			continue
		}
		suffix := strings.TrimPrefix(path, prefix)
		if suffix == "" || suffix == "/" {
			return "/mcp"
		}
		if suffix == "/messages" || suffix == "/sse" {
			return suffix
		}
		return "/mcp" + suffix
	}
	return "/mcp"
}

func (h *Handler) publicMCPPath() string {
	return "/api/v1/ai/mcp"
}

func (h *Handler) writeMCPConfigFields(ctx context.Context, out map[string]any) {
	out["mcp_url"] = h.publicMCPPath()
	out["enabled_tools"] = h.loadEnabledTools(ctx)
}

// MCPHealth proxies ultimail-mcp /health without auth (for compose healthchecks).
func (h *Handler) MCPHealth(w http.ResponseWriter, r *http.Request) {
	if h.cfg == nil || strings.TrimSpace(h.cfg.UltimailMCPURL) == "" {
		apiresponse.WriteJSON(w, http.StatusOK, map[string]any{"ok": false, "reason": "mcp not configured"})
		return
	}
	target := strings.TrimRight(strings.TrimSpace(h.cfg.UltimailMCPURL), "/") + "/health"
	resp, err := http.Get(target)
	if err != nil {
		apiresponse.WriteJSON(w, http.StatusOK, map[string]any{"ok": false, "reason": err.Error()})
		return
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	_, _ = w.Write(body)
}