feat: multi-provider support — Ollama Cloud + OpenCode Go
- Account.Provider field (ollama-cloud | opencode-go), backward compatible - Model-based routing: common models served by combined pool, unique models routed to their provider only - /go/v1/* path forces OpenCode Go provider (prefix stripped upstream) - Merged /v1/models endpoint returns union of both catalogs (44 models) - Failover: 429/402 → cooldown + failover; 5xx → retry without cooldown - CLI: accounts add --provider flag, list shows provider column - Body buffering: request body buffered (8 MiB cap) for failover replay - opencode integration: unified provider 'oc' with all merged models - 64 tests pass (unit + integration) - Verified: glm-5 → ollama, mimo-v2.5 → go, gpt-oss:20b → ollama
This commit is contained in:
Atte149
parent
98bbc96bf5
commit
4fe15324c8
11 changed files with 973 additions and 125 deletions
|
|
@ -1,6 +1,7 @@
|
|||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
|
|
@ -8,41 +9,58 @@ import (
|
|||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Atte149/ollama-proxy/internal/config"
|
||||
)
|
||||
|
||||
// Handler is the http.Handler that proxies incoming requests to Ollama Cloud
|
||||
// through the balancer, applying per-account Authorization and failover logic.
|
||||
// Handler is the http.Handler that proxies incoming requests to upstream
|
||||
// providers (Ollama Cloud and OpenCode Go) through the balancer, applying
|
||||
// per-account Authorization and failover logic.
|
||||
type Handler struct {
|
||||
Balancer *Balancer
|
||||
BaseURL string // upstream root, e.g. "https://ollama.com"
|
||||
Client *http.Client
|
||||
Log *slog.Logger
|
||||
Client *http.Client
|
||||
Retries int // max attempts per request (= number of accounts to try)
|
||||
// DefaultBaseURL, when non-empty, overrides the provider-default upstream
|
||||
// root for accounts that have no explicit BaseURL. This is primarily used
|
||||
// by the --base-url flag for local testing; in production accounts rely on
|
||||
// the provider default.
|
||||
DefaultBaseURL string
|
||||
}
|
||||
|
||||
// NewHandler builds a Handler with sensible HTTP client defaults (no timeout
|
||||
// on the overall request — streaming responses can be long; per-read timeout
|
||||
// is governed by the caller's context).
|
||||
// NewHandler builds a Handler with sensible HTTP client defaults.
|
||||
func NewHandler(b *Balancer, baseURL string, retries int, log *slog.Logger) *Handler {
|
||||
return &Handler{
|
||||
Balancer: b,
|
||||
BaseURL: strings.TrimRight(baseURL, "/"),
|
||||
Log: log,
|
||||
Client: &http.Client{
|
||||
// No overall timeout: streaming chat may take minutes. The caller's
|
||||
// request context (cancelled when the client disconnects) still
|
||||
// propagates through to upstream via NewWithContext.
|
||||
Timeout: 0,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
return http.ErrUseLastResponse // don't follow redirects automatically
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
},
|
||||
Log: log,
|
||||
Retries: retries,
|
||||
Retries: retries,
|
||||
DefaultBaseURL: baseURL,
|
||||
}
|
||||
}
|
||||
|
||||
// allowedPaths is the allowlist of upstream paths the proxy will forward.
|
||||
// Anything else returns 404 — we never proxy arbitrary paths.
|
||||
// upstreamBaseURL resolves the canonical upstream root for an account:
|
||||
// 1. The account's own BaseURL override (highest priority)
|
||||
// 2. The handler's DefaultBaseURL — but ONLY for ollama-cloud accounts,
|
||||
// since --base-url is an Ollama-specific flag used for local testing
|
||||
// 3. The provider default
|
||||
func (h *Handler) upstreamBaseURL(acct config.Account) string {
|
||||
if acct.BaseURL != "" {
|
||||
return strings.TrimRight(acct.BaseURL, "/")
|
||||
}
|
||||
if h.DefaultBaseURL != "" && acct.Provider == config.ProviderOllamaCloud {
|
||||
return strings.TrimRight(h.DefaultBaseURL, "/")
|
||||
}
|
||||
return strings.TrimRight(config.DefaultProviderBaseURL(acct.Provider), "/")
|
||||
}
|
||||
|
||||
// allowedPaths is the allowlist of paths the proxy will forward. The same set
|
||||
// is used for both providers; the /go/ prefix routes to OpenCode Go.
|
||||
var allowedPaths = map[string]bool{
|
||||
"/api/chat": true,
|
||||
"/api/generate": true,
|
||||
|
|
@ -58,13 +76,11 @@ var allowedPaths = map[string]bool{
|
|||
"/v1/files": true,
|
||||
}
|
||||
|
||||
// isAllowed reports whether a path (possibly with a trailing slash or query)
|
||||
// matches one of the allowed upstream endpoints.
|
||||
// isAllowed reports whether a path matches an allowed upstream endpoint.
|
||||
func isAllowed(path string) bool {
|
||||
if allowedPaths[path] {
|
||||
return true
|
||||
}
|
||||
// allow sub-paths like /v1/files/<id> under an allowed prefix.
|
||||
for prefix := range allowedPaths {
|
||||
if strings.HasPrefix(path, prefix+"/") {
|
||||
return true
|
||||
|
|
@ -73,9 +89,72 @@ func isAllowed(path string) bool {
|
|||
return false
|
||||
}
|
||||
|
||||
// ServeHTTP proxies a single client request to Ollama Cloud with failover.
|
||||
// routePath normalises the incoming path and returns (upstreamPath, forcedProvider).
|
||||
// Paths under /go/ force the OpenCode Go provider and have the /go prefix
|
||||
// stripped before forwarding. Otherwise the provider is chosen by the request
|
||||
// body model (or defaults to ollama-cloud when no model is present).
|
||||
func routePath(path string) (upstreamPath string, forced config.ProviderType, isGo bool) {
|
||||
if strings.HasPrefix(path, "/go/") || path == "/go" {
|
||||
return "/" + strings.TrimPrefix(path, "/go/"), config.ProviderOpenCodeGo, true
|
||||
}
|
||||
return path, "", false
|
||||
}
|
||||
|
||||
// modelFromBody extracts the "model" field from a buffered request body.
|
||||
// Returns "" when the body is nil, empty, or does not contain a parseable
|
||||
// JSON object with a model field. Does not consume the buffer (resets position).
|
||||
func modelFromBody(body *bufferedBody) string {
|
||||
if body == nil || len(body.data) == 0 {
|
||||
return ""
|
||||
}
|
||||
var probe struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
if err := json.Unmarshal(body.data, &probe); err != nil {
|
||||
return ""
|
||||
}
|
||||
return probe.Model
|
||||
}
|
||||
|
||||
// bufferedBody wraps a []byte so it can be replayed across failover attempts.
|
||||
type bufferedBody struct {
|
||||
data []byte
|
||||
pos int
|
||||
}
|
||||
|
||||
func (b *bufferedBody) Read(p []byte) (int, error) {
|
||||
if b.pos >= len(b.data) {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n := copy(p, b.data[b.pos:])
|
||||
b.pos += n
|
||||
return n, nil
|
||||
}
|
||||
func (b *bufferedBody) Close() error { return nil }
|
||||
func (b *bufferedBody) reset() { b.pos = 0 }
|
||||
|
||||
// bufferRequestBody reads the request body (up to a cap) into a buffer so it
|
||||
// can be replayed across failover attempts. Bodies larger than the cap are
|
||||
// not buffered and the request is not retried. Returns nil for streaming-safe
|
||||
// GET/DELETE requests.
|
||||
func bufferRequestBody(r *http.Request) *bufferedBody {
|
||||
if r.Body == nil || r.Method == http.MethodGet || r.Method == http.MethodDelete ||
|
||||
r.Method == http.MethodHead {
|
||||
return nil
|
||||
}
|
||||
const cap = 8 << 20 // 8 MiB
|
||||
data, err := io.ReadAll(io.LimitReader(r.Body, cap+1))
|
||||
if err != nil || len(data) > cap {
|
||||
return nil // too large or unreadable; disable retries by returning nil
|
||||
}
|
||||
return &bufferedBody{data: data}
|
||||
}
|
||||
|
||||
// ServeHTTP proxies a single client request to the upstream with failover.
|
||||
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if !isAllowed(r.URL.Path) {
|
||||
upstreamPath, forced, isGo := routePath(r.URL.Path)
|
||||
_ = isGo
|
||||
if !isAllowed(upstreamPath) {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
|
|
@ -84,32 +163,101 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
// Buffer the request body up front so we can both extract the model field
|
||||
// and replay the body across failover attempts. This must happen BEFORE
|
||||
// requestModel reads from r.Body.
|
||||
bodyBuf := bufferRequestBody(r)
|
||||
|
||||
// Determine eligible providers for this request.
|
||||
var eligible []config.ProviderType
|
||||
switch {
|
||||
case forced != "":
|
||||
// /go/* path forces the OpenCode Go provider.
|
||||
eligible = []config.ProviderType{forced}
|
||||
case upstreamPath == "/v1/models" || upstreamPath == "/api/tags":
|
||||
// Listing endpoints: handled specially below (merged response).
|
||||
h.handleModelsEndpoint(w, r, upstreamPath)
|
||||
return
|
||||
case upstreamPath == "/api/version":
|
||||
// Version is Ollama-specific; route to Ollama Cloud.
|
||||
eligible = []config.ProviderType{config.ProviderOllamaCloud}
|
||||
default:
|
||||
// Chat/generate/etc.: route by the model in the body.
|
||||
model := modelFromBody(bodyBuf)
|
||||
if model == "" {
|
||||
eligible = nil // any provider
|
||||
} else {
|
||||
eligible = ProvidersForModel(model)
|
||||
}
|
||||
}
|
||||
|
||||
maxAttempts := h.Retries
|
||||
if maxAttempts < 1 {
|
||||
maxAttempts = 1
|
||||
}
|
||||
if maxAttempts > h.Balancer.Len() {
|
||||
maxAttempts = h.Balancer.Len()
|
||||
if n := h.Balancer.Len(); maxAttempts > n {
|
||||
maxAttempts = n
|
||||
}
|
||||
|
||||
h.proxyWithFailover(w, r, maxAttempts)
|
||||
h.proxyWithFailover(w, r, upstreamPath, eligible, bodyBuf, maxAttempts)
|
||||
}
|
||||
|
||||
// proxyWithFailover tries up to maxAttempts accounts. The first attempt that
|
||||
// begins streaming (returns headers + a 2xx) commits: we copy the rest of the
|
||||
// response to the client verbatim and stop retrying. Pre-stream 429/5xx move
|
||||
// on to the next account.
|
||||
func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxAttempts int) {
|
||||
// handleModelsEndpoint returns a merged list of all known models for
|
||||
// /v1/models (OpenAI format) or /api/tags (Ollama format).
|
||||
func (h *Handler) handleModelsEndpoint(w http.ResponseWriter, r *http.Request, path string) {
|
||||
now := time.Now().Unix()
|
||||
models := AllModels()
|
||||
if path == "/v1/models" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
var sb strings.Builder
|
||||
sb.WriteString(`{"object":"list","data":[`)
|
||||
for i, m := range models {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
// The id is the raw model ID (used in API requests). The name
|
||||
// includes a provider label for unique models so users can tell
|
||||
// them apart in the opencode model picker.
|
||||
name := m
|
||||
if !IsCommonModel(m) {
|
||||
name = m + " (" + ProviderLabel(ProvidersForModel(m)[0]) + ")"
|
||||
}
|
||||
fmt.Fprintf(&sb, `{"id":%q,"object":"model","created":%d,"owned_by":"ollama-proxy"}`,
|
||||
m, now)
|
||||
_ = name // name is conveyed via opencode config, not this endpoint
|
||||
}
|
||||
sb.WriteString(`]}`)
|
||||
_, _ = w.Write([]byte(sb.String()))
|
||||
return
|
||||
}
|
||||
// /api/tags — Ollama format
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
var sb strings.Builder
|
||||
sb.WriteString(`{"models":[`)
|
||||
for i, m := range models {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&sb, `{"name":%q,"model":%q,"modified_at":"","size":0,"digest":""}`, m, m)
|
||||
}
|
||||
sb.WriteString(`]}`)
|
||||
_, _ = w.Write([]byte(sb.String()))
|
||||
}
|
||||
|
||||
// proxyWithFailover tries up to maxAttempts accounts eligible for the request.
|
||||
func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request,
|
||||
upstreamPath string, eligible []config.ProviderType, body *bufferedBody,
|
||||
maxAttempts int) {
|
||||
ctx := r.Context()
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt < maxAttempts; attempt++ {
|
||||
acct, err := h.Balancer.Next()
|
||||
acct, err := h.Balancer.NextFor(eligible)
|
||||
if err != nil {
|
||||
// No accounts available.
|
||||
switch err.(type) {
|
||||
case ErrAllCooldown:
|
||||
// All accounts rate-limited upstream → propagate 429 to client.
|
||||
msg := err.Error()
|
||||
if lastErr != nil {
|
||||
msg = msg + "; last upstream error: " + lastErr.Error()
|
||||
|
|
@ -125,43 +273,58 @@ func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxA
|
|||
}
|
||||
}
|
||||
|
||||
// Build the upstream request, substituting the account's key.
|
||||
upstreamURL := h.BaseURL + r.URL.RequestURI()
|
||||
// Per-account BaseURL override wins when set.
|
||||
if acct.BaseURL != "" {
|
||||
upstreamURL = strings.TrimRight(acct.BaseURL, "/") + r.URL.RequestURI()
|
||||
// Build the upstream request.
|
||||
base := h.upstreamBaseURL(acct)
|
||||
// Build the upstream URI from the routed path + original query string.
|
||||
uri := upstreamPath
|
||||
if r.URL.RawQuery != "" {
|
||||
uri += "?" + r.URL.RawQuery
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL, r.Body)
|
||||
upstreamURL := base + uri
|
||||
|
||||
var bodyReader io.Reader
|
||||
if body != nil {
|
||||
body.reset()
|
||||
bodyReader = body
|
||||
} else if r.Body != nil {
|
||||
bodyReader = r.Body
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL, bodyReader)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("build request: %w", err)
|
||||
continue
|
||||
}
|
||||
// Copy headers, replacing Authorization with the account key.
|
||||
copyHeaders(req.Header, r.Header)
|
||||
req.Header.Set("Authorization", "Bearer "+acct.APIKey)
|
||||
req.Host = "" // let the URL determine Host
|
||||
// OpenCode Go uses OpenAI-compatible JSON; ensure content-type is set.
|
||||
if req.Header.Get("Content-Type") == "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
// Let the URL determine the Host header (do not clear it; clearing can
|
||||
// cause some CDNs to reject the request with 401/403).
|
||||
|
||||
start := time.Now()
|
||||
resp, err := h.Client.Do(req)
|
||||
latency := time.Since(start)
|
||||
if err != nil {
|
||||
// Network/timeout error: log and try next account.
|
||||
h.Log.Warn("upstream request failed",
|
||||
"account", acct.Name, "account_id", acct.ID,
|
||||
"provider", string(acct.Provider),
|
||||
"path", r.URL.Path, "latency_ms", latency.Milliseconds(), "err", err)
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
|
||||
// Pre-stream decision: 429 or 5xx → cooldown + failover.
|
||||
if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
||||
bodyPreview := readPreview(resp.Body, 256)
|
||||
// Pre-stream decision: retriable errors → cooldown + failover.
|
||||
if isRetriable(resp.StatusCode) {
|
||||
bodyPreview := readPreview(resp.Body, 512)
|
||||
resp.Body.Close()
|
||||
h.Log.Warn("upstream rejected, failover",
|
||||
"account", acct.Name, "account_id", acct.ID,
|
||||
"provider", string(acct.Provider),
|
||||
"path", r.URL.Path, "status", resp.StatusCode,
|
||||
"latency_ms", latency.Milliseconds(), "preview", bodyPreview)
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
if shouldCooldown(resp.StatusCode, bodyPreview) {
|
||||
h.Balancer.MarkCooldown(acct.ID)
|
||||
}
|
||||
lastErr = fmt.Errorf("upstream %d: %s", resp.StatusCode, bodyPreview)
|
||||
|
|
@ -172,17 +335,29 @@ func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxA
|
|||
h.copyResponse(w, resp)
|
||||
h.Log.Info("proxied",
|
||||
"account", acct.Name, "account_id", acct.ID,
|
||||
"provider", string(acct.Provider),
|
||||
"path", r.URL.Path, "status", resp.StatusCode,
|
||||
"latency_ms", latency.Milliseconds(), "stream", isStreaming(resp))
|
||||
return
|
||||
}
|
||||
|
||||
// Exhausted retries. Distinguish "everything is rate-limited" (→ 429) from
|
||||
// "mixed upstream errors" (→ 502).
|
||||
// Exhausted retries.
|
||||
if h.Balancer.Len() > 0 {
|
||||
allCooldown := true
|
||||
status := h.Balancer.Status()
|
||||
for _, a := range h.Balancer.Accounts() {
|
||||
if eligible != nil && len(eligible) > 0 {
|
||||
matched := false
|
||||
for _, p := range eligible {
|
||||
if a.Provider == p {
|
||||
matched = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if !status[a.ID].InCooldown {
|
||||
allCooldown = false
|
||||
break
|
||||
|
|
@ -204,25 +379,21 @@ func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxA
|
|||
writeError(w, http.StatusBadGateway, "all accounts failed")
|
||||
}
|
||||
|
||||
// copyResponse streams the upstream response body to the client. For chunked /
|
||||
// SSE responses we flush after every read so tokens reach the client immediately.
|
||||
// copyResponse streams the upstream response body to the client.
|
||||
func (h *Handler) copyResponse(w http.ResponseWriter, resp *http.Response) {
|
||||
defer resp.Body.Close()
|
||||
// Copy headers (except hop-by-hop ones).
|
||||
for k, vs := range resp.Header {
|
||||
for _, v := range vs {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
|
||||
flusher, _ := w.(http.Flusher)
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := resp.Body.Read(buf)
|
||||
if n > 0 {
|
||||
if _, werr := w.Write(buf[:n]); werr != nil {
|
||||
// client went away; stop copying silently
|
||||
return
|
||||
}
|
||||
if flusher != nil {
|
||||
|
|
@ -239,12 +410,12 @@ func (h *Handler) copyResponse(w http.ResponseWriter, resp *http.Response) {
|
|||
}
|
||||
|
||||
// copyHeaders duplicates src into dst, dropping hop-by-hop headers and any
|
||||
// Authorization that the client may have sent (we always set our own).
|
||||
// Authorization the client may have sent.
|
||||
func copyHeaders(dst, src http.Header) {
|
||||
hopByHop := []string{
|
||||
"Connection", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization",
|
||||
"Te", "Trailers", "Transfer-Encoding", "Upgrade",
|
||||
"Authorization", // always overwritten by the account key
|
||||
"Authorization",
|
||||
}
|
||||
for k, vs := range src {
|
||||
skip := false
|
||||
|
|
@ -263,8 +434,6 @@ func copyHeaders(dst, src http.Header) {
|
|||
}
|
||||
}
|
||||
|
||||
// readPreview reads up to n bytes from r and returns them as a string, always
|
||||
// closing the reader.
|
||||
func readPreview(r io.ReadCloser, n int) string {
|
||||
defer r.Close()
|
||||
buf := make([]byte, n)
|
||||
|
|
@ -272,21 +441,15 @@ func readPreview(r io.ReadCloser, n int) string {
|
|||
return string(buf[:m])
|
||||
}
|
||||
|
||||
// isStreaming reports whether the response is streaming (SSE or NDJSON/chunked).
|
||||
func isStreaming(resp *http.Response) bool {
|
||||
ct := resp.Header.Get("Content-Type")
|
||||
if strings.Contains(ct, "text/event-stream") {
|
||||
return true
|
||||
}
|
||||
if strings.Contains(ct, "application/x-ndjson") {
|
||||
if strings.Contains(ct, "text/event-stream") || strings.Contains(ct, "application/x-ndjson") {
|
||||
return true
|
||||
}
|
||||
te := resp.Header.Get("Transfer-Encoding")
|
||||
return strings.Contains(strings.ToLower(te), "chunked")
|
||||
}
|
||||
|
||||
// writeError emits a JSON-formatted error to the client, mirroring Ollama's
|
||||
// error shape so OpenAI-compatible clients can parse it.
|
||||
func writeError(w http.ResponseWriter, status int, msg string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue