proxy.go

package main

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

const (
	copilotChatURL  = "https://api.githubcopilot.com/chat/completions"
	editorVersion   = "vscode/1.96.2"
	copilotUA       = "GitHubCopilotChat/0.26.7"
)

// openAIRequest is used to decode the model field for usage tracking / routing.
type openAIRequest struct {
	Model  string `json:"model"`
	Stream bool   `json:"stream"`
}

// makeOpenAIHandler returns an HTTP handler that proxies OpenAI-compatible
// chat-completion requests to the GitHub Copilot API.
func makeOpenAIHandler(tm *TokenManager, tracker *UsageTracker) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}

		body, err := io.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "reading request body: "+err.Error(), http.StatusBadRequest)
			return
		}
		r.Body.Close()

		var req openAIRequest
		if err := json.Unmarshal(body, &req); err != nil {
			http.Error(w, "parsing request body: "+err.Error(), http.StatusBadRequest)
			return
		}

		token, err := tm.Token()
		if err != nil {
			logger.Printf("token error: %v", err)
			http.Error(w, "failed to obtain Copilot token", http.StatusServiceUnavailable)
			return
		}

		upstream, err := http.NewRequest(http.MethodPost, copilotChatURL, bytes.NewReader(body))
		if err != nil {
			http.Error(w, "creating upstream request: "+err.Error(), http.StatusInternalServerError)
			return
		}
		upstream.Header.Set("Authorization", "Bearer "+token)
		upstream.Header.Set("Content-Type", "application/json")
		upstream.Header.Set("Editor-Version", editorVersion)
		upstream.Header.Set("User-Agent", copilotUA)
		upstream.Header.Set("Copilot-Integration-Id", "vscode-chat")

		client := &http.Client{Timeout: 120 * time.Second}
		resp, err := client.Do(upstream)
		if err != nil {
			logger.Printf("upstream request error: %v", err)
			http.Error(w, "upstream error: "+err.Error(), http.StatusBadGateway)
			return
		}
		defer resp.Body.Close()

		// Copy upstream headers to client.
		for k, vs := range resp.Header {
			for _, v := range vs {
				w.Header().Add(k, v)
			}
		}
		w.WriteHeader(resp.StatusCode)

		if req.Stream {
			// Stream directly; track usage from SSE if possible.
			proxyStreamOpenAI(w, resp.Body, req.Model, tracker)
		} else {
			respBody, err := io.ReadAll(resp.Body)
			if err != nil {
				logger.Printf("reading upstream response: %v", err)
				return
			}
			w.Write(respBody)
			if resp.StatusCode == http.StatusOK {
				trackOpenAIUsage(respBody, req.Model, tracker)
			}
		}
	}
}

// proxyStreamOpenAI copies SSE lines from the upstream to the client and tries
// to record usage from the final data chunk.
func proxyStreamOpenAI(w http.ResponseWriter, body io.Reader, model string, tracker *UsageTracker) {
	flusher, canFlush := w.(http.Flusher)

	// We accumulate the last non-[DONE] data line to try to extract usage.
	var lastDataLine string

	buf := make([]byte, 4096)
	pending := ""
	for {
		n, err := body.Read(buf)
		if n > 0 {
			pending += string(buf[:n])
			// Process complete lines.
			for {
				idx := strings.Index(pending, "\n")
				if idx < 0 {
					break
				}
				line := pending[:idx]
				pending = pending[idx+1:]

				// Track potential usage lines.
				trimmed := strings.TrimPrefix(line, "data: ")
				if trimmed != "[DONE]" && strings.HasPrefix(line, "data: ") {
					lastDataLine = trimmed
				}

				fmt.Fprint(w, line+"\n")
			}
			if canFlush {
				flusher.Flush()
			}
		}
		if err != nil {
			break
		}
	}
	// Flush any remaining bytes.
	if pending != "" {
		fmt.Fprint(w, pending)
		if canFlush {
			flusher.Flush()
		}
	}

	// Attempt to parse usage from the last data line.
	if lastDataLine != "" {
		trackOpenAIUsage([]byte(lastDataLine), model, tracker)
	}
}

// trackOpenAIUsage decodes an OpenAI response body and records token usage.
func trackOpenAIUsage(body []byte, model string, tracker *UsageTracker) {
	var resp struct {
		Usage *struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(body, &resp); err != nil || resp.Usage == nil {
		return
	}
	if err := tracker.Record(model, resp.Usage.PromptTokens, resp.Usage.CompletionTokens); err != nil {
		logger.Printf("usage record error: %v", err)
	}
}