vecna/pkg/forwardproxy/forwardproxy.go

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