proxx/handler.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

// Config holds the application configuration
type Config struct {
	Port        int    `yaml:"port"`
	UpstreamURL string `yaml:"upstream_url"`
}

var config *Config

// blockedHeaders are headers that should never be forwarded to upstream
// for security/privacy reasons. These headers could leak internal URLs,
// session information, or other sensitive data.
var blockedHeaders = map[string]bool{
	"Referer":        true, // Don't leak internal URLs to external API
	"Cookie":         true, // Don't forward session cookies
	"Authorization":   true, // Already extracted and sent as x-api-key
	"X-Forwarded-For": true, // Don't leak client IP
	"X-Real-Ip":       true, // Don't leak client IP
	"X-Forwarded-Host": true, // Don't leak internal hostnames
}

// ClaudeCodeHeaders returns the headers to mimic claude-code CLI
func ClaudeCodeHeaders(apiKey, sessionID string) map[string]string {
	return map[string]string{
		"User-Agent":               "claude-cli/1.0.18 (pro, cli)",
		"x-api-key":                apiKey,
		"x-app":                    "cli",
		"anthropic-version":        "2023-06-01",
		"anthropic-beta":           "interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,context-management-2025-06-27",
		"X-Claude-Code-Session-Id": sessionID,
		"content-type":             "application/json",
	}
}

func handleModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", "method_not_allowed")
		return
	}

	models := []map[string]interface{}{
		{"id": "claude-opus-4-5", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
		{"id": "claude-sonnet-4-20250514", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
		{"id": "claude-3-5-sonnet-20241022", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
		{"id": "claude-3-opus-20240229", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
		{"id": "claude-3-sonnet-20240229", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
		{"id": "claude-3-haiku-20240307", "object": "model", "created": 1234567890, "owned_by": "anthropic"},
	}

	response := map[string]interface{}{
		"object": "list",
		"data":   models,
	}

	w.Header().Set("content-type", "application/json")
	json.NewEncoder(w).Encode(response)
}

func handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", "method_not_allowed")
		return
	}

	// Extract Bearer token
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
		writeError(w, http.StatusUnauthorized, "Missing or invalid Authorization header", "authentication_error", "missing_authorization")
		return
	}
	apiKey := strings.TrimPrefix(authHeader, "Bearer ")

	// Read body
	body, err := io.ReadAll(r.Body)
	if err != nil {
		writeError(w, http.StatusBadRequest, "Failed to read request body", "invalid_request_error", "body_read_error")
		return
	}

	// Decode request
	var req ChatCompletionRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeError(w, http.StatusBadRequest, "Invalid JSON in request body", "invalid_request_error", "json_decode_error")
		return
	}

	// Get session ID from context (set by main)
	sessionID := r.Context().Value(sessionIDKey).(string)

	// Convert to Anthropic format
	anthropicReq := ConvertOpenAIRequest(&req)
	anthropicReq.Stream = true // Always stream from upstream for reliability

	log.Printf("[debug] Sending to upstream %s, model=%s", config.UpstreamURL, req.Model)

	// Proxy to upstream (always streaming)
	resp, err := proxyToUpstream(anthropicReq, apiKey, sessionID, true)
	if err != nil {
		writeError(w, http.StatusBadGateway, fmt.Sprintf("Upstream request failed: %v", err), "upstream_error", "proxy_error")
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		log.Printf("[debug] Upstream error status %d: %s", resp.StatusCode, string(respBody))
		writeError(w, http.StatusBadGateway, fmt.Sprintf("Upstream returned error: %s", string(respBody)), "upstream_error", fmt.Sprintf("status_%d", resp.StatusCode))
		return
	}

	isStream := req.Stream != nil && *req.Stream

	if isStream {
		// Stream to client
		w.Header().Set("content-type", "text/event-stream")
		w.Header().Set("cache-control", "no-cache")
		w.Header().Set("connection", "keep-alive")

		scanner := bufio.NewScanner(resp.Body)
		buf := make([]byte, 0, 64*1024)
		scanner.Buffer(buf, 1024*1024)

		created := int64(time.Now().Unix())
		chunkID := "chatcmpl-" + randomString(8)

		for scanner.Scan() {
			line := scanner.Text()
			if !strings.HasPrefix(line, "data: ") {
				continue
			}
			dataStr := strings.TrimPrefix(line, "data: ")
			if dataStr == "[DONE]" {
				break
			}

			var event struct {
				Type string `json:"type"`
			}
			if err := json.Unmarshal([]byte(dataStr), &event); err != nil {
				continue
			}

			openAIData := processAnthropicEvent(dataStr, chunkID, req.Model, created)
			if openAIData != "" {
				fmt.Fprintf(w, "data: %s\n\n", openAIData)
				w.(http.Flusher).Flush()
			}
		}
		fmt.Fprintf(w, "data: [DONE]\n\n")
		w.(http.Flusher).Flush()
	} else {
		// Non-streaming to client: accumulate upstream stream into a single response
		scanner := bufio.NewScanner(resp.Body)
		buf := make([]byte, 0, 64*1024)
		scanner.Buffer(buf, 1024*1024)

		var textContent string
		var stopReason string
		var usage AnthropicUsage
		msgID := ""

		for scanner.Scan() {
			line := scanner.Text()
			if !strings.HasPrefix(line, "data: ") {
				continue
			}
			dataStr := strings.TrimPrefix(line, "data: ")
			if dataStr == "[DONE]" {
				break
			}

			var event struct {
				Type string `json:"type"`
			}
			if err := json.Unmarshal([]byte(dataStr), &event); err != nil {
				continue
			}

			switch event.Type {
			case "message_start":
				var msgStart struct {
					Message struct {
						Id    string         `json:"id"`
						Usage AnthropicUsage `json:"usage"`
					} `json:"message"`
				}
				if json.Unmarshal([]byte(dataStr), &msgStart) == nil {
					msgID = msgStart.Message.Id
					usage.InputTokens = msgStart.Message.Usage.InputTokens
				}

			case "content_block_delta":
				var delta struct {
					Delta struct {
						Text string `json:"text,omitempty"`
					} `json:"delta"`
				}
				if json.Unmarshal([]byte(dataStr), &delta) == nil {
					textContent += delta.Delta.Text
				}

			case "message_delta":
				var msgDelta struct {
					Delta      struct{}               `json:"delta"`
					Usage      *AnthropicUsage        `json:"usage,omitempty"`
					StopReason string                 `json:"stop_reason,omitempty"`
				}
				if json.Unmarshal([]byte(dataStr), &msgDelta) == nil {
					stopReason = msgDelta.StopReason
					if msgDelta.Usage != nil {
						usage.OutputTokens = msgDelta.Usage.OutputTokens
					}
				}
			}
		}

		openAIResp := &ChatCompletionResponse{
			ID:      msgID,
			Object:  "chat.completion",
			Created: time.Now().Unix(),
			Model:   req.Model,
			Choices: []Choice{
				{
					Index:        0,
					Message:      Message{Role: "assistant", Content: textContent},
					FinishReason: mapStopReason(stopReason),
				},
			},
			Usage: Usage{
				PromptTokens:     usage.InputTokens,
				CompletionTokens: usage.OutputTokens,
				TotalTokens:      usage.InputTokens + usage.OutputTokens,
			},
		}

		w.Header().Set("content-type", "application/json")
		json.NewEncoder(w).Encode(openAIResp)
	}
}

func processAnthropicEvent(dataStr, chunkID, model string, created int64) string {
	var event struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal([]byte(dataStr), &event); err != nil {
		return ""
	}

	switch event.Type {
	case "content_block_delta":
		var delta struct {
			Index int `json:"index"`
			Delta struct {
				Text           string `json:"text,omitempty"`
				InputJSONDelta string `json:"input_json_delta,omitempty"`
			} `json:"delta"`
		}
		json.Unmarshal([]byte(dataStr), &delta)

		if delta.Delta.Text != "" {
			chunk := StreamChunk{
				ID:      chunkID,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []StreamChoice{
					{
						Index: delta.Index,
						Delta: Delta{
							Content: delta.Delta.Text,
						},
					},
				},
			}
			data, _ := json.Marshal(chunk)
			return string(data)
		}
		if delta.Delta.InputJSONDelta != "" {
			chunk := StreamChunk{
				ID:      chunkID,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []StreamChoice{
					{
						Index: delta.Index,
						Delta: Delta{
							ToolCalls: []StreamToolCall{
								{
									Index: delta.Index,
									Function: StreamFunction{
										Arguments: delta.Delta.InputJSONDelta,
									},
								},
							},
						},
					},
				},
			}
			data, _ := json.Marshal(chunk)
			return string(data)
		}

	case "message_delta":
		var msgDelta struct {
			StopReason string `json:"stop_reason,omitempty"`
		}
		json.Unmarshal([]byte(dataStr), &msgDelta)
		chunk := StreamChunk{
			ID:      chunkID,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []StreamChoice{
				{
					Index:        0,
					Delta:        Delta{},
					FinishReason: mapStopReason(msgDelta.StopReason),
				},
			},
		}
		data, _ := json.Marshal(chunk)
		return string(data)
	}

	return ""
}

func proxyToUpstream(req *AnthropicRequest, apiKey, sessionID string, stream bool) (*http.Response, error) {
	bodyBytes, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	upstreamURL := config.UpstreamURL + "/v1/messages"
	httpReq, err := http.NewRequest(http.MethodPost, upstreamURL, bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	headers := ClaudeCodeHeaders(apiKey, sessionID)
	for k, v := range headers {
		httpReq.Header.Set(k, v)
	}

	if stream {
		httpReq.Header.Set("anthropic-sse-beta", "output-2025-05-14")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
	defer cancel()

	httpReq = httpReq.WithContext(ctx)

	client := &http.Client{}
	return client.Do(httpReq)
}

func writeError(w http.ResponseWriter, code int, message, errType, errCode string) {
	w.Header().Set("content-type", "application/json")
	w.WriteHeader(code)
	resp := map[string]interface{}{
		"error": map[string]string{
			"message": message,
			"type":    errType,
			"code":    errCode,
		},
	}
	json.NewEncoder(w).Encode(resp)
}

func randomString(n int) string {
	const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, n)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := range b {
		b[i] = letters[r.Intn(len(letters))]
	}
	return string(b)
}