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:
parent
98bbc96bf5
commit
4fe15324c8
11 changed files with 973 additions and 125 deletions
35
AGENTS.md
35
AGENTS.md
|
|
@ -2,7 +2,7 @@
|
||||||
|
|
||||||
## Контекст
|
## Контекст
|
||||||
|
|
||||||
`ollama-proxy` — Go-прокси для Ollama Cloud (`https://ollama.com`). Балансирует запросы по нескольким API-ключам с round-robin и failover на 429/5xx. Слушает `127.0.0.1:11435`, экспонирует нативный Ollama API (`/api/*`) и OpenAI-совместимый (`/v1/*`) passthrough. Деплой — systemd unit `ollama-proxy.service`. Используется opencode через кастомный провайдер `ocp`.
|
`ollama-proxy` — Go-прокси для Ollama Cloud (`https://ollama.com`) и OpenCode Go (`https://opencode.ai/zen/go/v1`). Балансирует запросы по нескольким API-ключам с round-robin и failover на 429/402/5xx. Модели общие для обоих провайдеров идут одним пулом, уникальные — маршрутизируются к нужному провайдеру. Слушает `127.0.0.1:11435`, экспонирует OpenAI-совместимый (`/v1/*`) и нативный Ollama (`/api/*`) API. Деплой — systemd unit `ollama-proxy.service`. Используется opencode через кастомный провайдер `oc`.
|
||||||
|
|
||||||
## Стек
|
## Стек
|
||||||
|
|
||||||
|
|
@ -40,43 +40,50 @@ systemd/ollama-proxy.service — systemd unit
|
||||||
|
|
||||||
## Типичные задачи
|
## Типичные задачи
|
||||||
|
|
||||||
- **Добавить аккаунт:** `ollama-proxy accounts add --name <alias>` (промпт) или `... add <key> --name <alias>`
|
- **Добавить аккаунт:** `ollama-proxy accounts add --name <alias> --provider ollama-cloud|opencode-go` (промпт) или `... add <key> --name <alias> --provider <P>`
|
||||||
- **Перезапустить после правки аккаунтов:** `sudo systemctl restart ollama-proxy` (прокси читает accounts.json при старте; runtime-обновления не поддерживаются)
|
- **Перезапустить после правки аккаунтов:** `sudo systemctl restart ollama-proxy` (прокси читает accounts.json при старте; runtime-обновления не поддерживаются)
|
||||||
- **Посмотреть логи:** `journalctl -u ollama-proxy -f` — JSON-структурированные (account, account_id, path, status, latency_ms, stream)
|
- **Посмотреть логи:** `journalctl -u ollama-proxy -f` — JSON-структурированные (account, account_id, provider, path, status, latency_ms, stream)
|
||||||
- **Проверить, что прокси жив:** `curl http://127.0.0.1:11435/api/version` → должен вернуть JSON от ollama.com
|
- **Проверить, что прокси жив:** `curl http://127.0.0.1:11435/api/version` → должен вернуть JSON от ollama.com
|
||||||
- **Проверить список моделей:** `curl http://127.0.0.1:11435/v1/models | jq '.data | length'` (~35 моделей)
|
- **Проверить список моделей:** `curl http://127.0.0.1:11435/v1/models | jq '.data | length'` (~44 модели, merged из Ollama Cloud + OpenCode Go)
|
||||||
- **Запустить opencode с пулом:** `openlama` (обёртка в `~/.local/bin/openlama`, выставляет NO_PROXY и запускает opencode)
|
- **Запустить opencode с пулом:** `openlama` (обёртка в `~/.local/bin/openlama`, выставляет NO_PROXY и запускает opencode)
|
||||||
- **Обновить список моделей в opencode-конфиге:** `curl -sS http://127.0.0.1:11435/v1/models | jq -r '.data[].id'` → обновить `models` блок в `~/.config/opencode/opencode.json`
|
- **Обновить список моделей в opencode-конфиге:** `curl -sS http://127.0.0.1:11435/v1/models | jq -r '.data[].id'` → обновить `models` блок в `~/.config/opencode/openlama.json`
|
||||||
|
|
||||||
## Интеграция с opencode
|
## Интеграция с opencode
|
||||||
|
|
||||||
**Два отдельных конфига:** `opencode` видит только marg/fireworks, `openlama` видит только ocp.
|
**Два отдельных конфига:** `opencode` видит только marg/fireworks, `openlama` видит только `oc` (объединённый пул Ollama Cloud + OpenCode Go).
|
||||||
|
|
||||||
- `~/.config/opencode/opencode.json` — глобальный конфиг с `marg` + `fireworks` (без `ocp`)
|
- `~/.config/opencode/opencode.json` — глобальный конфиг с `marg` + `fireworks` (без `oc`)
|
||||||
- `~/.config/opencode/openlama.json` — отдельный конфиг с `ocp` + `"enabled_providers": ["ocp"]`. Содержит те же instructions/mcp/plugins/agent что глобальный, но только провайдер `ocp`.
|
- `~/.config/opencode/openlama.json` — отдельный конфиг с `oc` + `"enabled_providers": ["oc"]`. Содержит те же instructions/mcp/plugins/agent что глобальный, но только провайдер `oc`.
|
||||||
- `~/.local/bin/openlama` — обёртка, выставляет `OPENCODE_CONFIG=~/.config/opencode/openlama.json` + `NO_PROXY=127.0.0.1,localhost,::1` + `exec opencode "$@"`
|
- `~/.local/bin/openlama` — обёртка, выставляет `OPENCODE_CONFIG=~/.config/opencode/openlama.json` + `NO_PROXY=127.0.0.1,localhost,::1` + `exec opencode "$@"`
|
||||||
|
|
||||||
openlama.json (ключевая часть):
|
openlama.json (ключевая часть):
|
||||||
```json
|
```json
|
||||||
{
|
{
|
||||||
"enabled_providers": ["ocp"],
|
"enabled_providers": ["oc"],
|
||||||
"provider": {
|
"provider": {
|
||||||
"ocp": {
|
"oc": {
|
||||||
"npm": "@ai-sdk/openai-compatible",
|
"npm": "@ai-sdk/openai-compatible",
|
||||||
"name": "Ollama Cloud (pool)",
|
"name": "Ollama + OpenCode Go (pool)",
|
||||||
"options": {
|
"options": {
|
||||||
"baseURL": "http://127.0.0.1:11435/v1",
|
"baseURL": "http://127.0.0.1:11435/v1",
|
||||||
"apiKey": "proxy"
|
"apiKey": "proxy"
|
||||||
},
|
},
|
||||||
"models": {
|
"models": {
|
||||||
"gpt-oss:20b": { "name": "gpt-oss:20b" },
|
"glm-5": { "name": "glm-5" },
|
||||||
"...": { "name": "..." }
|
"gpt-oss:20b": { "name": "gpt-oss:20b (ollama)" },
|
||||||
|
"mimo-v2.5": { "name": "mimo-v2.5 (go)" }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
`OPENCODE_CONFIG` мерджится с глобальным конфигом, но `enabled_providers: ["ocp"]` перекрывает глобальный — `marg`/`fireworks` скрыты.
|
`OPENCODE_CONFIG` мерджится с глобальным конфигом, но `enabled_providers: ["oc"]` перекрывает глобальный — `marg`/`fireworks` скрыты.
|
||||||
|
|
||||||
|
**Маршрутизация моделей:** прокси парсит `model` из тела запроса и выбирает eligible провайдеров:
|
||||||
|
- Общие модели (glm-5, kimi-k2.6, deepseek-v4-flash и др.) → пул всех аккаунтов обоих провайдеров
|
||||||
|
- Уникальные Ollama (gpt-oss:20b, qwen3-coder, gemma, nemotron) → только ollama-cloud аккаунты
|
||||||
|
- Уникальные Go (mimo-v2.5, qwen3.5-plus, qwen3.7-max) → только opencode-go аккаунты
|
||||||
|
- Также: `/go/v1/*` path force-рует opencode-go провайдер
|
||||||
|
|
||||||
**Модель не сделана дефолтной** — выбирается через `/models` в TUI.
|
**Модель не сделана дефолтной** — выбирается через `/models` в TUI.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -37,16 +37,21 @@ func RunAccounts(args []string) int {
|
||||||
|
|
||||||
func printAccountsUsage() {
|
func printAccountsUsage() {
|
||||||
fmt.Print(`Usage:
|
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 list
|
||||||
ollama-proxy accounts remove <id|name>
|
ollama-proxy accounts remove <id|name>
|
||||||
ollama-proxy accounts set-base-url <url>
|
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 {
|
func accountsAdd(args []string) int {
|
||||||
var apiKey, name string
|
var apiKey, name string
|
||||||
|
provider := config.ProviderOllamaCloud
|
||||||
for i := 0; i < len(args); i++ {
|
for i := 0; i < len(args); i++ {
|
||||||
switch args[i] {
|
switch args[i] {
|
||||||
case "--name":
|
case "--name":
|
||||||
|
|
@ -56,8 +61,15 @@ func accountsAdd(args []string) int {
|
||||||
}
|
}
|
||||||
name = args[i+1]
|
name = args[i+1]
|
||||||
i++
|
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":
|
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
|
return 0
|
||||||
default:
|
default:
|
||||||
if apiKey == "" {
|
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 == "" {
|
if name == "" {
|
||||||
name = promptString("Account name (alias)", "acct"+config.NewID()[:4])
|
name = promptString("Account name (alias)", "acct"+config.NewID()[:4])
|
||||||
}
|
}
|
||||||
|
|
@ -78,7 +95,14 @@ func accountsAdd(args []string) int {
|
||||||
}
|
}
|
||||||
|
|
||||||
if apiKey == "" {
|
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()))
|
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||||
fmt.Println()
|
fmt.Println()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
@ -108,6 +132,7 @@ func accountsAdd(args []string) int {
|
||||||
acct := config.Account{
|
acct := config.Account{
|
||||||
ID: config.NewID(),
|
ID: config.NewID(),
|
||||||
Name: name,
|
Name: name,
|
||||||
|
Provider: provider,
|
||||||
APIKey: apiKey,
|
APIKey: apiKey,
|
||||||
Created: nowFn(),
|
Created: nowFn(),
|
||||||
}
|
}
|
||||||
|
|
@ -116,7 +141,8 @@ func accountsAdd(args []string) int {
|
||||||
fmt.Fprintln(os.Stderr, "save accounts:", err)
|
fmt.Fprintln(os.Stderr, "save accounts:", err)
|
||||||
return 1
|
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
|
return 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -133,15 +159,21 @@ func accountsList(args []string) int {
|
||||||
}
|
}
|
||||||
w := bufio.NewWriter(os.Stdout)
|
w := bufio.NewWriter(os.Stdout)
|
||||||
defer w.Flush()
|
defer w.Flush()
|
||||||
fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", "ID", "NAME", "KEY", "CREATED")
|
fmt.Fprintf(w, "%-10s %-14s %-16s %-18s %-22s\n", "ID", "PROVIDER", "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",
|
||||||
|
strings.Repeat("-", 8), strings.Repeat("-", 12), strings.Repeat("-", 14),
|
||||||
|
strings.Repeat("-", 16), strings.Repeat("-", 20))
|
||||||
for _, a := range af.Accounts {
|
for _, a := range af.Accounts {
|
||||||
id := a.ID
|
id := a.ID
|
||||||
if len(id) > 8 {
|
if len(id) > 8 {
|
||||||
id = 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")
|
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 != "" {
|
if af.BaseURL != "" {
|
||||||
fmt.Fprintf(w, "\nDefault upstream: %s\n", af.BaseURL)
|
fmt.Fprintf(w, "\nDefault upstream: %s\n", af.BaseURL)
|
||||||
|
|
|
||||||
|
|
@ -14,12 +14,43 @@ import (
|
||||||
"time"
|
"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 {
|
type Account struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
Provider ProviderType `json:"provider,omitempty"` // empty = ollama-cloud (backward compat)
|
||||||
APIKey string `json:"api_key"`
|
APIKey string `json:"api_key"`
|
||||||
BaseURL string `json:"base_url,omitempty"` // overrides ServerConfig.BaseURL; empty = use default
|
BaseURL string `json:"base_url,omitempty"` // overrides the provider default; empty = use default
|
||||||
Created time.Time `json:"created"`
|
Created time.Time `json:"created"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -73,6 +104,13 @@ func LoadAccounts() (*AccountsFile, error) {
|
||||||
if err := json.Unmarshal(data, &af); err != nil {
|
if err := json.Unmarshal(data, &af); err != nil {
|
||||||
return nil, fmt.Errorf("parse accounts: %w", err)
|
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
|
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) {
|
func TestLoadAccounts_MissingFile(t *testing.T) {
|
||||||
withTempAccountsPath(t)
|
withTempAccountsPath(t)
|
||||||
af, err := LoadAccounts()
|
af, err := LoadAccounts()
|
||||||
|
|
|
||||||
|
|
@ -27,9 +27,14 @@ type Balancer struct {
|
||||||
// is applied uniformly to all accounts. At least one account is required;
|
// is applied uniformly to all accounts. At least one account is required;
|
||||||
// otherwise Next returns ErrNoAccounts on every call.
|
// otherwise Next returns ErrNoAccounts on every call.
|
||||||
func NewBalancer(accounts []config.Account, cooldown time.Duration) *Balancer {
|
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))
|
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{
|
return &Balancer{
|
||||||
accounts: accts,
|
accounts: accts,
|
||||||
cooldown: make(map[string]time.Time),
|
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
|
// rotates starting from the round-robin counter so consecutive calls land on
|
||||||
// different accounts when possible. Returns ErrNoAccounts or ErrAllCooldown
|
// different accounts when possible. Returns ErrNoAccounts or ErrAllCooldown
|
||||||
// when nothing is available.
|
// 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) {
|
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 {
|
if len(b.accounts) == 0 {
|
||||||
return config.Account{}, ErrNoAccounts{}
|
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()
|
now := time.Now()
|
||||||
var earliest time.Time
|
var earliest time.Time
|
||||||
for i := 0; i < len(b.accounts); i++ {
|
for i := 0; i < len(b.accounts); i++ {
|
||||||
idx := int(b.rr.Add(1)) % len(b.accounts)
|
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
|
// idx is computed from the *new* counter value; we Add-then-mod so each
|
||||||
// call advances even if this loop iteration rejects the candidate.
|
// call advances even if this loop iteration rejects the candidate.
|
||||||
if !b.isCooldown(b.accounts[idx].ID, now) {
|
if !allAllowed && !allowed[acct.Provider] {
|
||||||
return b.accounts[idx], nil
|
continue
|
||||||
}
|
}
|
||||||
if earliest.IsZero() || b.cooldownUntil(b.accounts[idx].ID).Before(earliest) {
|
if !b.isCooldown(acct.ID, now) {
|
||||||
earliest = b.cooldownUntil(b.accounts[idx].ID)
|
return acct, nil
|
||||||
|
}
|
||||||
|
if earliest.IsZero() || b.cooldownUntil(acct.ID).Before(earliest) {
|
||||||
|
earliest = b.cooldownUntil(acct.ID)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return config.Account{}, ErrAllCooldown{Until: earliest}
|
return config.Account{}, ErrAllCooldown{Until: earliest}
|
||||||
|
|
|
||||||
|
|
@ -11,11 +11,25 @@ import (
|
||||||
func mkAccts(n int) []config.Account {
|
func mkAccts(n int) []config.Account {
|
||||||
out := make([]config.Account, n)
|
out := make([]config.Account, n)
|
||||||
for i := range out {
|
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
|
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) {
|
func TestBalancer_Next_RoundRobin(t *testing.T) {
|
||||||
b := NewBalancer(mkAccts(3), 60*time.Second)
|
b := NewBalancer(mkAccts(3), 60*time.Second)
|
||||||
seen := make(map[string]int)
|
seen := make(map[string]int)
|
||||||
|
|
@ -158,3 +172,67 @@ func TestBalancer_CooldownExpires(t *testing.T) {
|
||||||
t.Errorf("Next after cooldown expiry: %v", err)
|
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
|
package proxy
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
|
@ -8,41 +9,58 @@ import (
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Handler is the http.Handler that proxies incoming requests to Ollama Cloud
|
// Handler is the http.Handler that proxies incoming requests to upstream
|
||||||
// through the balancer, applying per-account Authorization and failover logic.
|
// providers (Ollama Cloud and OpenCode Go) through the balancer, applying
|
||||||
|
// per-account Authorization and failover logic.
|
||||||
type Handler struct {
|
type Handler struct {
|
||||||
Balancer *Balancer
|
Balancer *Balancer
|
||||||
BaseURL string // upstream root, e.g. "https://ollama.com"
|
|
||||||
Client *http.Client
|
|
||||||
Log *slog.Logger
|
Log *slog.Logger
|
||||||
|
Client *http.Client
|
||||||
Retries int // max attempts per request (= number of accounts to try)
|
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
|
// NewHandler builds a Handler with sensible HTTP client defaults.
|
||||||
// on the overall request — streaming responses can be long; per-read timeout
|
|
||||||
// is governed by the caller's context).
|
|
||||||
func NewHandler(b *Balancer, baseURL string, retries int, log *slog.Logger) *Handler {
|
func NewHandler(b *Balancer, baseURL string, retries int, log *slog.Logger) *Handler {
|
||||||
return &Handler{
|
return &Handler{
|
||||||
Balancer: b,
|
Balancer: b,
|
||||||
BaseURL: strings.TrimRight(baseURL, "/"),
|
Log: log,
|
||||||
Client: &http.Client{
|
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,
|
Timeout: 0,
|
||||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
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.
|
// upstreamBaseURL resolves the canonical upstream root for an account:
|
||||||
// Anything else returns 404 — we never proxy arbitrary paths.
|
// 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{
|
var allowedPaths = map[string]bool{
|
||||||
"/api/chat": true,
|
"/api/chat": true,
|
||||||
"/api/generate": true,
|
"/api/generate": true,
|
||||||
|
|
@ -58,13 +76,11 @@ var allowedPaths = map[string]bool{
|
||||||
"/v1/files": true,
|
"/v1/files": true,
|
||||||
}
|
}
|
||||||
|
|
||||||
// isAllowed reports whether a path (possibly with a trailing slash or query)
|
// isAllowed reports whether a path matches an allowed upstream endpoint.
|
||||||
// matches one of the allowed upstream endpoints.
|
|
||||||
func isAllowed(path string) bool {
|
func isAllowed(path string) bool {
|
||||||
if allowedPaths[path] {
|
if allowedPaths[path] {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
// allow sub-paths like /v1/files/<id> under an allowed prefix.
|
|
||||||
for prefix := range allowedPaths {
|
for prefix := range allowedPaths {
|
||||||
if strings.HasPrefix(path, prefix+"/") {
|
if strings.HasPrefix(path, prefix+"/") {
|
||||||
return true
|
return true
|
||||||
|
|
@ -73,9 +89,72 @@ func isAllowed(path string) bool {
|
||||||
return false
|
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) {
|
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)
|
http.NotFound(w, r)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -84,32 +163,101 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
return
|
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
|
maxAttempts := h.Retries
|
||||||
if maxAttempts < 1 {
|
if maxAttempts < 1 {
|
||||||
maxAttempts = 1
|
maxAttempts = 1
|
||||||
}
|
}
|
||||||
if maxAttempts > h.Balancer.Len() {
|
if n := h.Balancer.Len(); maxAttempts > n {
|
||||||
maxAttempts = h.Balancer.Len()
|
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
|
// handleModelsEndpoint returns a merged list of all known models for
|
||||||
// begins streaming (returns headers + a 2xx) commits: we copy the rest of the
|
// /v1/models (OpenAI format) or /api/tags (Ollama format).
|
||||||
// response to the client verbatim and stop retrying. Pre-stream 429/5xx move
|
func (h *Handler) handleModelsEndpoint(w http.ResponseWriter, r *http.Request, path string) {
|
||||||
// on to the next account.
|
now := time.Now().Unix()
|
||||||
func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxAttempts int) {
|
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()
|
ctx := r.Context()
|
||||||
var lastErr error
|
var lastErr error
|
||||||
|
|
||||||
for attempt := 0; attempt < maxAttempts; attempt++ {
|
for attempt := 0; attempt < maxAttempts; attempt++ {
|
||||||
acct, err := h.Balancer.Next()
|
acct, err := h.Balancer.NextFor(eligible)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// No accounts available.
|
|
||||||
switch err.(type) {
|
switch err.(type) {
|
||||||
case ErrAllCooldown:
|
case ErrAllCooldown:
|
||||||
// All accounts rate-limited upstream → propagate 429 to client.
|
|
||||||
msg := err.Error()
|
msg := err.Error()
|
||||||
if lastErr != nil {
|
if lastErr != nil {
|
||||||
msg = msg + "; last upstream error: " + lastErr.Error()
|
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.
|
// Build the upstream request.
|
||||||
upstreamURL := h.BaseURL + r.URL.RequestURI()
|
base := h.upstreamBaseURL(acct)
|
||||||
// Per-account BaseURL override wins when set.
|
// Build the upstream URI from the routed path + original query string.
|
||||||
if acct.BaseURL != "" {
|
uri := upstreamPath
|
||||||
upstreamURL = strings.TrimRight(acct.BaseURL, "/") + r.URL.RequestURI()
|
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 {
|
if err != nil {
|
||||||
lastErr = fmt.Errorf("build request: %w", err)
|
lastErr = fmt.Errorf("build request: %w", err)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
// Copy headers, replacing Authorization with the account key.
|
|
||||||
copyHeaders(req.Header, r.Header)
|
copyHeaders(req.Header, r.Header)
|
||||||
req.Header.Set("Authorization", "Bearer "+acct.APIKey)
|
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()
|
start := time.Now()
|
||||||
resp, err := h.Client.Do(req)
|
resp, err := h.Client.Do(req)
|
||||||
latency := time.Since(start)
|
latency := time.Since(start)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// Network/timeout error: log and try next account.
|
|
||||||
h.Log.Warn("upstream request failed",
|
h.Log.Warn("upstream request failed",
|
||||||
"account", acct.Name, "account_id", acct.ID,
|
"account", acct.Name, "account_id", acct.ID,
|
||||||
|
"provider", string(acct.Provider),
|
||||||
"path", r.URL.Path, "latency_ms", latency.Milliseconds(), "err", err)
|
"path", r.URL.Path, "latency_ms", latency.Milliseconds(), "err", err)
|
||||||
lastErr = err
|
lastErr = err
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
// Pre-stream decision: 429 or 5xx → cooldown + failover.
|
// Pre-stream decision: retriable errors → cooldown + failover.
|
||||||
if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
if isRetriable(resp.StatusCode) {
|
||||||
bodyPreview := readPreview(resp.Body, 256)
|
bodyPreview := readPreview(resp.Body, 512)
|
||||||
resp.Body.Close()
|
resp.Body.Close()
|
||||||
h.Log.Warn("upstream rejected, failover",
|
h.Log.Warn("upstream rejected, failover",
|
||||||
"account", acct.Name, "account_id", acct.ID,
|
"account", acct.Name, "account_id", acct.ID,
|
||||||
|
"provider", string(acct.Provider),
|
||||||
"path", r.URL.Path, "status", resp.StatusCode,
|
"path", r.URL.Path, "status", resp.StatusCode,
|
||||||
"latency_ms", latency.Milliseconds(), "preview", bodyPreview)
|
"latency_ms", latency.Milliseconds(), "preview", bodyPreview)
|
||||||
if resp.StatusCode == http.StatusTooManyRequests {
|
if shouldCooldown(resp.StatusCode, bodyPreview) {
|
||||||
h.Balancer.MarkCooldown(acct.ID)
|
h.Balancer.MarkCooldown(acct.ID)
|
||||||
}
|
}
|
||||||
lastErr = fmt.Errorf("upstream %d: %s", resp.StatusCode, bodyPreview)
|
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.copyResponse(w, resp)
|
||||||
h.Log.Info("proxied",
|
h.Log.Info("proxied",
|
||||||
"account", acct.Name, "account_id", acct.ID,
|
"account", acct.Name, "account_id", acct.ID,
|
||||||
|
"provider", string(acct.Provider),
|
||||||
"path", r.URL.Path, "status", resp.StatusCode,
|
"path", r.URL.Path, "status", resp.StatusCode,
|
||||||
"latency_ms", latency.Milliseconds(), "stream", isStreaming(resp))
|
"latency_ms", latency.Milliseconds(), "stream", isStreaming(resp))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Exhausted retries. Distinguish "everything is rate-limited" (→ 429) from
|
// Exhausted retries.
|
||||||
// "mixed upstream errors" (→ 502).
|
|
||||||
if h.Balancer.Len() > 0 {
|
if h.Balancer.Len() > 0 {
|
||||||
allCooldown := true
|
allCooldown := true
|
||||||
status := h.Balancer.Status()
|
status := h.Balancer.Status()
|
||||||
for _, a := range h.Balancer.Accounts() {
|
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 {
|
if !status[a.ID].InCooldown {
|
||||||
allCooldown = false
|
allCooldown = false
|
||||||
break
|
break
|
||||||
|
|
@ -204,25 +379,21 @@ func (h *Handler) proxyWithFailover(w http.ResponseWriter, r *http.Request, maxA
|
||||||
writeError(w, http.StatusBadGateway, "all accounts failed")
|
writeError(w, http.StatusBadGateway, "all accounts failed")
|
||||||
}
|
}
|
||||||
|
|
||||||
// copyResponse streams the upstream response body to the client. For chunked /
|
// copyResponse streams the upstream response body to the client.
|
||||||
// SSE responses we flush after every read so tokens reach the client immediately.
|
|
||||||
func (h *Handler) copyResponse(w http.ResponseWriter, resp *http.Response) {
|
func (h *Handler) copyResponse(w http.ResponseWriter, resp *http.Response) {
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
// Copy headers (except hop-by-hop ones).
|
|
||||||
for k, vs := range resp.Header {
|
for k, vs := range resp.Header {
|
||||||
for _, v := range vs {
|
for _, v := range vs {
|
||||||
w.Header().Add(k, v)
|
w.Header().Add(k, v)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
w.WriteHeader(resp.StatusCode)
|
w.WriteHeader(resp.StatusCode)
|
||||||
|
|
||||||
flusher, _ := w.(http.Flusher)
|
flusher, _ := w.(http.Flusher)
|
||||||
buf := make([]byte, 4096)
|
buf := make([]byte, 4096)
|
||||||
for {
|
for {
|
||||||
n, err := resp.Body.Read(buf)
|
n, err := resp.Body.Read(buf)
|
||||||
if n > 0 {
|
if n > 0 {
|
||||||
if _, werr := w.Write(buf[:n]); werr != nil {
|
if _, werr := w.Write(buf[:n]); werr != nil {
|
||||||
// client went away; stop copying silently
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if flusher != nil {
|
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
|
// 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) {
|
func copyHeaders(dst, src http.Header) {
|
||||||
hopByHop := []string{
|
hopByHop := []string{
|
||||||
"Connection", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization",
|
"Connection", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization",
|
||||||
"Te", "Trailers", "Transfer-Encoding", "Upgrade",
|
"Te", "Trailers", "Transfer-Encoding", "Upgrade",
|
||||||
"Authorization", // always overwritten by the account key
|
"Authorization",
|
||||||
}
|
}
|
||||||
for k, vs := range src {
|
for k, vs := range src {
|
||||||
skip := false
|
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 {
|
func readPreview(r io.ReadCloser, n int) string {
|
||||||
defer r.Close()
|
defer r.Close()
|
||||||
buf := make([]byte, n)
|
buf := make([]byte, n)
|
||||||
|
|
@ -272,21 +441,15 @@ func readPreview(r io.ReadCloser, n int) string {
|
||||||
return string(buf[:m])
|
return string(buf[:m])
|
||||||
}
|
}
|
||||||
|
|
||||||
// isStreaming reports whether the response is streaming (SSE or NDJSON/chunked).
|
|
||||||
func isStreaming(resp *http.Response) bool {
|
func isStreaming(resp *http.Response) bool {
|
||||||
ct := resp.Header.Get("Content-Type")
|
ct := resp.Header.Get("Content-Type")
|
||||||
if strings.Contains(ct, "text/event-stream") {
|
if strings.Contains(ct, "text/event-stream") || strings.Contains(ct, "application/x-ndjson") {
|
||||||
return true
|
|
||||||
}
|
|
||||||
if strings.Contains(ct, "application/x-ndjson") {
|
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
te := resp.Header.Get("Transfer-Encoding")
|
te := resp.Header.Get("Transfer-Encoding")
|
||||||
return strings.Contains(strings.ToLower(te), "chunked")
|
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) {
|
func writeError(w http.ResponseWriter, status int, msg string) {
|
||||||
w.Header().Set("Content-Type", "application/json")
|
w.Header().Set("Content-Type", "application/json")
|
||||||
w.WriteHeader(status)
|
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 {
|
func newTestHandler(t *testing.T, upstreamURL string, accounts []config.Account, retries int) *Handler {
|
||||||
t.Helper()
|
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)
|
b := NewBalancer(accounts, 50*time.Millisecond)
|
||||||
return NewHandler(b, upstreamURL, retries, log.New("debug"))
|
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
|
package proxy
|
||||||
|
|
||||||
// retry.go holds the failover helpers used by handler.go. The retry loop itself
|
import (
|
||||||
// lives inside Handler.proxyWithFailover (see handler.go) because it needs
|
"net/http"
|
||||||
// tight control over the moment a response starts streaming vs. is rejected
|
"strings"
|
||||||
// pre-stream. This file documents the invariants the loop must maintain.
|
)
|
||||||
|
|
||||||
// Streaming invariant
|
// isRetriable reports whether an upstream status code should trigger failover
|
||||||
// ===================
|
// to the next account before any bytes are streamed to the client.
|
||||||
// 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.
|
|
||||||
//
|
//
|
||||||
// Pre-stream failover
|
// 429 — rate limit / usage limit (retryable cooldown)
|
||||||
// -------------------
|
// 402 — payment required; treated as retryable when the body text matches a
|
||||||
// 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).
|
|
||||||
//
|
//
|
||||||
// Once copyResponse has called WriteHeader, no further retries are possible
|
// usage-window pattern ("usage limit", "limit reached", "exhausted",
|
||||||
// for this request.
|
// "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