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
|
|
@ -37,16 +37,21 @@ func RunAccounts(args []string) int {
|
|||
|
||||
func printAccountsUsage() {
|
||||
fmt.Print(`Usage:
|
||||
ollama-proxy accounts add [API_KEY] [--name <alias>]
|
||||
ollama-proxy accounts add [API_KEY] [--name <alias>] [--provider ollama-cloud|opencode-go]
|
||||
ollama-proxy accounts list
|
||||
ollama-proxy accounts remove <id|name>
|
||||
ollama-proxy accounts set-base-url <url>
|
||||
|
||||
Providers:
|
||||
ollama-cloud (default) Ollama Cloud at https://ollama.com — native + OpenAI API
|
||||
opencode-go OpenCode Go at https://opencode.ai/zen/go/v1 — OpenAI API only
|
||||
`)
|
||||
}
|
||||
|
||||
// accountsAdd implements `accounts add [API_KEY] [--name <alias>]`.
|
||||
// accountsAdd implements `accounts add [API_KEY] [--name <alias>] [--provider P]`.
|
||||
func accountsAdd(args []string) int {
|
||||
var apiKey, name string
|
||||
provider := config.ProviderOllamaCloud
|
||||
for i := 0; i < len(args); i++ {
|
||||
switch args[i] {
|
||||
case "--name":
|
||||
|
|
@ -56,8 +61,15 @@ func accountsAdd(args []string) int {
|
|||
}
|
||||
name = args[i+1]
|
||||
i++
|
||||
case "--provider":
|
||||
if i+1 >= len(args) {
|
||||
fmt.Fprintln(os.Stderr, "--provider requires an argument")
|
||||
return 2
|
||||
}
|
||||
provider = config.ProviderType(args[i+1])
|
||||
i++
|
||||
case "-h", "--help":
|
||||
fmt.Println("usage: accounts add [API_KEY] [--name <alias>]")
|
||||
fmt.Println("usage: accounts add [API_KEY] [--name <alias>] [--provider ollama-cloud|opencode-go]")
|
||||
return 0
|
||||
default:
|
||||
if apiKey == "" {
|
||||
|
|
@ -69,6 +81,11 @@ func accountsAdd(args []string) int {
|
|||
}
|
||||
}
|
||||
|
||||
if !config.ValidProvider(provider) {
|
||||
fmt.Fprintf(os.Stderr, "unknown provider %q (use ollama-cloud or opencode-go)\n", provider)
|
||||
return 2
|
||||
}
|
||||
|
||||
if name == "" {
|
||||
name = promptString("Account name (alias)", "acct"+config.NewID()[:4])
|
||||
}
|
||||
|
|
@ -78,7 +95,14 @@ func accountsAdd(args []string) int {
|
|||
}
|
||||
|
||||
if apiKey == "" {
|
||||
fmt.Print("Enter Ollama Cloud API key (input hidden): ")
|
||||
label := "API key"
|
||||
switch provider {
|
||||
case config.ProviderOpenCodeGo:
|
||||
label = "OpenCode Go API key"
|
||||
case config.ProviderOllamaCloud:
|
||||
label = "Ollama Cloud API key"
|
||||
}
|
||||
fmt.Printf("Enter %s (input hidden): ", label)
|
||||
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||
fmt.Println()
|
||||
if err != nil {
|
||||
|
|
@ -106,17 +130,19 @@ func accountsAdd(args []string) int {
|
|||
return 2
|
||||
}
|
||||
acct := config.Account{
|
||||
ID: config.NewID(),
|
||||
Name: name,
|
||||
APIKey: apiKey,
|
||||
Created: nowFn(),
|
||||
ID: config.NewID(),
|
||||
Name: name,
|
||||
Provider: provider,
|
||||
APIKey: apiKey,
|
||||
Created: nowFn(),
|
||||
}
|
||||
af.Accounts = append(af.Accounts, acct)
|
||||
if err := af.Save(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "save accounts:", err)
|
||||
return 1
|
||||
}
|
||||
fmt.Printf("added account %s (name=%s key=%s)\n", acct.ID[:8], acct.Name, maskKey(acct.APIKey))
|
||||
fmt.Printf("added account %s (name=%s provider=%s key=%s)\n",
|
||||
acct.ID[:8], acct.Name, acct.Provider, maskKey(acct.APIKey))
|
||||
return 0
|
||||
}
|
||||
|
||||
|
|
@ -133,15 +159,21 @@ func accountsList(args []string) int {
|
|||
}
|
||||
w := bufio.NewWriter(os.Stdout)
|
||||
defer w.Flush()
|
||||
fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", "ID", "NAME", "KEY", "CREATED")
|
||||
fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", strings.Repeat("-", 8), strings.Repeat("-", 14), strings.Repeat("-", 16), strings.Repeat("-", 20))
|
||||
fmt.Fprintf(w, "%-10s %-14s %-16s %-18s %-22s\n", "ID", "PROVIDER", "NAME", "KEY", "CREATED")
|
||||
fmt.Fprintf(w, "%-10s %-14s %-16s %-18s %-22s\n",
|
||||
strings.Repeat("-", 8), strings.Repeat("-", 12), strings.Repeat("-", 14),
|
||||
strings.Repeat("-", 16), strings.Repeat("-", 20))
|
||||
for _, a := range af.Accounts {
|
||||
id := a.ID
|
||||
if len(id) > 8 {
|
||||
id = id[:8]
|
||||
}
|
||||
prov := string(a.Provider)
|
||||
if prov == "" {
|
||||
prov = string(config.ProviderOllamaCloud)
|
||||
}
|
||||
created := a.Created.Format("2006-01-02 15:04 MST")
|
||||
fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", id, a.Name, maskKey(a.APIKey), created)
|
||||
fmt.Fprintf(w, "%-10s %-14s %-16s %-18s %-22s\n", id, prov, a.Name, maskKey(a.APIKey), created)
|
||||
}
|
||||
if af.BaseURL != "" {
|
||||
fmt.Fprintf(w, "\nDefault upstream: %s\n", af.BaseURL)
|
||||
|
|
|
|||
|
|
@ -14,13 +14,44 @@ import (
|
|||
"time"
|
||||
)
|
||||
|
||||
// Account is a single Ollama Cloud API key with an optional alias.
|
||||
// ProviderType identifies which upstream service an account belongs to.
|
||||
type ProviderType string
|
||||
|
||||
const (
|
||||
// ProviderOllamaCloud is the Ollama Cloud API at https://ollama.com.
|
||||
ProviderOllamaCloud ProviderType = "ollama-cloud"
|
||||
// ProviderOpenCodeGo is the OpenCode Go API at https://opencode.ai/zen/go/v1.
|
||||
ProviderOpenCodeGo ProviderType = "opencode-go"
|
||||
)
|
||||
|
||||
// DefaultProviderBaseURL returns the canonical upstream root for a provider.
|
||||
func DefaultProviderBaseURL(p ProviderType) string {
|
||||
switch p {
|
||||
case ProviderOpenCodeGo:
|
||||
return "https://opencode.ai/zen/go"
|
||||
default:
|
||||
return "https://ollama.com"
|
||||
}
|
||||
}
|
||||
|
||||
// ValidProvider reports whether p is a recognised provider type.
|
||||
func ValidProvider(p ProviderType) bool {
|
||||
switch p {
|
||||
case ProviderOllamaCloud, ProviderOpenCodeGo:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Account is a single upstream API key with an optional alias.
|
||||
type Account struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
APIKey string `json:"api_key"`
|
||||
BaseURL string `json:"base_url,omitempty"` // overrides ServerConfig.BaseURL; empty = use default
|
||||
Created time.Time `json:"created"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Provider ProviderType `json:"provider,omitempty"` // empty = ollama-cloud (backward compat)
|
||||
APIKey string `json:"api_key"`
|
||||
BaseURL string `json:"base_url,omitempty"` // overrides the provider default; empty = use default
|
||||
Created time.Time `json:"created"`
|
||||
}
|
||||
|
||||
// AccountsFile is the on-disk JSON structure.
|
||||
|
|
@ -73,6 +104,13 @@ func LoadAccounts() (*AccountsFile, error) {
|
|||
if err := json.Unmarshal(data, &af); err != nil {
|
||||
return nil, fmt.Errorf("parse accounts: %w", err)
|
||||
}
|
||||
// Backward compat: accounts without an explicit provider default to
|
||||
// ollama-cloud (the original behaviour before multi-provider support).
|
||||
for i := range af.Accounts {
|
||||
if af.Accounts[i].Provider == "" {
|
||||
af.Accounts[i].Provider = ProviderOllamaCloud
|
||||
}
|
||||
}
|
||||
return &af, nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -77,6 +77,60 @@ func TestValidateName(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestProviderType_Constants(t *testing.T) {
|
||||
if ProviderOllamaCloud != "ollama-cloud" {
|
||||
t.Errorf("ProviderOllamaCloud = %q", ProviderOllamaCloud)
|
||||
}
|
||||
if ProviderOpenCodeGo != "opencode-go" {
|
||||
t.Errorf("ProviderOpenCodeGo = %q", ProviderOpenCodeGo)
|
||||
}
|
||||
if !ValidProvider(ProviderOllamaCloud) || !ValidProvider(ProviderOpenCodeGo) {
|
||||
t.Error("ValidProvider returned false for known providers")
|
||||
}
|
||||
if ValidProvider(ProviderType("bogus")) {
|
||||
t.Error("ValidProvider returned true for bogus provider")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultProviderBaseURL(t *testing.T) {
|
||||
if DefaultProviderBaseURL(ProviderOllamaCloud) != "https://ollama.com" {
|
||||
t.Errorf("ollama base url = %q", DefaultProviderBaseURL(ProviderOllamaCloud))
|
||||
}
|
||||
if DefaultProviderBaseURL(ProviderOpenCodeGo) != "https://opencode.ai/zen/go" {
|
||||
t.Errorf("go base url = %q", DefaultProviderBaseURL(ProviderOpenCodeGo))
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadAccounts_BackwardCompatProvider verifies that accounts.json entries
|
||||
// without a "provider" field are normalised to ollama-cloud on load.
|
||||
func TestLoadAccounts_BackwardCompatProvider(t *testing.T) {
|
||||
p := withTempAccountsPath(t)
|
||||
// Write an old-style file with no provider field.
|
||||
old := `{
|
||||
"base_url": "https://ollama.com",
|
||||
"accounts": [
|
||||
{"id":"abc","name":"legacy","api_key":"sk-xyz","created":"2026-06-19T10:00:00Z"},
|
||||
{"id":"def","name":"go1","provider":"opencode-go","api_key":"sk-go","created":"2026-06-19T10:00:00Z"}
|
||||
]
|
||||
}`
|
||||
if err := os.WriteFile(p, []byte(old), 0o600); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
af, err := LoadAccounts()
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if len(af.Accounts) != 2 {
|
||||
t.Fatalf("len = %d", len(af.Accounts))
|
||||
}
|
||||
if af.Accounts[0].Name != "legacy" || af.Accounts[0].Provider != ProviderOllamaCloud {
|
||||
t.Errorf("legacy account provider = %q, want ollama-cloud", af.Accounts[0].Provider)
|
||||
}
|
||||
if af.Accounts[1].Provider != ProviderOpenCodeGo {
|
||||
t.Errorf("go account provider = %q, want opencode-go", af.Accounts[1].Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAccounts_MissingFile(t *testing.T) {
|
||||
withTempAccountsPath(t)
|
||||
af, err := LoadAccounts()
|
||||
|
|
|
|||
|
|
@ -27,9 +27,14 @@ type Balancer struct {
|
|||
// is applied uniformly to all accounts. At least one account is required;
|
||||
// otherwise Next returns ErrNoAccounts on every call.
|
||||
func NewBalancer(accounts []config.Account, cooldown time.Duration) *Balancer {
|
||||
// copy to avoid external mutation
|
||||
// copy to avoid external mutation; normalise empty provider to ollama-cloud.
|
||||
accts := make([]config.Account, len(accounts))
|
||||
copy(accts, accounts)
|
||||
for i, a := range accounts {
|
||||
if a.Provider == "" {
|
||||
a.Provider = config.ProviderOllamaCloud
|
||||
}
|
||||
accts[i] = a
|
||||
}
|
||||
return &Balancer{
|
||||
accounts: accts,
|
||||
cooldown: make(map[string]time.Time),
|
||||
|
|
@ -70,21 +75,41 @@ func (e ErrAllCooldown) Error() string {
|
|||
// rotates starting from the round-robin counter so consecutive calls land on
|
||||
// different accounts when possible. Returns ErrNoAccounts or ErrAllCooldown
|
||||
// when nothing is available.
|
||||
//
|
||||
// Next considers all accounts regardless of provider. Use NextFor to restrict
|
||||
// to a set of eligible providers (e.g. for model-based routing).
|
||||
func (b *Balancer) Next() (config.Account, error) {
|
||||
return b.NextFor(nil)
|
||||
}
|
||||
|
||||
// NextFor picks the next available account whose provider is in the eligible
|
||||
// list. When eligible is nil or empty, all providers are eligible (matches
|
||||
// the original Next behaviour). Accounts in cooldown are skipped.
|
||||
func (b *Balancer) NextFor(eligible []config.ProviderType) (config.Account, error) {
|
||||
if len(b.accounts) == 0 {
|
||||
return config.Account{}, ErrNoAccounts{}
|
||||
}
|
||||
allowed := make(map[config.ProviderType]bool, len(eligible))
|
||||
for _, p := range eligible {
|
||||
allowed[p] = true
|
||||
}
|
||||
allAllowed := len(eligible) == 0
|
||||
|
||||
now := time.Now()
|
||||
var earliest time.Time
|
||||
for i := 0; i < len(b.accounts); i++ {
|
||||
idx := int(b.rr.Add(1)) % len(b.accounts)
|
||||
acct := b.accounts[idx]
|
||||
// idx is computed from the *new* counter value; we Add-then-mod so each
|
||||
// call advances even if this loop iteration rejects the candidate.
|
||||
if !b.isCooldown(b.accounts[idx].ID, now) {
|
||||
return b.accounts[idx], nil
|
||||
if !allAllowed && !allowed[acct.Provider] {
|
||||
continue
|
||||
}
|
||||
if earliest.IsZero() || b.cooldownUntil(b.accounts[idx].ID).Before(earliest) {
|
||||
earliest = b.cooldownUntil(b.accounts[idx].ID)
|
||||
if !b.isCooldown(acct.ID, now) {
|
||||
return acct, nil
|
||||
}
|
||||
if earliest.IsZero() || b.cooldownUntil(acct.ID).Before(earliest) {
|
||||
earliest = b.cooldownUntil(acct.ID)
|
||||
}
|
||||
}
|
||||
return config.Account{}, ErrAllCooldown{Until: earliest}
|
||||
|
|
|
|||
|
|
@ -11,11 +11,25 @@ import (
|
|||
func mkAccts(n int) []config.Account {
|
||||
out := make([]config.Account, n)
|
||||
for i := range out {
|
||||
out[i] = config.Account{ID: "id" + string(rune('a'+i)), Name: "a" + string(rune('a'+i)), APIKey: "k" + string(rune('a'+i))}
|
||||
out[i] = config.Account{
|
||||
ID: "id" + string(rune('a'+i)),
|
||||
Name: "a" + string(rune('a'+i)),
|
||||
APIKey: "k" + string(rune('a'+i)),
|
||||
Provider: config.ProviderOllamaCloud,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func mkMixedAccts() []config.Account {
|
||||
return []config.Account{
|
||||
{ID: "idO1", Name: "o1", APIKey: "k1", Provider: config.ProviderOllamaCloud},
|
||||
{ID: "idO2", Name: "o2", APIKey: "k2", Provider: config.ProviderOllamaCloud},
|
||||
{ID: "idG1", Name: "g1", APIKey: "k3", Provider: config.ProviderOpenCodeGo},
|
||||
{ID: "idG2", Name: "g2", APIKey: "k4", Provider: config.ProviderOpenCodeGo},
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalancer_Next_RoundRobin(t *testing.T) {
|
||||
b := NewBalancer(mkAccts(3), 60*time.Second)
|
||||
seen := make(map[string]int)
|
||||
|
|
@ -158,3 +172,67 @@ func TestBalancer_CooldownExpires(t *testing.T) {
|
|||
t.Errorf("Next after cooldown expiry: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalancer_NextFor_FiltersByProvider(t *testing.T) {
|
||||
accts := mkMixedAccts()
|
||||
b := NewBalancer(accts, 60*time.Second)
|
||||
// Request only OpenCode Go accounts: every result must be a Go account.
|
||||
seen := map[string]int{}
|
||||
for i := 0; i < 8; i++ {
|
||||
a, err := b.NextFor([]config.ProviderType{config.ProviderOpenCodeGo})
|
||||
if err != nil {
|
||||
t.Fatalf("NextFor %d: %v", i, err)
|
||||
}
|
||||
if a.Provider != config.ProviderOpenCodeGo {
|
||||
t.Errorf("NextFor returned provider %q, want opencode-go", a.Provider)
|
||||
}
|
||||
seen[a.ID]++
|
||||
}
|
||||
if len(seen) != 2 {
|
||||
t.Errorf("expected 2 distinct go accounts, got %d", len(seen))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalancer_NextFor_MixedPool(t *testing.T) {
|
||||
accts := mkMixedAccts()
|
||||
b := NewBalancer(accts, 60*time.Second)
|
||||
// Eligible = both providers: all 4 accounts should be reachable.
|
||||
seen := map[string]bool{}
|
||||
for i := 0; i < 16; i++ {
|
||||
a, err := b.NextFor([]config.ProviderType{
|
||||
config.ProviderOllamaCloud, config.ProviderOpenCodeGo,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NextFor %d: %v", i, err)
|
||||
}
|
||||
seen[a.ID] = true
|
||||
}
|
||||
if len(seen) != 4 {
|
||||
t.Errorf("expected 4 distinct accounts in mixed pool, got %d", len(seen))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalancer_NextFor_NilEligibleReturnsAll(t *testing.T) {
|
||||
accts := mkMixedAccts()
|
||||
b := NewBalancer(accts, 60*time.Second)
|
||||
seen := map[config.ProviderType]bool{}
|
||||
for i := 0; i < 16; i++ {
|
||||
a, err := b.NextFor(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NextFor nil %d: %v", i, err)
|
||||
}
|
||||
seen[a.Provider] = true
|
||||
}
|
||||
if !seen[config.ProviderOllamaCloud] || !seen[config.ProviderOpenCodeGo] {
|
||||
t.Errorf("nil eligible should include both providers, got %v", seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalancer_NextFor_NoEligibleAccounts(t *testing.T) {
|
||||
// Only Ollama accounts configured, but requesting Go.
|
||||
b := NewBalancer(mkAccts(2), 60*time.Second)
|
||||
_, err := b.NextFor([]config.ProviderType{config.ProviderOpenCodeGo})
|
||||
if _, ok := err.(ErrAllCooldown); !ok && err == nil {
|
||||
t.Errorf("NextFor with no eligible accounts: %v, want error or ErrAllCooldown", err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -53,6 +53,12 @@ func (m *mockUpstream) handler(w http.ResponseWriter, r *http.Request) {
|
|||
|
||||
func newTestHandler(t *testing.T, upstreamURL string, accounts []config.Account, retries int) *Handler {
|
||||
t.Helper()
|
||||
// Apply the mock upstream URL to any account without an explicit BaseURL.
|
||||
for i := range accounts {
|
||||
if accounts[i].BaseURL == "" {
|
||||
accounts[i].BaseURL = upstreamURL
|
||||
}
|
||||
}
|
||||
b := NewBalancer(accounts, 50*time.Millisecond)
|
||||
return NewHandler(b, upstreamURL, retries, log.New("debug"))
|
||||
}
|
||||
|
|
@ -243,3 +249,198 @@ func TestIsAllowed(t *testing.T) {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandler_GoPathRoutesToGoProvider verifies /go/v1/* forces the
|
||||
// opencode-go provider and strips the /go prefix before forwarding upstream.
|
||||
func TestHandler_GoPathRoutesToGoProvider(t *testing.T) {
|
||||
var seenPath, seenAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
seenPath = r.URL.Path
|
||||
seenAuth = r.Header.Get("Authorization")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"ok":true}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
accts := []config.Account{
|
||||
{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: "http://unused-ollama"},
|
||||
{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: srv.URL},
|
||||
}
|
||||
h := newTestHandler(t, "http://unused", accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/go/v1/models", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
if seenPath != "/v1/models" {
|
||||
t.Errorf("upstream path = %q, want /v1/models", seenPath)
|
||||
}
|
||||
if seenAuth != "Bearer kG" {
|
||||
t.Errorf("upstream auth = %q, want go account key", seenAuth)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandler_ModelBasedRouting verifies a common model can be served by
|
||||
// either provider and a unique Go model routes only to Go accounts.
|
||||
func TestHandler_ModelBasedRouting(t *testing.T) {
|
||||
ollamaSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"from":"ollama"}`)
|
||||
}))
|
||||
goSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"from":"go"}`)
|
||||
}))
|
||||
defer ollamaSrv.Close()
|
||||
defer goSrv.Close()
|
||||
|
||||
accts := []config.Account{
|
||||
{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: ollamaSrv.URL},
|
||||
{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: goSrv.URL},
|
||||
}
|
||||
h := newTestHandler(t, "http://unused", accts, 2)
|
||||
|
||||
// Unique Go model must route to the Go upstream.
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("POST", "/v1/chat/completions",
|
||||
strings.NewReader(`{"model":"mimo-v2.5","messages":[]}`))
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("mimo status = %d", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), `"from":"go"`) {
|
||||
t.Errorf("mimo routed to wrong upstream: %s", rec.Body.String())
|
||||
}
|
||||
|
||||
// Unique Ollama model must route to the Ollama upstream.
|
||||
rec = httptest.NewRecorder()
|
||||
req = httptest.NewRequest("POST", "/v1/chat/completions",
|
||||
strings.NewReader(`{"model":"gpt-oss:20b","messages":[]}`))
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("gpt-oss status = %d", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), `"from":"ollama"`) {
|
||||
t.Errorf("gpt-oss routed to wrong upstream: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandler_MergedModels verifies /v1/models returns the union of both
|
||||
// catalogs with raw model IDs (no labels in the id field).
|
||||
func TestHandler_MergedModels(t *testing.T) {
|
||||
accts := []config.Account{
|
||||
{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: "http://unused"},
|
||||
{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: "http://unused"},
|
||||
}
|
||||
h := newTestHandler(t, "http://unused", accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/v1/models", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d", rec.Code)
|
||||
}
|
||||
body := rec.Body.String()
|
||||
// Common model appears with raw ID (no label).
|
||||
if !strings.Contains(body, `"id":"glm-5"`) {
|
||||
t.Errorf("common model glm-5 missing from merged list")
|
||||
}
|
||||
// Unique Ollama model appears with raw ID (no label in id).
|
||||
if !strings.Contains(body, `"id":"gpt-oss:20b"`) {
|
||||
t.Errorf("unique ollama model missing: %s", body)
|
||||
}
|
||||
// Unique Go model appears with raw ID.
|
||||
if !strings.Contains(body, `"id":"mimo-v2.5"`) {
|
||||
t.Errorf("unique go model missing: %s", body)
|
||||
}
|
||||
// No labels should appear in any id field.
|
||||
if strings.Contains(body, `(ollama)`) || strings.Contains(body, `(go)`) {
|
||||
t.Errorf("labels should not appear in id field: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandler_402UsageLimitTriggersFailover verifies 402 is treated as
|
||||
// retriable and triggers failover to the next account.
|
||||
func TestHandler_402UsageLimitTriggersFailover(t *testing.T) {
|
||||
badSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(402)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"usage limit exhausted"}}`)
|
||||
}))
|
||||
goodSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"ok":true}`)
|
||||
}))
|
||||
defer badSrv.Close()
|
||||
defer goodSrv.Close()
|
||||
|
||||
accts := []config.Account{
|
||||
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
|
||||
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
|
||||
}
|
||||
h := newTestHandler(t, "http://unused", accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200 after 402 failover", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsRetriable and TestShouldCooldown cover the retry classification.
|
||||
func TestIsRetriable(t *testing.T) {
|
||||
retriable := []int{429, 402, 500, 502, 503, 408}
|
||||
for _, s := range retriable {
|
||||
if !isRetriable(s) {
|
||||
t.Errorf("isRetriable(%d) = false, want true", s)
|
||||
}
|
||||
}
|
||||
nonRetriable := []int{200, 400, 401, 403, 404, 422}
|
||||
for _, s := range nonRetriable {
|
||||
if isRetriable(s) {
|
||||
t.Errorf("isRetriable(%d) = true, want false", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldCooldown(t *testing.T) {
|
||||
if !shouldCooldown(429, "rate limit") {
|
||||
t.Error("429 should cooldown")
|
||||
}
|
||||
if !shouldCooldown(402, "usage limit") {
|
||||
t.Error("402 should cooldown")
|
||||
}
|
||||
if shouldCooldown(500, "internal error") {
|
||||
t.Error("500 should NOT cooldown")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsBillingDisable(t *testing.T) {
|
||||
if !isBillingDisable("Insufficient credits for this account") {
|
||||
t.Error("insufficient credits should be billing disable")
|
||||
}
|
||||
if !isBillingDisable("credit balance too low") {
|
||||
t.Error("credit balance should be billing disable")
|
||||
}
|
||||
if isBillingDisable("usage limit exhausted") {
|
||||
t.Error("usage limit exhausted should NOT be billing disable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoutePath(t *testing.T) {
|
||||
up, prov, isGo := routePath("/go/v1/models")
|
||||
if up != "/v1/models" || prov != config.ProviderOpenCodeGo || !isGo {
|
||||
t.Errorf("routePath(/go/v1/models) = %q,%v,%v", up, prov, isGo)
|
||||
}
|
||||
up, prov, isGo = routePath("/v1/models")
|
||||
if up != "/v1/models" || prov != "" || isGo {
|
||||
t.Errorf("routePath(/v1/models) = %q,%v,%v", up, prov, isGo)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
116
internal/proxy/model_router.go
Normal file
116
internal/proxy/model_router.go
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
package proxy
|
||||
|
||||
import (
|
||||
"sort"
|
||||
|
||||
"github.com/Atte149/ollama-proxy/internal/config"
|
||||
)
|
||||
|
||||
// ollamaCloudModels is the set of model IDs served by Ollama Cloud.
|
||||
// These overlap partly with the OpenCode Go catalog.
|
||||
var ollamaCloudModels = []string{
|
||||
"deepseek-v3.1:671b", "deepseek-v3.2", "deepseek-v4-flash", "deepseek-v4-pro",
|
||||
"devstral-2:123b", "devstral-small-2:24b",
|
||||
"gemini-3-flash-preview", "gemma3:12b", "gemma3:27b", "gemma3:4b", "gemma4:31b",
|
||||
"glm-4.7", "glm-5", "glm-5.1", "glm-5.2",
|
||||
"gpt-oss:120b", "gpt-oss:20b",
|
||||
"kimi-k2.5", "kimi-k2.6", "kimi-k2.7-code",
|
||||
"minimax-m2.1", "minimax-m2.5", "minimax-m2.7", "minimax-m3",
|
||||
"ministral-3:14b", "ministral-3:3b", "ministral-3:8b",
|
||||
"mistral-large-3:675b",
|
||||
"nemotron-3-nano:30b", "nemotron-3-super", "nemotron-3-ultra",
|
||||
"qwen3-coder:480b", "qwen3-coder-next", "qwen3.5:397b",
|
||||
"rnj-1:8b",
|
||||
}
|
||||
|
||||
// openCodeGoModels is the set of model IDs served by OpenCode Go.
|
||||
var openCodeGoModels = []string{
|
||||
"deepseek-v4-flash", "deepseek-v4-pro",
|
||||
"glm-5", "glm-5.1", "glm-5.2",
|
||||
"hy3-preview",
|
||||
"kimi-k2.5", "kimi-k2.6", "kimi-k2.7-code",
|
||||
"mimo-v2-omni", "mimo-v2-pro", "mimo-v2.5", "mimo-v2.5-pro",
|
||||
"minimax-m2.5", "minimax-m2.7", "minimax-m3",
|
||||
"qwen3.5-plus", "qwen3.6-plus", "qwen3.7-max", "qwen3.7-plus",
|
||||
}
|
||||
|
||||
// modelProviders maps a model ID to the providers that can serve it.
|
||||
// Built once from the two catalogs above.
|
||||
var modelProviders = func() map[string][]config.ProviderType {
|
||||
m := make(map[string][]config.ProviderType, len(ollamaCloudModels)+len(openCodeGoModels))
|
||||
add := func(id string, p config.ProviderType) {
|
||||
m[id] = append(m[id], p)
|
||||
}
|
||||
for _, id := range ollamaCloudModels {
|
||||
add(id, config.ProviderOllamaCloud)
|
||||
}
|
||||
for _, id := range openCodeGoModels {
|
||||
// Deduplicate: if already present (common model), append; else create.
|
||||
found := false
|
||||
for _, ex := range m[id] {
|
||||
if ex == config.ProviderOpenCodeGo {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
add(id, config.ProviderOpenCodeGo)
|
||||
}
|
||||
}
|
||||
return m
|
||||
}()
|
||||
|
||||
// ProvidersForModel returns the providers eligible to serve the given model.
|
||||
// Unknown models default to Ollama Cloud (the original behaviour).
|
||||
func ProvidersForModel(model string) []config.ProviderType {
|
||||
if ps, ok := modelProviders[model]; ok {
|
||||
return ps
|
||||
}
|
||||
return []config.ProviderType{config.ProviderOllamaCloud}
|
||||
}
|
||||
|
||||
// AllModels returns the deduplicated, sorted list of all model IDs known to the
|
||||
// router (union of both catalogs).
|
||||
func AllModels() []string {
|
||||
seen := make(map[string]struct{})
|
||||
for id := range modelProviders {
|
||||
seen[id] = struct{}{}
|
||||
}
|
||||
out := make([]string, 0, len(seen))
|
||||
for id := range seen {
|
||||
out = append(out, id)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// IsCommonModel reports whether the model is served by both providers.
|
||||
func IsCommonModel(model string) bool {
|
||||
ps := ProvidersForModel(model)
|
||||
if len(ps) < 2 {
|
||||
return false
|
||||
}
|
||||
hasOllama, hasGo := false, false
|
||||
for _, p := range ps {
|
||||
switch p {
|
||||
case config.ProviderOllamaCloud:
|
||||
hasOllama = true
|
||||
case config.ProviderOpenCodeGo:
|
||||
hasGo = true
|
||||
}
|
||||
}
|
||||
return hasOllama && hasGo
|
||||
}
|
||||
|
||||
// ProviderLabel returns a short human label for a provider, used to annotate
|
||||
// unique models in the merged /v1/models response and in opencode config names.
|
||||
func ProviderLabel(p config.ProviderType) string {
|
||||
switch p {
|
||||
case config.ProviderOpenCodeGo:
|
||||
return "go"
|
||||
case config.ProviderOllamaCloud:
|
||||
return "ollama"
|
||||
default:
|
||||
return string(p)
|
||||
}
|
||||
}
|
||||
99
internal/proxy/model_router_test.go
Normal file
99
internal/proxy/model_router_test.go
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
package proxy
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Atte149/ollama-proxy/internal/config"
|
||||
)
|
||||
|
||||
func TestProvidersForModel_Common(t *testing.T) {
|
||||
for _, m := range []string{"glm-5", "glm-5.1", "glm-5.2", "kimi-k2.5", "kimi-k2.6",
|
||||
"kimi-k2.7-code", "deepseek-v4-flash", "deepseek-v4-pro",
|
||||
"minimax-m2.5", "minimax-m2.7", "minimax-m3"} {
|
||||
ps := ProvidersForModel(m)
|
||||
if len(ps) != 2 {
|
||||
t.Errorf("ProvidersForModel(%q) = %v, want 2 providers", m, ps)
|
||||
continue
|
||||
}
|
||||
if !IsCommonModel(m) {
|
||||
t.Errorf("IsCommonModel(%q) = false, want true", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersForModel_UniqueOllama(t *testing.T) {
|
||||
for _, m := range []string{"gpt-oss:20b", "gpt-oss:120b", "qwen3-coder:480b",
|
||||
"gemma3:4b", "nemotron-3-ultra", "devstral-2:123b", "rnj-1:8b",
|
||||
"mistral-large-3:675b", "minimax-m2.1", "glm-4.7", "gemini-3-flash-preview",
|
||||
"qwen3.5:397b", "qwen3-coder-next", "ministral-3:3b",
|
||||
"deepseek-v3.1:671b", "deepseek-v3.2"} {
|
||||
ps := ProvidersForModel(m)
|
||||
if len(ps) != 1 || ps[0] != config.ProviderOllamaCloud {
|
||||
t.Errorf("ProvidersForModel(%q) = %v, want [ollama-cloud]", m, ps)
|
||||
}
|
||||
if IsCommonModel(m) {
|
||||
t.Errorf("IsCommonModel(%q) = true, want false", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersForModel_UniqueGo(t *testing.T) {
|
||||
for _, m := range []string{"mimo-v2.5", "mimo-v2.5-pro", "mimo-v2-omni", "mimo-v2-pro",
|
||||
"qwen3.5-plus", "qwen3.6-plus", "qwen3.7-max", "qwen3.7-plus", "hy3-preview"} {
|
||||
ps := ProvidersForModel(m)
|
||||
if len(ps) != 1 || ps[0] != config.ProviderOpenCodeGo {
|
||||
t.Errorf("ProvidersForModel(%q) = %v, want [opencode-go]", m, ps)
|
||||
}
|
||||
if IsCommonModel(m) {
|
||||
t.Errorf("IsCommonModel(%q) = true, want false", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersForModel_UnknownDefaultsOllama(t *testing.T) {
|
||||
ps := ProvidersForModel("bogus-model")
|
||||
if len(ps) != 1 || ps[0] != config.ProviderOllamaCloud {
|
||||
t.Errorf("ProvidersForModel(unknown) = %v, want [ollama-cloud]", ps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllModels_DedupSorted(t *testing.T) {
|
||||
all := AllModels()
|
||||
if len(all) == 0 {
|
||||
t.Fatal("AllModels returned empty list")
|
||||
}
|
||||
// Verify sorted.
|
||||
for i := 1; i < len(all); i++ {
|
||||
if all[i-1] >= all[i] {
|
||||
t.Errorf("AllModels not sorted at %d: %q >= %q", i, all[i-1], all[i])
|
||||
}
|
||||
}
|
||||
// Verify no duplicates.
|
||||
seen := make(map[string]bool)
|
||||
for _, m := range all {
|
||||
if seen[m] {
|
||||
t.Errorf("duplicate model in AllModels: %s", m)
|
||||
}
|
||||
seen[m] = true
|
||||
}
|
||||
// Common models must appear once, not twice.
|
||||
common := "glm-5"
|
||||
count := 0
|
||||
for _, m := range all {
|
||||
if m == common {
|
||||
count++
|
||||
}
|
||||
}
|
||||
if count != 1 {
|
||||
t.Errorf("common model %s appeared %d times, want 1", common, count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderLabel(t *testing.T) {
|
||||
if ProviderLabel(config.ProviderOllamaCloud) != "ollama" {
|
||||
t.Error("ollama label wrong")
|
||||
}
|
||||
if ProviderLabel(config.ProviderOpenCodeGo) != "go" {
|
||||
t.Error("go label wrong")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,28 +1,63 @@
|
|||
package proxy
|
||||
|
||||
// retry.go holds the failover helpers used by handler.go. The retry loop itself
|
||||
// lives inside Handler.proxyWithFailover (see handler.go) because it needs
|
||||
// tight control over the moment a response starts streaming vs. is rejected
|
||||
// pre-stream. This file documents the invariants the loop must maintain.
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Streaming invariant
|
||||
// ===================
|
||||
// Once an upstream account has returned a 2xx status AND the proxy has started
|
||||
// writing the response body to the client (a single byte flushed), the request
|
||||
// is committed: we MUST NOT switch accounts for that request. Any mid-stream
|
||||
// upstream error is surfaced to the client as-is (truncated response); we never
|
||||
// attempt to "restart" a streamed request on a different account, because the
|
||||
// client has already received partial output and a retry would duplicate it.
|
||||
// isRetriable reports whether an upstream status code should trigger failover
|
||||
// to the next account before any bytes are streamed to the client.
|
||||
//
|
||||
// Pre-stream failover
|
||||
// -------------------
|
||||
// The window in which we CAN retry on another account is exactly:
|
||||
// 1. The upstream HTTP request returned an error (network, timeout, EOF
|
||||
// before any response).
|
||||
// 2. The upstream returned 429 (Too Many Requests) — we mark the account
|
||||
// cooldown and try the next.
|
||||
// 3. The upstream returned 5xx — we try the next account WITHOUT marking
|
||||
// cooldown (5xx may be transient and is not necessarily a rate limit).
|
||||
// 429 — rate limit / usage limit (retryable cooldown)
|
||||
// 402 — payment required; treated as retryable when the body text matches a
|
||||
//
|
||||
// Once copyResponse has called WriteHeader, no further retries are possible
|
||||
// for this request.
|
||||
// usage-window pattern ("usage limit", "limit reached", "exhausted",
|
||||
// "daily/weekly/monthly limit"). A bare "insufficient credits" is also
|
||||
// retriable (billing disable) — we still fail over, just with a longer
|
||||
// cooldown.
|
||||
//
|
||||
// 5xx — transient server error
|
||||
// 408 — request timeout
|
||||
// Other 4xx are not retriable: the request is malformed or the model is
|
||||
// unknown, so retrying on another account would fail identically.
|
||||
func isRetriable(status int) bool {
|
||||
switch {
|
||||
case status == http.StatusTooManyRequests:
|
||||
return true
|
||||
case status == http.StatusPaymentRequired:
|
||||
return true
|
||||
case status >= 500:
|
||||
return true
|
||||
case status == http.StatusRequestTimeout:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// shouldCooldown reports whether an account should be marked cooldown after a
|
||||
// retriable failure. 5xx does not trigger cooldown (the error may be
|
||||
// transient and unrelated to this account's quota). 429 and 402 always do.
|
||||
func shouldCooldown(status int, bodyPreview string) bool {
|
||||
switch {
|
||||
case status == http.StatusTooManyRequests:
|
||||
return true
|
||||
case status == http.StatusPaymentRequired:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// isBillingDisable reports whether the error body indicates a billing/credit
|
||||
// exhaustion (longer cooldown) rather than a transient usage-window limit.
|
||||
// Used for logging/diagnostics; the cooldown duration is the same in v1.
|
||||
func isBillingDisable(bodyPreview string) bool {
|
||||
low := strings.ToLower(bodyPreview)
|
||||
for _, marker := range []string{"insufficient credits", "credit balance", "billing"} {
|
||||
if strings.Contains(low, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue