remap.go

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"regexp"
	"strings"
	"time"
)

// googleModelRegex extracts the model name from Google provider paths
var googleModelRegex = regexp.MustCompile(`/api/provider/google/.+/models/([^/:]+):(generateContent|streamGenerateContent)`)

// parseGoogleProviderRequest checks if this is a Google provider request and extracts the model name
func parseGoogleProviderRequest(r *http.Request) (model string, streaming bool, ok bool) {
	matches := googleModelRegex.FindStringSubmatch(r.URL.Path)
	if matches == nil {
		return "", false, false
	}
	model = matches[1]
	streaming = matches[2] == "streamGenerateContent" || r.URL.Query().Get("alt") == "sse"
	return model, streaming, true
}

// isUnsupportedProviderRequest checks if this is a request to a provider we don't support
// (anything that's not anthropic or openai)
func isUnsupportedProviderRequest(r *http.Request) (provider string, ok bool) {
	if !strings.HasPrefix(r.URL.Path, "/api/provider/") {
		return "", false
	}
	rest := strings.TrimPrefix(r.URL.Path, "/api/provider/")
	parts := strings.SplitN(rest, "/", 2)
	if len(parts) == 0 {
		return "", false
	}
	provider = parts[0]
	if provider == "anthropic" || provider == "openai" {
		return "", false
	}
	return provider, true
}

// handleRemappedRequest handles a request that needs model remapping
func (ph *ProxyHandler) handleRemappedRequest(w http.ResponseWriter, r *http.Request, reqID uint64, model string, streaming bool, remap ModelRemapConfig, isExplicit bool) {
	start := time.Now()

	if ph.gateway == nil || !ph.gateway.IsReady() {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprint(w, `{"error":"provider_not_ready","message":"Provider gateway is starting up. Please retry."}`)
		return
	}

	if !isExplicit {
		slog.Warn("unmapped model, using fallback", "reqID", reqID, "model", model, "targetProvider", remap.Provider, "targetModel", remap.To)
	}

	// Read the original request body
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, `{"error":"failed to read request body"}`, http.StatusBadRequest)
		slog.Error("failed to read body", "reqID", reqID, "error", err)
		ph.metrics.errors.Add(1)
		return
	}

	// Parse the Google GenAI request
	var googleReq map[string]interface{}
	if err := json.Unmarshal(body, &googleReq); err != nil {
		http.Error(w, `{"error":"invalid JSON body"}`, http.StatusBadRequest)
		slog.Error("invalid JSON", "reqID", reqID, "error", err)
		ph.metrics.errors.Add(1)
		return
	}

	// Translate based on target provider
	var translatedBody []byte
	switch remap.Provider {
	case "anthropic":
		translatedBody, err = translateGoogleToAnthropic(googleReq, remap.To, streaming)
	case "openai":
		translatedBody, err = translateGoogleToOpenAI(googleReq, remap.To, streaming)
	default:
		http.Error(w, `{"error":"unsupported target provider"}`, http.StatusInternalServerError)
		return
	}
	if err != nil {
		http.Error(w, fmt.Sprintf(`{"error":"translation failed: %s"}`, err.Error()), http.StatusInternalServerError)
		slog.Error("translation failed", "reqID", reqID, "error", err)
		ph.metrics.errors.Add(1)
		return
	}
	if err := validateTranslatedToolResults(remap.Provider, translatedBody); err != nil {
		http.Error(w, fmt.Sprintf(`{"error":"translated request validation failed: %s"}`, err.Error()), http.StatusInternalServerError)
		slog.Error("translated request validation failed", "reqID", reqID, "provider", remap.Provider, "error", err)
		ph.metrics.errors.Add(1)
		return
	}

	// Build the outbound request — route through the embedded provider gateway
	targetPath := TargetPathForProvider(remap.Provider)
	targetURL := ph.gateway.targetURL + targetPath
	outReq, err := http.NewRequestWithContext(r.Context(), "POST", targetURL, bytes.NewReader(translatedBody))
	if err != nil {
		http.Error(w, `{"error":"failed to create request"}`, http.StatusInternalServerError)
		return
	}
	outReq.Header.Set("Content-Type", "application/json")

	// Forward auth headers
	if auth := r.Header.Get("Authorization"); auth != "" {
		outReq.Header.Set("Authorization", auth)
	}
	for _, cookie := range r.Cookies() {
		outReq.AddCookie(cookie)
	}
	for name, values := range r.Header {
		if strings.HasPrefix(name, "X-Amp-") {
			for _, v := range values {
				outReq.Header.Set(name, v)
			}
		}
	}

	// Set Anthropic-specific headers
	if remap.Provider == "anthropic" {
		outReq.Header.Set("Anthropic-Version", "2023-06-01")
	}

	if len(translatedBody) <= 2000 {
		slog.Debug("remap translated body", "reqID", reqID, "body", string(translatedBody))
	} else {
		slog.Debug("remap translated body", "reqID", reqID, "body", string(translatedBody[:1000])+"...", "totalBytes", len(translatedBody))
	}

	slog.Info("remap request", "reqID", reqID, "targetURL", targetURL, "model", remap.To, "stream", streaming)

	// Make the request
	resp, err := ph.gateway.Do(outReq)
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadGateway)
		fmt.Fprintf(w, `{"error":"provider_unavailable","message":"Provider backend error: %s"}`, err.Error())
		slog.Error("upstream request failed", "reqID", reqID, "error", err)
		ph.metrics.errors.Add(1)
		return
	}
	defer resp.Body.Close()

	if !streaming {
		ph.handleNonStreamingResponse(w, resp, reqID, remap, start)
	} else {
		ph.handleStreamingResponse(w, resp, reqID, remap, start)
	}
}

func validateTranslatedToolResults(provider string, translatedBody []byte) error {
	var translatedReq map[string]interface{}
	if err := json.Unmarshal(translatedBody, &translatedReq); err != nil {
		return fmt.Errorf("invalid translated request JSON: %w", err)
	}

	switch provider {
	case "anthropic":
		return validateAnthropicToolResults(translatedReq)
	case "openai":
		return validateOpenAIToolResults(translatedReq)
	default:
		return nil
	}
}

func validateAnthropicToolResults(translatedReq map[string]interface{}) error {
	messages, ok := translatedReq["messages"].([]interface{})
	if !ok {
		return nil
	}

	seenToolResults := make(map[string]struct{})
	for _, rawMsg := range messages {
		msg, ok := rawMsg.(map[string]interface{})
		if !ok {
			continue
		}

		contentBlocks, ok := msg["content"].([]interface{})
		if !ok {
			continue
		}

		for _, rawBlock := range contentBlocks {
			block, ok := rawBlock.(map[string]interface{})
			if !ok {
				continue
			}
			if blockType, _ := block["type"].(string); blockType != "tool_result" {
				continue
			}

			toolUseID, _ := block["tool_use_id"].(string)
			if toolUseID == "" {
				return fmt.Errorf("tool_result missing tool_use_id")
			}
			if toolUseID == "toolu_unknown" {
				return fmt.Errorf("tool_result has unresolved tool_use_id")
			}
			if _, exists := seenToolResults[toolUseID]; exists {
				return fmt.Errorf("duplicate tool_result for tool_use_id %q", toolUseID)
			}
			seenToolResults[toolUseID] = struct{}{}
		}
	}

	return nil
}

func validateOpenAIToolResults(translatedReq map[string]interface{}) error {
	messages, ok := translatedReq["messages"].([]interface{})
	if !ok {
		return nil
	}

	seenToolResults := make(map[string]struct{})
	for _, rawMsg := range messages {
		msg, ok := rawMsg.(map[string]interface{})
		if !ok {
			continue
		}
		if role, _ := msg["role"].(string); role != "tool" {
			continue
		}

		toolCallID, _ := msg["tool_call_id"].(string)
		if toolCallID == "" {
			return fmt.Errorf("tool message missing tool_call_id")
		}
		if toolCallID == "call_unknown" {
			return fmt.Errorf("tool message has unresolved tool_call_id")
		}
		if _, exists := seenToolResults[toolCallID]; exists {
			return fmt.Errorf("duplicate tool_result for tool_call_id %q", toolCallID)
		}
		seenToolResults[toolCallID] = struct{}{}
	}

	return nil
}

func (ph *ProxyHandler) handleNonStreamingResponse(w http.ResponseWriter, resp *http.Response, reqID uint64, remap ModelRemapConfig, start time.Time) {
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		http.Error(w, `{"error":"failed to read upstream response"}`, http.StatusBadGateway)
		return
	}

	if resp.StatusCode >= 400 {
		slog.Error("upstream error", "reqID", reqID, "status", resp.StatusCode, "body", string(respBody))
		ph.metrics.errors.Add(1)
		for k, v := range resp.Header {
			for _, vv := range v {
				w.Header().Add(k, vv)
			}
		}
		w.WriteHeader(resp.StatusCode)
		w.Write(respBody)
		slog.Info("response remap", "reqID", reqID, "status", resp.StatusCode, "statusText", http.StatusText(resp.StatusCode), "bytes", len(respBody), "elapsed", time.Since(start).Round(time.Millisecond))
		return
	}

	var upstreamResp map[string]interface{}
	if err := json.Unmarshal(respBody, &upstreamResp); err != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(resp.StatusCode)
		w.Write(respBody)
		return
	}

	var translated []byte
	switch remap.Provider {
	case "anthropic":
		translated, err = translateAnthropicToGoogle(upstreamResp)
	case "openai":
		translated, err = translateOpenAIToGoogle(upstreamResp)
	}
	if err != nil {
		slog.Warn("response translation failed, forwarding raw", "reqID", reqID, "error", err)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(resp.StatusCode)
		w.Write(respBody)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	w.Write(translated)
	slog.Info("response remap", "reqID", reqID, "status", 200, "statusText", "OK", "bytes", len(translated), "elapsed", time.Since(start).Round(time.Millisecond))
}

func (ph *ProxyHandler) handleStreamingResponse(w http.ResponseWriter, resp *http.Response, reqID uint64, remap ModelRemapConfig, start time.Time) {
	if resp.StatusCode >= 400 {
		respBody, _ := io.ReadAll(resp.Body)
		w.WriteHeader(resp.StatusCode)
		w.Write(respBody)
		slog.Error("response remap error", "reqID", reqID, "status", resp.StatusCode, "statusText", http.StatusText(resp.StatusCode), "elapsed", time.Since(start).Round(time.Millisecond))
		ph.metrics.errors.Add(1)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	flusher, canFlush := w.(http.Flusher)
	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 256*1024), 256*1024)

	var totalBytes int64
	var eventType string

	for scanner.Scan() {
		line := scanner.Text()

		if strings.HasPrefix(line, "event: ") {
			eventType = strings.TrimPrefix(line, "event: ")
			continue
		}

		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")

		if data == "[DONE]" {
			break
		}

		var chunk map[string]interface{}
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}

		var googleChunk []byte
		var err error
		switch remap.Provider {
		case "anthropic":
			googleChunk, err = translateAnthropicStreamChunk(chunk, eventType)
		case "openai":
			googleChunk, err = translateOpenAIStreamChunk(chunk)
		}
		if err != nil || googleChunk == nil {
			continue
		}

		sseFrame := fmt.Sprintf("data: %s\n\n", googleChunk)
		n, _ := w.Write([]byte(sseFrame))
		totalBytes += int64(n)
		if canFlush {
			flusher.Flush()
		}
	}

	slog.Info("response remap stream", "reqID", reqID, "status", 200, "statusText", "OK", "bytes", totalBytes, "elapsed", time.Since(start).Round(time.Millisecond))
}