proxx/handler.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"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

// 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)

	// Check streaming
	isStream := req.Stream != nil && *req.Stream

	// Proxy to upstream
	resp, err := proxyToUpstream(anthropicReq, apiKey, sessionID, isStream)
	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)
		writeError(w, http.StatusBadGateway, fmt.Sprintf("Upstream returned error: %s", string(respBody)), "upstream_error", fmt.Sprintf("status_%d", resp.StatusCode))
		return
	}

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

		// Convert SSE stream on the fly
		reader := io.NopCloser(resp.Body)
		scanner := bufio.NewScanner(reader)
		// Use a larger buffer for potential large JSON lines
		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]" {
				fmt.Fprintf(w, "data: [DONE]\n\n")
				w.(http.Flusher).Flush()
				continue
			}

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

			// Process event and write OpenAI SSE
			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 response
		respBody, err := io.ReadAll(resp.Body)
		if err != nil {
			writeError(w, http.StatusBadGateway, "Failed to read upstream response", "upstream_error", "body_read_error")
			return
		}

		var anthropicResp AnthropicResponse
		if err := json.Unmarshal(respBody, &anthropicResp); err != nil {
			writeError(w, http.StatusBadGateway, "Failed to parse upstream response", "upstream_error", "json_decode_error")
			return
		}

		openAIResp := ConvertAnthropicResponse(&anthropicResp, req.Model)

		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 "message_start":
		return ""

	case "content_block_start":
		return ""

	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(), 120*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)
}