free-ide-proxy/internal/proxy/proxy.go

package proxy

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

// ---------- OpenAI request/response types ----------

// OpenAIChatRequest is the incoming OpenAI-format chat completion request.
type OpenAIChatRequest struct {
	Model       string          `json:"model"`
	Messages    []OpenAIMessage `json:"messages"`
	Stream      bool            `json:"stream,omitempty"`
	MaxTokens   int             `json:"max_tokens,omitempty"`
	Temperature float64         `json:"temperature,omitempty"`
	Tools       []OpenAITool    `json:"tools,omitempty"`
}

// OpenAIMessage represents a chat message.
type OpenAIMessage struct {
	Role        string           `json:"role"`
	Content     any              `json:"content"`
	Name        string           `json:"name,omitempty"`
	ToolCallID  string           `json:"tool_call_id,omitempty"`
	ToolCalls   []OpenAIToolCall `json:"tool_calls,omitempty"`
}

// OpenAIToolCall represents a tool call from the assistant.
type OpenAIToolCall struct {
	ID       string           `json:"id"`
	Type     string           `json:"type"`
	Function OpenAIToolFunc   `json:"function"`
}

// OpenAIToolFunc is the function part of a tool call.
type OpenAIToolFunc struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// OpenAITool defines a function tool.
type OpenAITool struct {
	Type     string         `json:"type"`
	Function OpenAIFunction `json:"function"`
}

// OpenAIFunction is the function definition within a tool.
type OpenAIFunction struct {
	Name        string         `json:"name"`
	Description string         `json:"description,omitempty"`
	Parameters  map[string]any `json:"parameters,omitempty"`
}

// OpenAIChatResponse is the non-streaming response.
type OpenAIChatResponse struct {
	ID      string             `json:"id"`
	Object  string             `json:"object"`
	Created int64              `json:"created"`
	Model   string             `json:"model"`
	Choices []OpenAIChoice     `json:"choices"`
	Usage   *OpenAIUsage       `json:"usage,omitempty"`
}

// OpenAIChoice is a single completion choice.
type OpenAIChoice struct {
	Index        int            `json:"index"`
	Message      OpenAIMessage  `json:"message,omitempty"`
	Delta        *OpenAIDelta   `json:"delta,omitempty"`
	FinishReason string         `json:"finish_reason,omitempty"`
}

// OpenAIDelta is a streaming delta.
type OpenAIDelta struct {
	Role      string           `json:"role,omitempty"`
	Content   string           `json:"content,omitempty"`
	ToolCalls []OpenAIToolCall `json:"tool_calls,omitempty"`
}

// OpenAIUsage holds token usage info.
type OpenAIUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// OpenAIError is a standard error response.
type OpenAIError struct {
	Error OpenAIErrorDetail `json:"error"`
}

// OpenAIErrorDetail holds error info.
type OpenAIErrorDetail struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    string `json:"code,omitempty"`
}

// ---------- Proxy forwarding logic ----------

// ForwardChatCompletion routes the OpenAI request to the matching backend and
// streams or collects the response.
func (p *Proxy) ForwardChatCompletion(w http.ResponseWriter, r *http.Request, body []byte) error {
	var req OpenAIChatRequest
	if err := json.Unmarshal(body, &req); err != nil {
		return fmt.Errorf("invalid request body: %w", err)
	}

	// Route to the backend that owns this model (default backend = pass-through).
	backend, resolved, known := p.resolveModel(req.Model)
	if !known {
		fmt.Printf("[proxy] unknown model '%s', forwarding as-is via %s\n", req.Model, backend.Name())
	} else {
		fmt.Printf("[proxy] model '%s' -> '%s' via %s\n", req.Model, resolved, backend.Name())
	}

	// Update the body with resolved model
	if resolved != req.Model {
		var reqMap map[string]any
		json.Unmarshal(body, &reqMap)
		reqMap["model"] = resolved
		body, _ = json.Marshal(reqMap)
		req.Model = resolved
	}

	// Generate request/session IDs for upstream
	requestID := RequestID()

	// Session ID keyed by the client's proxy API key (or default)
	userKey := "default"
	if p.apiKey != "" {
		userKey = r.Header.Get("Authorization")
		if userKey == "" {
			userKey = r.Header.Get("x-api-key")
		}
	}
	sessionID := p.SessionID(userKey)

	// Build upstream request to the chosen backend
	upstreamReq, err := http.NewRequestWithContext(r.Context(), "POST", backend.ChatURL(), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("failed to create upstream request: %w", err)
	}

	for k, v := range backend.Headers(requestID, sessionID) {
		upstreamReq.Header.Set(k, v)
	}

	// Execute
	client := &http.Client{Timeout: 10 * time.Minute}
	resp, err := client.Do(upstreamReq)
	if err != nil {
		return fmt.Errorf("upstream request failed: %w", err)
	}
	defer resp.Body.Close()

	// Handle non-200 responses
	if resp.StatusCode != 200 {
		fmt.Printf("[proxy] upstream status=%d for model='%s' via %s\n", resp.StatusCode, resolved, backend.Name())
		if resp.StatusCode == http.StatusTooManyRequests {
			p.markRateLimited(resolved)
		}
		bodyBytes, _ := io.ReadAll(resp.Body)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(resp.StatusCode)

		// Try to parse upstream error, wrap as standard OpenAI error
		var upstreamErr map[string]any
		if json.Unmarshal(bodyBytes, &upstreamErr) == nil {
			errMsg := "upstream error"
			if msg, ok := upstreamErr["message"]; ok {
				errMsg = fmt.Sprintf("%v", msg)
			}
			json.NewEncoder(w).Encode(OpenAIError{
				Error: OpenAIErrorDetail{
					Message: errMsg,
					Type:    "upstream_error",
				},
			})
			return nil
		}

		w.Write(bodyBytes)
		return nil
	}

	// Handle streaming vs non-streaming
	if req.Stream {
		return p.streamResponse(w, resp)
	}
	return p.nonStreamResponse(w, resp, req.Model)
}

// streamResponse forwards SSE chunks directly from upstream to client.
func (p *Proxy) streamResponse(w http.ResponseWriter, resp *http.Response) error {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(200)

	flusher, ok := w.(http.Flusher)
	if !ok {
		return fmt.Errorf("streaming not supported")
	}

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

	for scanner.Scan() {
		line := scanner.Text()
		// Write the line. If the client disconnects, stop silently.
		if _, err := fmt.Fprint(w, line+"\n"); err != nil {
			return nil
		}
		flusher.Flush()
	}

	// Scanner errors (e.g., connection reset) are not real errors
	// after streaming has started — the client likely disconnected.
	return nil
}

// nonStreamResponse collects the full upstream response and returns it as a single JSON.
func (p *Proxy) nonStreamResponse(w http.ResponseWriter, resp *http.Response, model string) error {
	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read upstream response: %w", err)
	}

	// If the upstream already returned a single JSON (non-streaming), pass it through
	contentType := resp.Header.Get("Content-Type")
	if !strings.Contains(contentType, "text/event-stream") {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		w.Write(bodyBytes)
		return nil
	}

	// Otherwise accumulate SSE chunks into a single response
	var response OpenAIChatResponse
	response.ID = "chatcmpl-" + randomHex(12)
	response.Object = "chat.completion"
	response.Created = time.Now().Unix()
	response.Model = model
	response.Choices = []OpenAIChoice{{
		Index: 0,
		Message: OpenAIMessage{
			Role: "assistant",
		},
	}}

	var contentBuf strings.Builder
	scanner := bufio.NewScanner(bytes.NewReader(bodyBytes))
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			continue
		}

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

		// Extract content from delta
		if choices, ok := chunk["choices"].([]any); ok && len(choices) > 0 {
			choice := choices[0].(map[string]any)
			if delta, ok := choice["delta"].(map[string]any); ok {
				if content, ok := delta["content"].(string); ok {
					contentBuf.WriteString(content)
				}
			}
			if fr, ok := choice["finish_reason"].(string); ok && fr != "" {
				response.Choices[0].FinishReason = fr
			}
		}
	}

	response.Choices[0].Message.Content = contentBuf.String()
	if response.Choices[0].FinishReason == "" {
		response.Choices[0].FinishReason = "stop"
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	return json.NewEncoder(w).Encode(response)
}

// WriteSSE writes a JSON object as an SSE data line.
func WriteSSE(w io.Writer, v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(w, "data: %s\n\n", data)
	return err
}