feat(server): add passthrough proxy for OpenAI-compatible endpoints

* implement proxy handler for various OpenAI API routes * add error handling for request body and response streaming * introduce new error response format for API compatibility * add tests for recover middleware to handle panics gracefully
2026-08-07 10:47:38 +00:00 · 2026-08-01 20:15:59 +02:00
parent 52902a8383
commit 1953a4f4f9
11 changed files with 1087 additions and 41 deletions
@@ -0,0 +1,315 @@
+// Package forwardproxy forwards raw HTTP requests (chat completions, models list,
+// audio, images, etc.) to a named target's backing endpoints, using the same
+// priority/cooldown failover algorithm as embedclient.TargetRouter but operating
+// on raw bytes instead of typed embedding requests — no vecna-specific processing
+// is applied to the body.
+package forwardproxy
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"sync"
+	"time"
+)
+
+// Endpoint is a single backing URL for a target.
+type Endpoint struct {
+	URL      string
+	APIKey   string
+	Priority int
+}
+
+// RouterConfig holds tuning parameters for a Router.
+type RouterConfig struct {
+	TargetName       string
+	TimeoutSecs      int
+	CooldownSecs     int
+	PriorityDecay    int
+	PriorityRecovery int
+}
+
+// Result is the outcome of a single forwarded (non-streaming) request.
+type Result struct {
+	Status int
+	Header http.Header
+	Body   []byte
+	Target string // endpoint URL that served the request
+}
+
+type slot struct {
+	url             string
+	apiKey          string
+	initialPriority int
+
+	mu           sync.Mutex
+	priority     int
+	inflight     int
+	successCount int
+	lastFail     time.Time
+}
+
+// Router forwards raw HTTP requests to the best available endpoint for a target,
+// retrying against other endpoints on failure.
+type Router struct {
+	slots      []*slot
+	cfg        RouterConfig
+	httpClient *http.Client
+}
+
+// NewRouter builds a Router from a target's endpoints.
+func NewRouter(endpoints []Endpoint, cfg RouterConfig, httpClient *http.Client) (*Router, error) {
+	if len(endpoints) == 0 {
+		return nil, fmt.Errorf("forwardproxy.NewRouter: at least one endpoint required")
+	}
+	if httpClient == nil {
+		httpClient = http.DefaultClient
+	}
+	slots := make([]*slot, len(endpoints))
+	for i, ep := range endpoints {
+		slots[i] = &slot{url: ep.URL, apiKey: ep.APIKey, initialPriority: ep.Priority, priority: ep.Priority}
+	}
+	return &Router{slots: slots, cfg: cfg, httpClient: httpClient}, nil
+}
+
+// Do forwards method+path+body+headers to the target's base URL + path, retrying
+// against other endpoints on failure. Buffers the full response body, so it is
+// not suitable for streaming (SSE) responses — use DoStream for those.
+func (r *Router) Do(ctx context.Context, method, path string, body []byte, headers http.Header) (Result, error) {
+	tried := make(map[*slot]bool, len(r.slots))
+	var lastErr error
+
+	for range r.slots {
+		if ctx.Err() != nil {
+			return Result{}, ctx.Err()
+		}
+		s := r.pickExcluding(tried)
+		if s == nil {
+			break
+		}
+		tried[s] = true
+
+		timeout := time.Duration(r.cfg.TimeoutSecs) * time.Second
+		reqCtx, cancel := context.WithTimeout(ctx, timeout)
+		res, err := r.doOnce(reqCtx, s, method, path, body, headers)
+		cancel()
+
+		if err != nil {
+			r.onFailure(s)
+			lastErr = fmt.Errorf("forward [%s]: %w", s.url, err)
+			continue
+		}
+		r.onSuccess(s)
+		return res, nil
+	}
+
+	if lastErr != nil {
+		return Result{}, lastErr
+	}
+	return Result{}, fmt.Errorf("forward: no endpoints available")
+}
+
+// DoStream forwards a single request to the best available endpoint without
+// retry, since bytes may already be flushed to the client before a failure is
+// known. The caller must Close the returned body; doing so records success or
+// failure against the endpoint's priority score.
+func (r *Router) DoStream(ctx context.Context, method, path string, body []byte, headers http.Header) (status int, respHeader http.Header, respBody io.ReadCloser, err error) {
+	s := r.pickExcluding(nil)
+	if s == nil {
+		return 0, nil, nil, fmt.Errorf("forward: no endpoints available")
+	}
+
+	httpReq, err := r.buildRequest(ctx, s, method, path, body, headers)
+	if err != nil {
+		return 0, nil, nil, err
+	}
+
+	s.mu.Lock()
+	s.inflight++
+	s.mu.Unlock()
+
+	resp, err := r.httpClient.Do(httpReq) //nolint:bodyclose // closed by the caller via trackedBody.Close
+	if err != nil {
+		s.mu.Lock()
+		s.inflight--
+		s.mu.Unlock()
+		r.onFailure(s)
+		return 0, nil, nil, fmt.Errorf("forward [%s]: %w", s.url, err)
+	}
+
+	return resp.StatusCode, resp.Header, &trackedBody{ReadCloser: resp.Body, onClose: func() {
+		s.mu.Lock()
+		s.inflight--
+		s.mu.Unlock()
+		r.onSuccess(s)
+	}}, nil
+}
+
+type trackedBody struct {
+	io.ReadCloser
+	onClose func()
+	once    sync.Once
+}
+
+func (t *trackedBody) Close() error {
+	err := t.ReadCloser.Close()
+	t.once.Do(t.onClose)
+	return err
+}
+
+func (r *Router) doOnce(ctx context.Context, s *slot, method, path string, body []byte, headers http.Header) (Result, error) {
+	s.mu.Lock()
+	s.inflight++
+	s.mu.Unlock()
+	defer func() {
+		s.mu.Lock()
+		s.inflight--
+		s.mu.Unlock()
+	}()
+
+	httpReq, err := r.buildRequest(ctx, s, method, path, body, headers)
+	if err != nil {
+		return Result{}, err
+	}
+	resp, err := r.httpClient.Do(httpReq)
+	if err != nil {
+		return Result{}, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return Result{}, fmt.Errorf("read response: %w", err)
+	}
+	if resp.StatusCode >= 500 {
+		return Result{}, fmt.Errorf("upstream status %d", resp.StatusCode)
+	}
+	return Result{Status: resp.StatusCode, Header: resp.Header, Body: respBody, Target: s.url}, nil
+}
+
+func (r *Router) buildRequest(ctx context.Context, s *slot, method, path string, body []byte, headers http.Header) (*http.Request, error) {
+	url := s.url + path
+	var bodyReader io.Reader
+	if body != nil {
+		bodyReader = bytes.NewReader(body)
+	}
+	httpReq, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
+	if err != nil {
+		return nil, fmt.Errorf("build request: %w", err)
+	}
+	copyHeaders(httpReq.Header, headers, skipOutbound)
+	if s.apiKey != "" {
+		httpReq.Header.Set("Authorization", "Bearer "+s.apiKey)
+	}
+	return httpReq, nil
+}
+
+func (r *Router) pickExcluding(excluded map[*slot]bool) *slot {
+	cooldown := time.Duration(r.cfg.CooldownSecs) * time.Second
+	now := time.Now()
+
+	var best *slot
+	bestScore := -1 << 30
+
+	for _, s := range r.slots {
+		if excluded[s] {
+			continue
+		}
+		s.mu.Lock()
+		inCooldown := !s.lastFail.IsZero() && now.Sub(s.lastFail) < cooldown
+		score := s.priority - s.inflight
+		s.mu.Unlock()
+
+		if inCooldown {
+			continue
+		}
+		if best == nil || score > bestScore {
+			best = s
+			bestScore = score
+		}
+	}
+
+	// All non-excluded slots are in cooldown — fall back to the one with the oldest failure.
+	if best == nil {
+		var oldest time.Time
+		for _, s := range r.slots {
+			if excluded[s] {
+				continue
+			}
+			s.mu.Lock()
+			lf := s.lastFail
+			s.mu.Unlock()
+			if best == nil || lf.Before(oldest) {
+				best = s
+				oldest = lf
+			}
+		}
+	}
+
+	return best
+}
+
+func (r *Router) onSuccess(s *slot) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.successCount++
+	if r.cfg.PriorityRecovery > 0 && s.successCount%r.cfg.PriorityRecovery == 0 {
+		if s.priority < s.initialPriority {
+			s.priority++
+		}
+	}
+}
+
+func (r *Router) onFailure(s *slot) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.lastFail = time.Now()
+	s.priority -= r.cfg.PriorityDecay
+	if s.priority < 1 {
+		s.priority = 1
+	}
+}
+
+// hopByHop headers are never forwarded in either direction.
+var hopByHop = map[string]bool{
+	"Connection":        true,
+	"Proxy-Connection":  true,
+	"Keep-Alive":        true,
+	"Transfer-Encoding": true,
+	"Te":                true,
+	"Trailer":           true,
+	"Upgrade":           true,
+	"Content-Length":    true, // net/http (re)computes this from the body it's given
+}
+
+// skipOutbound headers are additionally stripped from client->upstream requests:
+// these authenticate the client to vecna, not vecna to the upstream, and are
+// replaced with the endpoint's own api key (as an Authorization Bearer header), if any.
+var skipOutbound = map[string]bool{
+	"Host":          true,
+	"Authorization": true,
+	"X-Api-Key":     true,
+	"Api-Key":       true,
+}
+
+// SkipResponseHeader reports whether a header should not be copied from the
+// upstream response back to the original client.
+func SkipResponseHeader(name string) bool {
+	return hopByHop[http.CanonicalHeaderKey(name)]
+}
+
+func copyHeaders(dst, src http.Header, extraSkip map[string]bool) {
+	for k, vv := range src {
+		ck := http.CanonicalHeaderKey(k)
+		if hopByHop[ck] || extraSkip[ck] {
+			continue
+		}
+		for _, v := range vv {
+			dst.Add(k, v)
+		}
+	}
+}