feat: ollama-proxy — multi-account Ollama Cloud reverse proxy
Go reverse proxy for ollama.com that balances requests across multiple API keys with round-robin and failover on 429/5xx. Exposes both native Ollama API (/api/*) and OpenAI-compatible (/v1/*) passthrough. - Round-robin balancer with per-account cooldown (60s default) - Pre-stream failover: 429 → cooldown + next account; 5xx → next account - Streaming invariant: once 2xx starts streaming, no account switch - SSE (text/event-stream) and NDJSON passthrough with http.Flusher - CLI: accounts add/list/remove/set-base-url, serve, version - Accounts stored in ~/.config/ollama-proxy/accounts.json (chmod 0600) - systemd unit (User=dueattendant149, 127.0.0.1:11435, Restart=always) - 43 tests (unit + integration with httptest upstream) - opencode integration: custom provider 'ocp' with explicit model list - Requires NO_PROXY=127.0.0.1,localhost when HTTP_PROXY is set
This commit is contained in:
Atte149
commit
3963eede70
23 changed files with 2534 additions and 0 deletions
4
.gitignore
vendored
Normal file
4
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,4 @@
|
||||||
|
ollama-proxy
|
||||||
|
*.test
|
||||||
|
*.out
|
||||||
|
/.scratch/
|
||||||
91
AGENTS.md
Normal file
91
AGENTS.md
Normal file
|
|
@ -0,0 +1,91 @@
|
||||||
|
# AGENTS.md — ollama-proxy
|
||||||
|
|
||||||
|
## Контекст
|
||||||
|
|
||||||
|
`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`.
|
||||||
|
|
||||||
|
## Стек
|
||||||
|
|
||||||
|
- Go 1.22 (используется `log/slog`, `net/http`, `net/http/httputil`)
|
||||||
|
- Единственная внешняя зависимость: `golang.org/x/term` (prompt без echo в `accounts add`)
|
||||||
|
- Go 1.18 из apt НЕ подходит — нужен `~/.local/go/bin/go` (Go 1.22.10). Перед сборкой: `export PATH="$HOME/.local/go/bin:$PATH"`
|
||||||
|
- Сборка: `make build` → `./ollama-proxy`
|
||||||
|
- Тесты: `go test ./...` (43 теста, включая integration-тесты с `httptest`)
|
||||||
|
- Линт: `go vet ./...`, `gofmt -l .`
|
||||||
|
|
||||||
|
## Структура
|
||||||
|
|
||||||
|
```
|
||||||
|
main.go — точка входа, делегирует в cli.Run
|
||||||
|
internal/config/
|
||||||
|
accounts.go — Account, AccountsFile, JSON load/save (chmod 0600, atomic)
|
||||||
|
server.go — ServerConfig (addr, base_url, cooldown, retries, log_level)
|
||||||
|
internal/log/logger.go — slog JSON logger с configurable level
|
||||||
|
internal/proxy/
|
||||||
|
balancer.go — round-robin + per-account cooldown (RWMutex)
|
||||||
|
handler.go — HTTP handler: reverse-proxy, Authorization подмена, SSE flush
|
||||||
|
retry.go — документация стриминг-инварианта
|
||||||
|
internal/cli/
|
||||||
|
root.go — dispatch подкоманд (serve / accounts / version)
|
||||||
|
accounts.go — add / list / remove / set-base-url
|
||||||
|
serve.go — запуск HTTP-сервера с signal handling
|
||||||
|
systemd/ollama-proxy.service — systemd unit
|
||||||
|
```
|
||||||
|
|
||||||
|
## Конфигурация
|
||||||
|
|
||||||
|
- Аккаунты: `~/.config/ollama-proxy/accounts.json` (XDG_CONFIG_HOME, chmod 0600, atomic write)
|
||||||
|
- Флаги `serve`: `--addr`, `--base-url`, `--cooldown`, `--retries`, `--log-level`
|
||||||
|
- Env vars: `OLLAMA_PROXY_ADDR`, `OLLAMA_PROXY_BASE_URL`, `OLLAMA_PROXY_COOLDOWN`, `OLLAMA_PROXY_RETRIES`, `OLLAMA_PROXY_LOG_LEVEL`
|
||||||
|
|
||||||
|
## Типичные задачи
|
||||||
|
|
||||||
|
- **Добавить аккаунт:** `ollama-proxy accounts add --name <alias>` (промпт) или `... add <key> --name <alias>`
|
||||||
|
- **Перезапустить после правки аккаунтов:** `sudo systemctl restart ollama-proxy` (прокси читает accounts.json при старте; runtime-обновления не поддерживаются)
|
||||||
|
- **Посмотреть логи:** `journalctl -u ollama-proxy -f` — JSON-структурированные (account, account_id, path, status, latency_ms, stream)
|
||||||
|
- **Проверить, что прокси жив:** `curl http://127.0.0.1:11435/api/version` → должен вернуть JSON от ollama.com
|
||||||
|
- **Проверить список моделей:** `curl http://127.0.0.1:11435/v1/models | jq '.data | length'` (~36 моделей)
|
||||||
|
|
||||||
|
## Интеграция с opencode
|
||||||
|
|
||||||
|
В `~/.config/opencode/opencode.json` в блоке `provider`:
|
||||||
|
```json
|
||||||
|
"ocp": {
|
||||||
|
"npm": "@ai-sdk/openai-compatible",
|
||||||
|
"name": "Ollama Cloud (pool)",
|
||||||
|
"options": {
|
||||||
|
"baseURL": "http://127.0.0.1:11435/v1",
|
||||||
|
"apiKey": "proxy"
|
||||||
|
},
|
||||||
|
"models": {
|
||||||
|
"gpt-oss:20b": { "name": "gpt-oss:20b" },
|
||||||
|
"...": { "name": "..." }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
Модель выбирается через `/models` в TUI. Пул не сделан дефолтной моделью.
|
||||||
|
|
||||||
|
**ВАЖНО: NO_PROXY.** Если в окружении установлен `HTTP_PROXY`/`HTTPS_PROXY` (например `http://127.0.0.1:2080`), opencode (Bun runtime) направит запросы к прокси через HTTP-прокси → 502 Bad Gateway. Нужно явно исключить localhost:
|
||||||
|
```sh
|
||||||
|
# В ~/.config/environment.d/no-proxy-local.conf (systemd/desktop sessions):
|
||||||
|
NO_PROXY=127.0.0.1,localhost,::1
|
||||||
|
no_proxy=127.0.0.1,localhost,::1
|
||||||
|
```
|
||||||
|
Или запускать opencode с `NO_PROXY=127.0.0.1,localhost opencode ...`.
|
||||||
|
|
||||||
|
**Модели в конфиге:** opencode НЕ умеет динамически резолвить модели через `/v1/models` для `@ai-sdk/openai-compatible` — `models: {}` даёт `ProviderModelNotFoundError`. Нужно явно перечислить модели. Обновить список: `curl -sS http://127.0.0.1:11435/v1/models | jq -r '.data[].id'` → вставить в `models` блок.
|
||||||
|
|
||||||
|
## WARN
|
||||||
|
|
||||||
|
- **Никогда** не запускай прокси на `0.0.0.0` без auth — он не имеет аутентификации, любой может использовать твои ключи
|
||||||
|
- `accounts.json` содержит ключи в plaintext — файл chmod 0600, не коммить в git
|
||||||
|
- После правки `accounts.json` нужен `sudo systemctl restart ollama-proxy` (прокси не перечитывает файл в рантайме)
|
||||||
|
- Если локальная `ollama` запущена на 11434 — это НЕ конфликтует (прокси на 11435), но `OLLAMA_HOST=127.0.0.1:11435` нужно явно указывать
|
||||||
|
- Стриминг-инвариант: после первого flush-а 2xx ответа переключение аккаунта НЕ происходит. Mid-stream ошибки → truncated response клиенту. Это by design (нельзя дублировать вывод)
|
||||||
|
- 5xx НЕ вызывает cooldown аккаунта (только 429). 5xx → failover на следующий аккаунт, но без penalty
|
||||||
|
|
||||||
|
## План отката
|
||||||
|
|
||||||
|
- `sudo systemctl stop ollama-proxy && sudo systemctl disable ollama-proxy`
|
||||||
|
- `sudo rm /usr/local/bin/ollama-proxy /etc/systemd/system/ollama-proxy.service && sudo systemctl daemon-reload`
|
||||||
|
- Убрать блок `ocp` из `~/.config/opencode/opencode.json`
|
||||||
28
Makefile
Normal file
28
Makefile
Normal file
|
|
@ -0,0 +1,28 @@
|
||||||
|
GO=go
|
||||||
|
|
||||||
|
.PHONY: build install uninstall vet fmt test clean
|
||||||
|
|
||||||
|
build:
|
||||||
|
$(GO) build -o ollama-proxy .
|
||||||
|
|
||||||
|
install: build
|
||||||
|
install -m 755 ollama-proxy /usr/local/bin/ollama-proxy
|
||||||
|
install -m 644 systemd/ollama-proxy.service /etc/systemd/system/ollama-proxy.service
|
||||||
|
systemctl daemon-reload
|
||||||
|
|
||||||
|
uninstall:
|
||||||
|
-systemctl disable --now ollama-proxy
|
||||||
|
rm -f /usr/local/bin/ollama-proxy /etc/systemd/system/ollama-proxy.service
|
||||||
|
systemctl daemon-reload
|
||||||
|
|
||||||
|
vet:
|
||||||
|
$(GO) vet ./...
|
||||||
|
|
||||||
|
fmt:
|
||||||
|
$(GO)fmt -l .
|
||||||
|
|
||||||
|
test:
|
||||||
|
$(GO) test ./...
|
||||||
|
|
||||||
|
clean:
|
||||||
|
rm -f ollama-proxy
|
||||||
140
README.md
Normal file
140
README.md
Normal file
|
|
@ -0,0 +1,140 @@
|
||||||
|
# ollama-proxy
|
||||||
|
|
||||||
|
Multi-account reverse proxy for [Ollama Cloud](https://ollama.com). Round-robin load balancing across multiple API keys with automatic failover on 429 (rate limit) and 5xx errors. Streaming (SSE + NDJSON) is passed through without buffering.
|
||||||
|
|
||||||
|
Listens on `127.0.0.1:11435` by default; exposes both the native Ollama API (`/api/*`) and the OpenAI-compatible API (`/v1/*`) that `ollama.com` serves natively. Designed to sit in front of OpenCode, `ollama` CLI, or any OpenAI-compatible client.
|
||||||
|
|
||||||
|
## Build
|
||||||
|
|
||||||
|
```sh
|
||||||
|
make build
|
||||||
|
```
|
||||||
|
|
||||||
|
Produces `./ollama-proxy`. Requires Go 1.21+ (uses `log/slog`).
|
||||||
|
|
||||||
|
## Install
|
||||||
|
|
||||||
|
```sh
|
||||||
|
sudo make install
|
||||||
|
sudo systemctl enable --now ollama-proxy
|
||||||
|
```
|
||||||
|
|
||||||
|
This installs the binary to `/usr/local/bin/ollama-proxy` and the systemd unit to `/etc/systemd/system/ollama-proxy.service`, then reloads systemd.
|
||||||
|
|
||||||
|
## Configure accounts
|
||||||
|
|
||||||
|
Accounts are stored in `~/.config/ollama-proxy/accounts.json` (chmod 0600). Manage them via the CLI:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
# Add an account (prompts for the key with echo disabled if not given on the CLI)
|
||||||
|
ollama-proxy accounts add --name work
|
||||||
|
ollama-proxy accounts add sk-xxxx-... --name personal
|
||||||
|
|
||||||
|
# List accounts (keys are masked)
|
||||||
|
ollama-proxy accounts list
|
||||||
|
|
||||||
|
# Remove by name or id prefix
|
||||||
|
ollama-proxy accounts remove work
|
||||||
|
ollama-proxy accounts remove 3e94140f
|
||||||
|
|
||||||
|
# Override the default upstream (rarely needed; default is https://ollama.com)
|
||||||
|
ollama-proxy accounts set-base-url https://ollama.com
|
||||||
|
```
|
||||||
|
|
||||||
|
API keys are created at <https://ollama.com/settings/keys>.
|
||||||
|
|
||||||
|
## Run as a service
|
||||||
|
|
||||||
|
The systemd unit runs the proxy on `127.0.0.1:11435` and restarts on crash. Logs go to journald:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
journalctl -u ollama-proxy -f
|
||||||
|
```
|
||||||
|
|
||||||
|
Configuration flags can be overridden via environment variables in the unit (or in a drop-in):
|
||||||
|
|
||||||
|
| Variable | Default | Description |
|
||||||
|
| ------------------------- | ------------------- | --------------------------------- |
|
||||||
|
| `OLLAMA_PROXY_ADDR` | `127.0.0.1:11435` | listen address |
|
||||||
|
| `OLLAMA_PROXY_BASE_URL` | `https://ollama.com`| upstream root |
|
||||||
|
| `OLLAMA_PROXY_COOLDOWN` | `60s` | per-account 429 cooldown window |
|
||||||
|
| `OLLAMA_PROXY_RETRIES` | `3` | max failover attempts |
|
||||||
|
| `OLLAMA_PROXY_LOG_LEVEL` | `info` | `debug` \| `info` \| `warn` \| `error` |
|
||||||
|
|
||||||
|
## Use with OpenCode
|
||||||
|
|
||||||
|
Add a custom provider to `~/.config/opencode/opencode.json`:
|
||||||
|
|
||||||
|
```json
|
||||||
|
"ocp": {
|
||||||
|
"npm": "@ai-sdk/openai-compatible",
|
||||||
|
"name": "Ollama Cloud (pool)",
|
||||||
|
"options": {
|
||||||
|
"baseURL": "http://127.0.0.1:11435/v1",
|
||||||
|
"apiKey": "proxy"
|
||||||
|
},
|
||||||
|
"models": {
|
||||||
|
"gpt-oss:20b": { "name": "gpt-oss:20b" },
|
||||||
|
"gpt-oss:120b": { "name": "gpt-oss:120b" },
|
||||||
|
"kimi-k2.6": { "name": "kimi-k2.6" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
`apiKey: "proxy"` is a placeholder — the proxy ignores the client's `Authorization` header and substitutes the selected account's key. **You must list models explicitly** — opencode does not auto-fetch `/v1/models` for `@ai-sdk/openai-compatible` providers. To get the full list:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
curl -sS http://127.0.0.1:11435/v1/models | jq -r '.data[].id'
|
||||||
|
```
|
||||||
|
|
||||||
|
### NO_PROXY (important)
|
||||||
|
|
||||||
|
If your environment has `HTTP_PROXY`/`HTTPS_PROXY` set (e.g. `http://127.0.0.1:2080`), opencode's Bun runtime will route requests to the proxy *through* the HTTP proxy, resulting in `502 Bad Gateway`. Exclude localhost from proxying:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
# ~/.config/environment.d/no-proxy-local.conf
|
||||||
|
NO_PROXY=127.0.0.1,localhost,::1
|
||||||
|
no_proxy=127.0.0.1,localhost,::1
|
||||||
|
```
|
||||||
|
|
||||||
|
Or launch opencode with `NO_PROXY=127.0.0.1,localhost opencode`.
|
||||||
|
|
||||||
|
Run `/models` in the OpenCode TUI to pick a model from the pool, then chat as usual.
|
||||||
|
|
||||||
|
## Use with the `ollama` CLI
|
||||||
|
|
||||||
|
```sh
|
||||||
|
OLLAMA_HOST=127.0.0.1:11435 ollama run gpt-oss:20b
|
||||||
|
```
|
||||||
|
|
||||||
|
## Endpoints
|
||||||
|
|
||||||
|
The proxy forwards (passthrough) any of these paths to `ollama.com`:
|
||||||
|
|
||||||
|
- Native Ollama: `/api/chat`, `/api/generate`, `/api/tags`, `/api/show`, `/api/ps`, `/api/version`, `/api/delete`
|
||||||
|
- OpenAI-compatible: `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/v1/embeddings`, `/v1/files/*`
|
||||||
|
|
||||||
|
Anything else returns 404.
|
||||||
|
|
||||||
|
## How failover works
|
||||||
|
|
||||||
|
1. For each request, the balancer picks the next account round-robin.
|
||||||
|
2. If the upstream returns `429` before any bytes are sent to the client, the account is put in cooldown (default 60s) and the next account is tried, up to `--retries` attempts.
|
||||||
|
3. If the upstream returns `5xx`, the next account is tried **without** cooldown (5xx may be transient).
|
||||||
|
4. Once a 2xx response has started streaming to the client, the request is committed: mid-stream upstream errors are surfaced as-is (a truncated response). The proxy never restarts a streamed request on a different account, because the client would see duplicated output.
|
||||||
|
5. If every account is in cooldown, the client receives `429 {"error":{"message":"all accounts rate-limited"}}`.
|
||||||
|
6. If every account fails for non-429 reasons (network, 5xx), the client receives `502 {"error":{"message":"all accounts failed: …"}}`.
|
||||||
|
|
||||||
|
## Tests
|
||||||
|
|
||||||
|
```sh
|
||||||
|
go test ./... # unit + integration
|
||||||
|
go vet ./...
|
||||||
|
gofmt -l .
|
||||||
|
```
|
||||||
|
|
||||||
|
The integration tests build the binary and exercise the full HTTP stack against an `httptest` upstream.
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT.
|
||||||
7
go.mod
Normal file
7
go.mod
Normal file
|
|
@ -0,0 +1,7 @@
|
||||||
|
module github.com/Atte149/ollama-proxy
|
||||||
|
|
||||||
|
go 1.22
|
||||||
|
|
||||||
|
require golang.org/x/term v0.21.0
|
||||||
|
|
||||||
|
require golang.org/x/sys v0.21.0 // indirect
|
||||||
4
go.sum
Normal file
4
go.sum
Normal file
|
|
@ -0,0 +1,4 @@
|
||||||
|
golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
|
||||||
|
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||||
|
golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA=
|
||||||
|
golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
|
||||||
223
internal/cli/accounts.go
Normal file
223
internal/cli/accounts.go
Normal file
|
|
@ -0,0 +1,223 @@
|
||||||
|
package cli
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
|
"golang.org/x/term"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RunAccounts dispatches to the accounts subcommand.
|
||||||
|
func RunAccounts(args []string) int {
|
||||||
|
if len(args) < 1 {
|
||||||
|
printAccountsUsage()
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
switch args[0] {
|
||||||
|
case "add":
|
||||||
|
return accountsAdd(args[1:])
|
||||||
|
case "list":
|
||||||
|
return accountsList(args[1:])
|
||||||
|
case "remove", "rm":
|
||||||
|
return accountsRemove(args[1:])
|
||||||
|
case "set-base-url":
|
||||||
|
return accountsSetBaseURL(args[1:])
|
||||||
|
case "-h", "--help", "help":
|
||||||
|
printAccountsUsage()
|
||||||
|
return 0
|
||||||
|
default:
|
||||||
|
fmt.Fprintf(os.Stderr, "unknown accounts subcommand: %s\n\n", args[0])
|
||||||
|
printAccountsUsage()
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func printAccountsUsage() {
|
||||||
|
fmt.Print(`Usage:
|
||||||
|
ollama-proxy accounts add [API_KEY] [--name <alias>]
|
||||||
|
ollama-proxy accounts list
|
||||||
|
ollama-proxy accounts remove <id|name>
|
||||||
|
ollama-proxy accounts set-base-url <url>
|
||||||
|
`)
|
||||||
|
}
|
||||||
|
|
||||||
|
// accountsAdd implements `accounts add [API_KEY] [--name <alias>]`.
|
||||||
|
func accountsAdd(args []string) int {
|
||||||
|
var apiKey, name string
|
||||||
|
for i := 0; i < len(args); i++ {
|
||||||
|
switch args[i] {
|
||||||
|
case "--name":
|
||||||
|
if i+1 >= len(args) {
|
||||||
|
fmt.Fprintln(os.Stderr, "--name requires an argument")
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
name = args[i+1]
|
||||||
|
i++
|
||||||
|
case "-h", "--help":
|
||||||
|
fmt.Println("usage: accounts add [API_KEY] [--name <alias>]")
|
||||||
|
return 0
|
||||||
|
default:
|
||||||
|
if apiKey == "" {
|
||||||
|
apiKey = args[i]
|
||||||
|
} else {
|
||||||
|
fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", args[i])
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if name == "" {
|
||||||
|
name = promptString("Account name (alias)", "acct"+config.NewID()[:4])
|
||||||
|
}
|
||||||
|
if err := config.ValidateName(name); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "invalid name:", err)
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
|
||||||
|
if apiKey == "" {
|
||||||
|
fmt.Print("Enter Ollama Cloud API key (input hidden): ")
|
||||||
|
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||||
|
fmt.Println()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "read key:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
apiKey = strings.TrimSpace(string(b))
|
||||||
|
}
|
||||||
|
if apiKey == "" {
|
||||||
|
fmt.Fprintln(os.Stderr, "empty API key, nothing to add")
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := config.EnsureAccountsDir(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "create config dir:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
af, err := config.LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "load accounts:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
if af.FindByName(name) != nil {
|
||||||
|
fmt.Fprintf(os.Stderr, "an account named %q already exists\n", name)
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
acct := config.Account{
|
||||||
|
ID: config.NewID(),
|
||||||
|
Name: name,
|
||||||
|
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))
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// accountsList prints a table of accounts.
|
||||||
|
func accountsList(args []string) int {
|
||||||
|
af, err := config.LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "load accounts:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
if len(af.Accounts) == 0 {
|
||||||
|
fmt.Println("no accounts configured. Run: ollama-proxy accounts add")
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
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))
|
||||||
|
for _, a := range af.Accounts {
|
||||||
|
id := a.ID
|
||||||
|
if len(id) > 8 {
|
||||||
|
id = id[:8]
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
if af.BaseURL != "" {
|
||||||
|
fmt.Fprintf(w, "\nDefault upstream: %s\n", af.BaseURL)
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// accountsRemove removes an account by id or name.
|
||||||
|
func accountsRemove(args []string) int {
|
||||||
|
if len(args) < 1 || args[0] == "-h" || args[0] == "--help" {
|
||||||
|
fmt.Println("usage: accounts remove <id|name>")
|
||||||
|
if len(args) < 1 {
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
ident := args[0]
|
||||||
|
af, err := config.LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "load accounts:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
if err := af.Remove(ident); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
if err := af.Save(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "save:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
fmt.Printf("removed account %s\n", ident)
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// accountsSetBaseURL sets the default upstream URL stored in accounts.json.
|
||||||
|
func accountsSetBaseURL(args []string) int {
|
||||||
|
if len(args) < 1 {
|
||||||
|
fmt.Println("usage: accounts set-base-url <url>")
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
url := args[0]
|
||||||
|
if err := config.EnsureAccountsDir(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "create config dir:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
af, err := config.LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "load accounts:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
af.BaseURL = url
|
||||||
|
if err := af.Save(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "save:", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
fmt.Printf("default upstream set to %s\n", url)
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// promptString reads a line from stdin, returning def when the user just hits
|
||||||
|
// enter. Used for the optional account name prompt.
|
||||||
|
func promptString(label, def string) string {
|
||||||
|
fmt.Printf("%s [%s]: ", label, def)
|
||||||
|
r := bufio.NewReader(os.Stdin)
|
||||||
|
line, _ := r.ReadString('\n')
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if line == "" {
|
||||||
|
return def
|
||||||
|
}
|
||||||
|
return line
|
||||||
|
}
|
||||||
|
|
||||||
|
// maskKey mirrors config.accounts.maskKey but is exposed here for display.
|
||||||
|
func maskKey(key string) string {
|
||||||
|
if len(key) <= 4 {
|
||||||
|
return "****"
|
||||||
|
}
|
||||||
|
return "****" + key[len(key)-4:]
|
||||||
|
}
|
||||||
241
internal/cli/integration_test.go
Normal file
241
internal/cli/integration_test.go
Normal file
|
|
@ -0,0 +1,241 @@
|
||||||
|
package cli
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// findBinary locates the ollama-proxy executable built by `go test`. The test
|
||||||
|
// binary is in the package dir; the project binary is one level up. We try both
|
||||||
|
// and fall back to building one on the fly.
|
||||||
|
func findBinary(t *testing.T) string {
|
||||||
|
t.Helper()
|
||||||
|
candidates := []string{
|
||||||
|
filepath.Join("..", "..", "ollama-proxy"),
|
||||||
|
"ollama-proxy",
|
||||||
|
}
|
||||||
|
for _, c := range candidates {
|
||||||
|
if abs, err := filepath.Abs(c); err == nil {
|
||||||
|
if fi, err := os.Stat(abs); err == nil && !fi.IsDir() {
|
||||||
|
return abs
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Build it now into a temp path.
|
||||||
|
out := filepath.Join(t.TempDir(), "ollama-proxy")
|
||||||
|
cmd := exec.Command("go", "build", "-o", out, ".")
|
||||||
|
cmd.Dir = filepath.Join("..", "..")
|
||||||
|
if out, err := cmd.CombinedOutput(); err != nil {
|
||||||
|
t.Fatalf("go build: %v\n%s", err, out)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// freePort picks an unused TCP port for the proxy to listen on.
|
||||||
|
func freePort(t *testing.T) string {
|
||||||
|
t.Helper()
|
||||||
|
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer l.Close()
|
||||||
|
return l.Addr().String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// startProxy starts the binary as a subprocess and returns a cancel function
|
||||||
|
// that stops it.
|
||||||
|
func startProxy(t *testing.T, bin, xdgDir, addr, upstream string) func() {
|
||||||
|
t.Helper()
|
||||||
|
cmd := exec.Command(bin, "serve",
|
||||||
|
"--addr", addr,
|
||||||
|
"--base-url", upstream,
|
||||||
|
"--cooldown", "1s",
|
||||||
|
"--retries", "3",
|
||||||
|
"--log-level", "debug",
|
||||||
|
)
|
||||||
|
cmd.Env = append(os.Environ(),
|
||||||
|
"XDG_CONFIG_HOME="+xdgDir,
|
||||||
|
"OLLAMA_PROXY_LOG_LEVEL=debug",
|
||||||
|
)
|
||||||
|
var stdout, stderr bytes.Buffer
|
||||||
|
cmd.Stdout = &stdout
|
||||||
|
cmd.Stderr = &stderr
|
||||||
|
if err := cmd.Start(); err != nil {
|
||||||
|
t.Fatalf("start proxy: %v", err)
|
||||||
|
}
|
||||||
|
// Wait for the port to accept connections.
|
||||||
|
deadline := time.Now().Add(3 * time.Second)
|
||||||
|
for time.Now().Before(deadline) {
|
||||||
|
c, err := net.Dial("tcp", addr)
|
||||||
|
if err == nil {
|
||||||
|
c.Close()
|
||||||
|
break
|
||||||
|
}
|
||||||
|
time.Sleep(20 * time.Millisecond)
|
||||||
|
}
|
||||||
|
return func() {
|
||||||
|
_ = cmd.Process.Signal(os.Interrupt)
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() { _ = cmd.Wait(); close(done) }()
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
_ = cmd.Process.Kill()
|
||||||
|
<-done
|
||||||
|
}
|
||||||
|
t.Logf("proxy stdout: %s", stdout.String())
|
||||||
|
t.Logf("proxy stderr: %s", stderr.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestEndToEnd_ProxyForwardsToUpstream builds the binary and verifies the full
|
||||||
|
// stack: account JSON → balancer → handler → real HTTP server → upstream.
|
||||||
|
func TestEndToEnd_ProxyForwardsToUpstream(t *testing.T) {
|
||||||
|
// Upstream that echoes the auth and path.
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
_, _ = fmt.Fprintf(w, `{"path":%q,"auth":%q}`, r.URL.Path, r.Header.Get("Authorization"))
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
bin := findBinary(t)
|
||||||
|
xdg := t.TempDir()
|
||||||
|
// Write accounts.json directly via the binary's `accounts add`.
|
||||||
|
for _, a := range []struct{ key, name string }{
|
||||||
|
{"sk-keyAAAAAAAA", "a1"},
|
||||||
|
{"sk-keyBBBBBBBB", "a2"},
|
||||||
|
} {
|
||||||
|
cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
|
||||||
|
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||||
|
if out, err := cmd.CombinedOutput(); err != nil {
|
||||||
|
t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
addr := freePort(t)
|
||||||
|
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
resp, err := http.Get("http://" + addr + "/api/version")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GET /api/version: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
if resp.StatusCode != 200 {
|
||||||
|
t.Fatalf("status = %d, body = %s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), "sk-key") {
|
||||||
|
t.Errorf("body = %s, want account key forwarded", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), "/api/version") {
|
||||||
|
t.Errorf("body = %s, want path echoed", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestEndToEnd_SSEStreaming verifies SSE passthrough through the real binary.
|
||||||
|
func TestEndToEnd_SSEStreaming(t *testing.T) {
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Header().Set("Transfer-Encoding", "chunked")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
flusher, _ := w.(http.Flusher)
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
_, _ = io.WriteString(w, "data: {\"i\":"+string(rune('0'+i))+"}\n\n")
|
||||||
|
if flusher != nil {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||||||
|
if flusher != nil {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
bin := findBinary(t)
|
||||||
|
xdg := t.TempDir()
|
||||||
|
cmd := exec.Command(bin, "accounts", "add", "sk-testkey", "--name", "a1")
|
||||||
|
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||||
|
if out, err := cmd.CombinedOutput(); err != nil {
|
||||||
|
t.Fatalf("accounts add: %v\n%s", err, out)
|
||||||
|
}
|
||||||
|
|
||||||
|
addr := freePort(t)
|
||||||
|
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
resp, err := http.Post("http://"+addr+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("POST: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
if !strings.Contains(string(body), "[DONE]") {
|
||||||
|
t.Errorf("body = %s, want [DONE]", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), `"i":0`) || !strings.Contains(string(body), `"i":1`) {
|
||||||
|
t.Errorf("body = %s, missing chunks", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestEndToEnd_FailoverOn429 verifies the binary fails over between accounts.
|
||||||
|
func TestEndToEnd_FailoverOn429(t *testing.T) {
|
||||||
|
// Upstream that 429s for the first account key and 200s for the second.
|
||||||
|
var firstSeen string
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
auth := r.Header.Get("Authorization")
|
||||||
|
if firstSeen == "" {
|
||||||
|
firstSeen = auth
|
||||||
|
}
|
||||||
|
if auth == "Bearer sk-bad" {
|
||||||
|
http.Error(w, "rate limited", 429)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
_, _ = io.WriteString(w, `{"ok":true}`)
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
bin := findBinary(t)
|
||||||
|
xdg := t.TempDir()
|
||||||
|
for _, a := range []struct{ key, name string }{
|
||||||
|
{"sk-bad", "bad"},
|
||||||
|
{"sk-good", "good"},
|
||||||
|
} {
|
||||||
|
cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
|
||||||
|
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||||
|
if out, err := cmd.CombinedOutput(); err != nil {
|
||||||
|
t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
addr := freePort(t)
|
||||||
|
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
resp, err := http.Get("http://" + addr + "/api/version")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GET: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
if resp.StatusCode != 200 {
|
||||||
|
t.Fatalf("status = %d, body = %s (want 200 after failover)", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), `"ok":true`) {
|
||||||
|
t.Errorf("body = %s, want good response", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
7
internal/cli/now.go
Normal file
7
internal/cli/now.go
Normal file
|
|
@ -0,0 +1,7 @@
|
||||||
|
package cli
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
// nowFn returns the current time. It is an indirection so tests can stub it
|
||||||
|
// via package-internal reassignment (see accounts_test.go).
|
||||||
|
var nowFn = func() time.Time { return time.Now().UTC() }
|
||||||
59
internal/cli/root.go
Normal file
59
internal/cli/root.go
Normal file
|
|
@ -0,0 +1,59 @@
|
||||||
|
// Package cli implements the ollama-proxy command-line interface: account
|
||||||
|
// management subcommands and the `serve` subcommand that starts the proxy.
|
||||||
|
package cli
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Run dispatches to a subcommand based on args[1]. It returns the process exit
|
||||||
|
// code, or an error that the caller should print before exiting 1.
|
||||||
|
func Run(args []string) int {
|
||||||
|
if len(args) < 2 {
|
||||||
|
printRootUsage()
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
switch args[1] {
|
||||||
|
case "serve":
|
||||||
|
return RunServe(args[2:])
|
||||||
|
case "accounts":
|
||||||
|
return RunAccounts(args[2:])
|
||||||
|
case "version":
|
||||||
|
fmt.Println("ollama-proxy", Version)
|
||||||
|
return 0
|
||||||
|
case "-h", "--help", "help":
|
||||||
|
printRootUsage()
|
||||||
|
return 0
|
||||||
|
default:
|
||||||
|
fmt.Fprintf(os.Stderr, "unknown subcommand: %s\n\n", args[1])
|
||||||
|
printRootUsage()
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Version is overwritten at build time via -ldflags "-X .../cli.Version=...".
|
||||||
|
var Version = "dev"
|
||||||
|
|
||||||
|
func printRootUsage() {
|
||||||
|
fmt.Print(`ollama-proxy — multi-account Ollama Cloud reverse proxy
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
ollama-proxy serve [flags] Start the proxy server
|
||||||
|
ollama-proxy accounts <subcommand> Manage Ollama Cloud API keys
|
||||||
|
ollama-proxy version Print version
|
||||||
|
|
||||||
|
Accounts subcommands:
|
||||||
|
accounts add [API_KEY] [--name <alias>] Add an account (prompts for key if not given)
|
||||||
|
accounts list List configured accounts
|
||||||
|
accounts remove <id|name> Remove an account
|
||||||
|
accounts set-base-url <url> Set the default upstream URL
|
||||||
|
|
||||||
|
Environment variables (serve):
|
||||||
|
OLLAMA_PROXY_ADDR listen address (default 127.0.0.1:11435)
|
||||||
|
OLLAMA_PROXY_BASE_URL upstream root (default https://ollama.com)
|
||||||
|
OLLAMA_PROXY_COOLDOWN 429 cooldown window (default 60s)
|
||||||
|
OLLAMA_PROXY_RETRIES max failover attempts (default 3)
|
||||||
|
OLLAMA_PROXY_LOG_LEVEL debug|info|warn|error (default info)
|
||||||
|
`)
|
||||||
|
}
|
||||||
109
internal/cli/serve.go
Normal file
109
internal/cli/serve.go
Normal file
|
|
@ -0,0 +1,109 @@
|
||||||
|
package cli
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"os/signal"
|
||||||
|
"syscall"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/log"
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/proxy"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RunServe starts the proxy HTTP server.
|
||||||
|
func RunServe(args []string) int {
|
||||||
|
fs := flag.NewFlagSet("serve", flag.ContinueOnError)
|
||||||
|
addr := fs.String("addr", "", "listen address (default 127.0.0.1:11435)")
|
||||||
|
baseURL := fs.String("base-url", "", "upstream root (default https://ollama.com)")
|
||||||
|
cooldown := fs.Duration("cooldown", 0, "per-account 429 cooldown window")
|
||||||
|
retries := fs.Int("retries", 0, "max failover attempts")
|
||||||
|
logLevel := fs.String("log-level", "", "debug|info|warn|error")
|
||||||
|
if err := fs.Parse(args); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, err)
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg := config.DefaultServerConfig()
|
||||||
|
if *addr != "" {
|
||||||
|
cfg.Addr = *addr
|
||||||
|
}
|
||||||
|
if *baseURL != "" {
|
||||||
|
cfg.BaseURL = *baseURL
|
||||||
|
}
|
||||||
|
if *cooldown > 0 {
|
||||||
|
cfg.Cooldown = *cooldown
|
||||||
|
}
|
||||||
|
if *retries > 0 {
|
||||||
|
cfg.Retries = *retries
|
||||||
|
}
|
||||||
|
if *logLevel != "" {
|
||||||
|
cfg.LogLevel = *logLevel
|
||||||
|
}
|
||||||
|
cfg.ApplyEnv()
|
||||||
|
if err := cfg.Validate(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "config error:", err)
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
|
||||||
|
logger := log.New(cfg.LogLevel)
|
||||||
|
|
||||||
|
af, err := config.LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
logger.Error("load accounts failed", "err", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
if len(af.Accounts) == 0 {
|
||||||
|
logger.Error("no accounts configured; run `ollama-proxy accounts add` first")
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
// Per-file BaseURL override wins over flag/env when set.
|
||||||
|
if af.BaseURL != "" {
|
||||||
|
cfg.BaseURL = af.BaseURL
|
||||||
|
}
|
||||||
|
|
||||||
|
balancer := proxy.NewBalancer(af.Accounts, cfg.Cooldown)
|
||||||
|
handler := proxy.NewHandler(balancer, cfg.BaseURL, cfg.Retries, logger)
|
||||||
|
|
||||||
|
srv := &http.Server{
|
||||||
|
Addr: cfg.Addr,
|
||||||
|
Handler: handler,
|
||||||
|
// No ReadTimeout/WriteTimeout: streaming chat may be idle for long
|
||||||
|
// periods between tokens. The http.Server's default idle timeout (60s)
|
||||||
|
// applies between reads; if that becomes a problem we'll add a per-read
|
||||||
|
// deadline via a wrapped Transport, not a global timeout.
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
logger.Info("ollama-proxy starting",
|
||||||
|
"addr", cfg.Addr, "upstream", cfg.BaseURL,
|
||||||
|
"accounts", balancer.Len(), "cooldown", cfg.Cooldown.String(),
|
||||||
|
"retries", cfg.Retries, "log_level", cfg.LogLevel)
|
||||||
|
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||||
|
logger.Error("server error", "err", err)
|
||||||
|
stop()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
<-ctx.Done()
|
||||||
|
logger.Info("shutdown signal received, draining…")
|
||||||
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||||
|
logger.Error("graceful shutdown failed", "err", err)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
logger.Info("stopped")
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// keep slog referenced even if logger.go ever drops it
|
||||||
|
var _ = slog.LevelInfo
|
||||||
218
internal/config/accounts.go
Normal file
218
internal/config/accounts.go
Normal file
|
|
@ -0,0 +1,218 @@
|
||||||
|
// Package config holds runtime configuration for ollama-proxy: the list of
|
||||||
|
// Ollama Cloud accounts and the server runtime options.
|
||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Account is a single Ollama Cloud 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"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// AccountsFile is the on-disk JSON structure.
|
||||||
|
type AccountsFile struct {
|
||||||
|
BaseURL string `json:"base_url,omitempty"` // default upstream URL
|
||||||
|
Accounts []Account `json:"accounts"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// AccountsPath returns the canonical accounts.json path under XDG_CONFIG_HOME
|
||||||
|
// (or ~/.config when XDG_CONFIG_HOME is unset). The parent directory must exist.
|
||||||
|
func AccountsPath() (string, error) {
|
||||||
|
dir := os.Getenv("XDG_CONFIG_HOME")
|
||||||
|
if dir == "" {
|
||||||
|
home, err := os.UserHomeDir()
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
dir = filepath.Join(home, ".config")
|
||||||
|
}
|
||||||
|
return filepath.Join(dir, "ollama-proxy", "accounts.json"), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// EnsureAccountsDir creates ~/.config/ollama-proxy with 0700 perms if missing.
|
||||||
|
func EnsureAccountsDir() error {
|
||||||
|
p, err := AccountsPath()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
|
||||||
|
return fmt.Errorf("create config dir: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadAccounts reads accounts.json from the default path. Returns an empty
|
||||||
|
// AccountsFile (not an error) when the file does not exist.
|
||||||
|
func LoadAccounts() (*AccountsFile, error) {
|
||||||
|
p, err := AccountsPath()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
data, err := os.ReadFile(p)
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, os.ErrNotExist) {
|
||||||
|
return &AccountsFile{}, nil
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("read accounts: %w", err)
|
||||||
|
}
|
||||||
|
var af AccountsFile
|
||||||
|
if err := json.Unmarshal(data, &af); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse accounts: %w", err)
|
||||||
|
}
|
||||||
|
return &af, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save writes accounts.json atomically with 0600 permissions. The temp file is
|
||||||
|
// written in the same directory and renamed, so a crash mid-write never leaves
|
||||||
|
// a truncated file. The parent directory must already exist (see EnsureAccountsDir).
|
||||||
|
func (af *AccountsFile) Save() error {
|
||||||
|
p, err := AccountsPath()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if af.Accounts == nil {
|
||||||
|
af.Accounts = []Account{}
|
||||||
|
}
|
||||||
|
data, err := json.MarshalIndent(af, "", " ")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal accounts: %w", err)
|
||||||
|
}
|
||||||
|
dir := filepath.Dir(p)
|
||||||
|
tmp, err := os.CreateTemp(dir, ".accounts.json.*")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("create temp: %w", err)
|
||||||
|
}
|
||||||
|
tmpName := tmp.Name()
|
||||||
|
defer func() {
|
||||||
|
if err != nil {
|
||||||
|
_ = os.Remove(tmpName)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
if err := tmp.Chmod(0o600); err != nil {
|
||||||
|
tmp.Close()
|
||||||
|
return fmt.Errorf("chmod temp: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := tmp.Write(data); err != nil {
|
||||||
|
tmp.Close()
|
||||||
|
return fmt.Errorf("write temp: %w", err)
|
||||||
|
}
|
||||||
|
if err := tmp.Close(); err != nil {
|
||||||
|
return fmt.Errorf("close temp: %w", err)
|
||||||
|
}
|
||||||
|
if err := os.Rename(tmpName, p); err != nil {
|
||||||
|
return fmt.Errorf("rename temp: %w", err)
|
||||||
|
}
|
||||||
|
if err := os.Chmod(p, 0o600); err != nil {
|
||||||
|
return fmt.Errorf("chmod final: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewID returns a random 16-byte hex string suitable as a stable account ID.
|
||||||
|
func NewID() string {
|
||||||
|
b := make([]byte, 16)
|
||||||
|
if _, err := rand.Read(b); err != nil {
|
||||||
|
panic("crypto/rand failed: " + err.Error())
|
||||||
|
}
|
||||||
|
return hex.EncodeToString(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
// shortID returns the first 8 chars of an ID, for display only.
|
||||||
|
func shortID(id string) string {
|
||||||
|
if len(id) < 8 {
|
||||||
|
return id
|
||||||
|
}
|
||||||
|
return id[:8]
|
||||||
|
}
|
||||||
|
|
||||||
|
// maskKey renders an API key as "****ABCD" (last 4 chars) for safe display.
|
||||||
|
func maskKey(key string) string {
|
||||||
|
if len(key) <= 4 {
|
||||||
|
return "****"
|
||||||
|
}
|
||||||
|
return "****" + key[len(key)-4:]
|
||||||
|
}
|
||||||
|
|
||||||
|
// FindByName returns the first account matching name (case-sensitive). Returns
|
||||||
|
// nil when no account matches.
|
||||||
|
func (af *AccountsFile) FindByName(name string) *Account {
|
||||||
|
for i := range af.Accounts {
|
||||||
|
if af.Accounts[i].Name == name {
|
||||||
|
return &af.Accounts[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// FindByID returns the account with the given ID prefix (matches short ID too).
|
||||||
|
func (af *AccountsFile) FindByID(idPrefix string) *Account {
|
||||||
|
for i := range af.Accounts {
|
||||||
|
if af.Accounts[i].ID == idPrefix || hasPrefix(af.Accounts[i].ID, idPrefix) {
|
||||||
|
return &af.Accounts[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func hasPrefix(s, prefix string) bool {
|
||||||
|
if len(prefix) > len(s) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return s[:len(prefix)] == prefix
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find returns the account matching the given identifier, which may be either
|
||||||
|
// a name or an ID/ID-prefix. Name match takes precedence.
|
||||||
|
func (af *AccountsFile) Find(ident string) *Account {
|
||||||
|
if a := af.FindByName(ident); a != nil {
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
return af.FindByID(ident)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove deletes the account with the given ID or name. Returns ErrNotFound
|
||||||
|
// when no account matched.
|
||||||
|
func (af *AccountsFile) Remove(ident string) error {
|
||||||
|
idx := -1
|
||||||
|
for i := range af.Accounts {
|
||||||
|
if af.Accounts[i].Name == ident || af.Accounts[i].ID == ident || hasPrefix(af.Accounts[i].ID, ident) {
|
||||||
|
idx = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if idx < 0 {
|
||||||
|
return ErrNotFound
|
||||||
|
}
|
||||||
|
af.Accounts = append(af.Accounts[:idx], af.Accounts[idx+1:]...)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ErrNotFound is returned when an account lookup or removal misses.
|
||||||
|
var ErrNotFound = errors.New("account not found")
|
||||||
|
|
||||||
|
// ValidateName returns an error when name contains characters outside
|
||||||
|
// [A-Za-z0-9_-]. Names must be safe for CLI display and log output.
|
||||||
|
func ValidateName(name string) error {
|
||||||
|
ok, err := regexp.MatchString(`^[A-Za-z0-9_-]+$`, name)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if !ok {
|
||||||
|
return errors.New("name must contain only letters, digits, '_' or '-'")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
265
internal/config/accounts_test.go
Normal file
265
internal/config/accounts_test.go
Normal file
|
|
@ -0,0 +1,265 @@
|
||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// withTempAccountsPath swaps AccountsPath to a temp dir for the duration of t.
|
||||||
|
// Returns the full path to accounts.json inside that dir (dir is created 0700).
|
||||||
|
func withTempAccountsPath(t *testing.T) string {
|
||||||
|
t.Helper()
|
||||||
|
dir := t.TempDir()
|
||||||
|
sub := filepath.Join(dir, "ollama-proxy")
|
||||||
|
if err := os.MkdirAll(sub, 0o700); err != nil {
|
||||||
|
t.Fatalf("mkdir: %v", err)
|
||||||
|
}
|
||||||
|
p := filepath.Join(sub, "accounts.json")
|
||||||
|
t.Setenv("XDG_CONFIG_HOME", dir)
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewID_UniqueAndLength(t *testing.T) {
|
||||||
|
ids := make(map[string]struct{}, 100)
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
id := NewID()
|
||||||
|
if len(id) != 32 {
|
||||||
|
t.Fatalf("NewID len = %d, want 32 (16 bytes hex)", len(id))
|
||||||
|
}
|
||||||
|
ids[id] = struct{}{}
|
||||||
|
}
|
||||||
|
if len(ids) != 100 {
|
||||||
|
t.Fatalf("NewID produced %d unique ids out of 100", len(ids))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMaskKey(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
in, want string
|
||||||
|
}{
|
||||||
|
{"abc", "****"},
|
||||||
|
{"abcd", "****"},
|
||||||
|
{"abcde", "****bcde"},
|
||||||
|
{"sk-1234567890abcdef", "****cdef"},
|
||||||
|
{"", "****"},
|
||||||
|
}
|
||||||
|
for _, c := range cases {
|
||||||
|
if got := maskKey(c.in); got != c.want {
|
||||||
|
t.Errorf("maskKey(%q) = %q, want %q", c.in, got, c.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShortID(t *testing.T) {
|
||||||
|
if got := shortID("0123456789abcdef"); got != "01234567" {
|
||||||
|
t.Errorf("shortID = %q, want 01234567", got)
|
||||||
|
}
|
||||||
|
if got := shortID("abc"); got != "abc" {
|
||||||
|
t.Errorf("shortID = %q, want abc", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestValidateName(t *testing.T) {
|
||||||
|
good := []string{"a", "work", "my-account_1", "ACC-2"}
|
||||||
|
for _, n := range good {
|
||||||
|
if err := ValidateName(n); err != nil {
|
||||||
|
t.Errorf("ValidateName(%q) = %v, want nil", n, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
bad := []string{"", "with space", "dollar$", "dot.in", "slah/s", "café"}
|
||||||
|
for _, n := range bad {
|
||||||
|
if err := ValidateName(n); err == nil {
|
||||||
|
t.Errorf("ValidateName(%q) = nil, want error", n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadAccounts_MissingFile(t *testing.T) {
|
||||||
|
withTempAccountsPath(t)
|
||||||
|
af, err := LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadAccounts on missing file: %v", err)
|
||||||
|
}
|
||||||
|
if af == nil {
|
||||||
|
t.Fatal("LoadAccounts returned nil")
|
||||||
|
}
|
||||||
|
if len(af.Accounts) != 0 {
|
||||||
|
t.Fatalf("got %d accounts, want 0", len(af.Accounts))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSaveLoad_RoundTrip(t *testing.T) {
|
||||||
|
p := withTempAccountsPath(t)
|
||||||
|
af := &AccountsFile{
|
||||||
|
BaseURL: "https://ollama.com",
|
||||||
|
Accounts: []Account{
|
||||||
|
{ID: NewID(), Name: "a1", APIKey: "sk-key1111111111", Created: parseTime(t, "2026-06-19T10:00:00Z")},
|
||||||
|
{ID: NewID(), Name: "a2", APIKey: "sk-key2222222222", BaseURL: "https://staging.example.com"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := af.Save(); err != nil {
|
||||||
|
t.Fatalf("Save: %v", err)
|
||||||
|
}
|
||||||
|
info, err := os.Stat(p)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("stat: %v", err)
|
||||||
|
}
|
||||||
|
if mode := info.Mode().Perm(); mode != 0o600 {
|
||||||
|
t.Errorf("file mode = %o, want 600", mode)
|
||||||
|
}
|
||||||
|
loaded, err := LoadAccounts()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadAccounts: %v", err)
|
||||||
|
}
|
||||||
|
if loaded.BaseURL != af.BaseURL {
|
||||||
|
t.Errorf("BaseURL = %q, want %q", loaded.BaseURL, af.BaseURL)
|
||||||
|
}
|
||||||
|
if len(loaded.Accounts) != 2 {
|
||||||
|
t.Fatalf("len = %d, want 2", len(loaded.Accounts))
|
||||||
|
}
|
||||||
|
if loaded.Accounts[0].Name != "a1" || loaded.Accounts[0].APIKey != "sk-key1111111111" {
|
||||||
|
t.Errorf("account[0] = %+v", loaded.Accounts[0])
|
||||||
|
}
|
||||||
|
if loaded.Accounts[1].BaseURL != "https://staging.example.com" {
|
||||||
|
t.Errorf("account[1].BaseURL = %q", loaded.Accounts[1].BaseURL)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSave_CreatesNoFileWhenDirMissing(t *testing.T) {
|
||||||
|
// AccountsPath targets a dir we never create; Save must fail, not create the dir.
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("XDG_CONFIG_HOME", dir)
|
||||||
|
// remove the ollama-proxy subdir if t.TempDir created it — it doesn't, so
|
||||||
|
// Save() should fail because parent dir does not exist.
|
||||||
|
af := &AccountsFile{Accounts: []Account{{ID: "x", Name: "n", APIKey: "k"}}}
|
||||||
|
// Manually point to a non-existent parent: use a path under dir that lacks dirs.
|
||||||
|
// EnsureAccountsDir not called, so Save should fail.
|
||||||
|
err := af.Save()
|
||||||
|
if err == nil {
|
||||||
|
// Some filesystems auto-create the rename target's parent? No — os.Rename
|
||||||
|
// requires the parent to exist. So this should error.
|
||||||
|
t.Skip("Save unexpectedly succeeded; skipping — but check filesystem semantics")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFind_ByNameOrIDPrefix(t *testing.T) {
|
||||||
|
af := &AccountsFile{
|
||||||
|
Accounts: []Account{
|
||||||
|
{ID: "abcdef0123456789", Name: "alpha", APIKey: "k1"},
|
||||||
|
{ID: "ffeeddccbbaa9988", Name: "beta", APIKey: "k2"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if a := af.Find("alpha"); a == nil || a.ID != "abcdef0123456789" {
|
||||||
|
t.Errorf("Find(name=alpha) = %+v", a)
|
||||||
|
}
|
||||||
|
if a := af.Find("abcdef"); a == nil || a.Name != "alpha" {
|
||||||
|
t.Errorf("Find(id-prefix=abcdef) = %+v", a)
|
||||||
|
}
|
||||||
|
if a := af.Find("nope"); a != nil {
|
||||||
|
t.Errorf("Find(nope) = %+v, want nil", a)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRemove_ByIndexIDAndName(t *testing.T) {
|
||||||
|
af := &AccountsFile{
|
||||||
|
Accounts: []Account{
|
||||||
|
{ID: "abcdef0123456789", Name: "alpha", APIKey: "k1"},
|
||||||
|
{ID: "ffeeddccbbaa9988", Name: "beta", APIKey: "k2"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := af.Remove("alpha"); err != nil {
|
||||||
|
t.Fatalf("Remove(alpha): %v", err)
|
||||||
|
}
|
||||||
|
if len(af.Accounts) != 1 || af.Accounts[0].Name != "beta" {
|
||||||
|
t.Errorf("after Remove(alpha): %+v", af.Accounts)
|
||||||
|
}
|
||||||
|
if err := af.Remove("ffee"); err != nil {
|
||||||
|
t.Fatalf("Remove(id-prefix): %v", err)
|
||||||
|
}
|
||||||
|
if len(af.Accounts) != 0 {
|
||||||
|
t.Errorf("after Remove(ffee): %+v", af.Accounts)
|
||||||
|
}
|
||||||
|
if err := af.Remove("missing"); err != ErrNotFound {
|
||||||
|
t.Errorf("Remove(missing) = %v, want ErrNotFound", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSave_Atomicity_TempFileCleanedOnError(t *testing.T) {
|
||||||
|
// Hard to force a mid-write failure without injection; instead verify that
|
||||||
|
// no leftover temp files exist after a successful Save.
|
||||||
|
p := withTempAccountsPath(t)
|
||||||
|
af := &AccountsFile{Accounts: []Account{{ID: "x", Name: "n", APIKey: "k"}}}
|
||||||
|
if err := af.Save(); err != nil {
|
||||||
|
t.Fatalf("Save: %v", err)
|
||||||
|
}
|
||||||
|
dir := filepath.Dir(p)
|
||||||
|
entries, err := os.ReadDir(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("readdir: %v", err)
|
||||||
|
}
|
||||||
|
for _, e := range entries {
|
||||||
|
name := e.Name()
|
||||||
|
if name != "accounts.json" {
|
||||||
|
t.Errorf("unexpected leftover file in config dir: %s", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAccountsPath_XDGMissing(t *testing.T) {
|
||||||
|
t.Setenv("XDG_CONFIG_HOME", "")
|
||||||
|
home := t.TempDir()
|
||||||
|
t.Setenv("HOME", home)
|
||||||
|
p, err := AccountsPath()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("AccountsPath: %v", err)
|
||||||
|
}
|
||||||
|
want := filepath.Join(home, ".config", "ollama-proxy", "accounts.json")
|
||||||
|
if p != want {
|
||||||
|
t.Errorf("AccountsPath = %q, want %q", p, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSave_LoadAccounts_CorruptFileErrors(t *testing.T) {
|
||||||
|
p := withTempAccountsPath(t)
|
||||||
|
if err := os.WriteFile(p, []byte("{not json"), 0o600); err != nil {
|
||||||
|
t.Fatalf("write corrupt: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := LoadAccounts(); err == nil {
|
||||||
|
t.Error("LoadAccounts on corrupt json returned nil error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure AccountsFile with nil accounts slice marshals to a valid JSON array
|
||||||
|
// rather than null, so consumers never see a missing field.
|
||||||
|
func TestSave_NilAccountsSliceMarshalsAsArray(t *testing.T) {
|
||||||
|
withTempAccountsPath(t)
|
||||||
|
af := &AccountsFile{}
|
||||||
|
if err := af.Save(); err != nil {
|
||||||
|
t.Fatalf("Save: %v", err)
|
||||||
|
}
|
||||||
|
p, _ := AccountsPath()
|
||||||
|
data, _ := os.ReadFile(p)
|
||||||
|
var raw map[string]json.RawMessage
|
||||||
|
if err := json.Unmarshal(data, &raw); err != nil {
|
||||||
|
t.Fatalf("unmarshal: %v", err)
|
||||||
|
}
|
||||||
|
accountsRaw, ok := raw["accounts"]
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("no 'accounts' field in saved json")
|
||||||
|
}
|
||||||
|
if string(accountsRaw) == "null" {
|
||||||
|
t.Errorf("accounts marshaled as null; want []")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseTime(t *testing.T, s string) time.Time {
|
||||||
|
t.Helper()
|
||||||
|
x, err := time.Parse(time.RFC3339, s)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("parseTime %s: %v", s, err)
|
||||||
|
}
|
||||||
|
return x
|
||||||
|
}
|
||||||
91
internal/config/server.go
Normal file
91
internal/config/server.go
Normal file
|
|
@ -0,0 +1,91 @@
|
||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ServerConfig holds the runtime options for the proxy server. Fields map to
|
||||||
|
// CLI flags and (for LogLevel and BaseURL) the OLLAMA_PROXY_* environment
|
||||||
|
// variables, so a systemd unit can set defaults without changing the binary.
|
||||||
|
type ServerConfig struct {
|
||||||
|
Addr string // listen address, e.g. "127.0.0.1:11435"
|
||||||
|
BaseURL string // upstream Ollama Cloud root, e.g. "https://ollama.com"
|
||||||
|
Cooldown time.Duration // per-account 429 cooldown window
|
||||||
|
Retries int // max failover attempts (= number of accounts)
|
||||||
|
LogLevel string // debug | info | warn | error
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultServerConfig returns the canonical defaults used when a flag or env
|
||||||
|
// var is not set.
|
||||||
|
func DefaultServerConfig() ServerConfig {
|
||||||
|
return ServerConfig{
|
||||||
|
Addr: "127.0.0.1:11435",
|
||||||
|
BaseURL: "https://ollama.com",
|
||||||
|
Cooldown: 60 * time.Second,
|
||||||
|
Retries: 3,
|
||||||
|
LogLevel: "info",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ApplyEnv overlays OLLAMA_PROXY_* environment variables on top of the current
|
||||||
|
// config. Empty / unparsable env vars are ignored (the existing value wins).
|
||||||
|
func (c *ServerConfig) ApplyEnv() {
|
||||||
|
if v := os.Getenv("OLLAMA_PROXY_ADDR"); v != "" {
|
||||||
|
c.Addr = v
|
||||||
|
}
|
||||||
|
if v := os.Getenv("OLLAMA_PROXY_BASE_URL"); v != "" {
|
||||||
|
c.BaseURL = v
|
||||||
|
}
|
||||||
|
if v := os.Getenv("OLLAMA_PROXY_COOLDOWN"); v != "" {
|
||||||
|
if d, err := time.ParseDuration(v); err == nil {
|
||||||
|
c.Cooldown = d
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if v := os.Getenv("OLLAMA_PROXY_RETRIES"); v != "" {
|
||||||
|
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||||
|
c.Retries = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if v := os.Getenv("OLLAMA_PROXY_LOG_LEVEL"); v != "" {
|
||||||
|
switch v {
|
||||||
|
case "debug", "info", "warn", "error":
|
||||||
|
c.LogLevel = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate returns an error when a field has an obviously invalid value.
|
||||||
|
func (c *ServerConfig) Validate() error {
|
||||||
|
if c.Addr == "" {
|
||||||
|
return errors.New("addr is required")
|
||||||
|
}
|
||||||
|
if c.BaseURL == "" {
|
||||||
|
return errors.New("base_url is required")
|
||||||
|
}
|
||||||
|
if !startsWithScheme(c.BaseURL, "http://") && !startsWithScheme(c.BaseURL, "https://") {
|
||||||
|
return fmt.Errorf("base_url must start with http:// or https://, got %q", c.BaseURL)
|
||||||
|
}
|
||||||
|
if c.Cooldown < 0 {
|
||||||
|
return errors.New("cooldown must be >= 0")
|
||||||
|
}
|
||||||
|
if c.Retries < 1 {
|
||||||
|
return errors.New("retries must be >= 1")
|
||||||
|
}
|
||||||
|
switch c.LogLevel {
|
||||||
|
case "debug", "info", "warn", "error":
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("log_level must be one of debug|info|warn|error, got %q", c.LogLevel)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func startsWithScheme(s, scheme string) bool {
|
||||||
|
if len(s) < len(scheme) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return s[:len(scheme)] == scheme
|
||||||
|
}
|
||||||
103
internal/config/server_test.go
Normal file
103
internal/config/server_test.go
Normal file
|
|
@ -0,0 +1,103 @@
|
||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDefaultServerConfig(t *testing.T) {
|
||||||
|
c := DefaultServerConfig()
|
||||||
|
if c.Addr != "127.0.0.1:11435" {
|
||||||
|
t.Errorf("Addr = %q", c.Addr)
|
||||||
|
}
|
||||||
|
if c.BaseURL != "https://ollama.com" {
|
||||||
|
t.Errorf("BaseURL = %q", c.BaseURL)
|
||||||
|
}
|
||||||
|
if c.Cooldown != 60*time.Second {
|
||||||
|
t.Errorf("Cooldown = %v", c.Cooldown)
|
||||||
|
}
|
||||||
|
if c.Retries != 3 {
|
||||||
|
t.Errorf("Retries = %d", c.Retries)
|
||||||
|
}
|
||||||
|
if c.LogLevel != "info" {
|
||||||
|
t.Errorf("LogLevel = %q", c.LogLevel)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestServerConfig_Validate(t *testing.T) {
|
||||||
|
good := DefaultServerConfig()
|
||||||
|
if err := good.Validate(); err != nil {
|
||||||
|
t.Errorf("default config invalid: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
mut func(*ServerConfig)
|
||||||
|
}{
|
||||||
|
{"empty addr", func(c *ServerConfig) { c.Addr = "" }},
|
||||||
|
{"empty base_url", func(c *ServerConfig) { c.BaseURL = "" }},
|
||||||
|
{"bad scheme", func(c *ServerConfig) { c.BaseURL = "ftp://x" }},
|
||||||
|
{"negative cooldown", func(c *ServerConfig) { c.Cooldown = -1 * time.Second }},
|
||||||
|
{"zero retries", func(c *ServerConfig) { c.Retries = 0 }},
|
||||||
|
{"bad log level", func(c *ServerConfig) { c.LogLevel = "trace" }},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
c := DefaultServerConfig()
|
||||||
|
tc.mut(&c)
|
||||||
|
if err := c.Validate(); err == nil {
|
||||||
|
t.Errorf("Validate(%s) = nil, want error", tc.name)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestServerConfig_ApplyEnv(t *testing.T) {
|
||||||
|
t.Setenv("OLLAMA_PROXY_ADDR", "0.0.0.0:9000")
|
||||||
|
t.Setenv("OLLAMA_PROXY_BASE_URL", "https://staging.example.com")
|
||||||
|
t.Setenv("OLLAMA_PROXY_COOLDOWN", "120s")
|
||||||
|
t.Setenv("OLLAMA_PROXY_RETRIES", "5")
|
||||||
|
t.Setenv("OLLAMA_PROXY_LOG_LEVEL", "debug")
|
||||||
|
|
||||||
|
c := DefaultServerConfig()
|
||||||
|
c.ApplyEnv()
|
||||||
|
|
||||||
|
if c.Addr != "0.0.0.0:9000" {
|
||||||
|
t.Errorf("Addr = %q", c.Addr)
|
||||||
|
}
|
||||||
|
if c.BaseURL != "https://staging.example.com" {
|
||||||
|
t.Errorf("BaseURL = %q", c.BaseURL)
|
||||||
|
}
|
||||||
|
if c.Cooldown != 120*time.Second {
|
||||||
|
t.Errorf("Cooldown = %v", c.Cooldown)
|
||||||
|
}
|
||||||
|
if c.Retries != 5 {
|
||||||
|
t.Errorf("Retries = %d", c.Retries)
|
||||||
|
}
|
||||||
|
if c.LogLevel != "debug" {
|
||||||
|
t.Errorf("LogLevel = %q", c.LogLevel)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestServerConfig_ApplyEnv_IgnoresInvalid(t *testing.T) {
|
||||||
|
t.Setenv("OLLAMA_PROXY_COOLDOWN", "not-a-duration")
|
||||||
|
t.Setenv("OLLAMA_PROXY_RETRIES", "NaN")
|
||||||
|
t.Setenv("OLLAMA_PROXY_LOG_LEVEL", "trace")
|
||||||
|
t.Setenv("OLLAMA_PROXY_ADDR", "")
|
||||||
|
|
||||||
|
c := DefaultServerConfig()
|
||||||
|
c.ApplyEnv()
|
||||||
|
|
||||||
|
if c.Cooldown != 60*time.Second {
|
||||||
|
t.Errorf("bad cooldown env should be ignored, got %v", c.Cooldown)
|
||||||
|
}
|
||||||
|
if c.Retries != 3 {
|
||||||
|
t.Errorf("bad retries env should be ignored, got %d", c.Retries)
|
||||||
|
}
|
||||||
|
if c.LogLevel != "info" {
|
||||||
|
t.Errorf("bad log_level env should be ignored, got %q", c.LogLevel)
|
||||||
|
}
|
||||||
|
if c.Addr != "127.0.0.1:11435" {
|
||||||
|
t.Errorf("empty addr env should keep default, got %q", c.Addr)
|
||||||
|
}
|
||||||
|
}
|
||||||
29
internal/log/logger.go
Normal file
29
internal/log/logger.go
Normal file
|
|
@ -0,0 +1,29 @@
|
||||||
|
// Package log provides a thin slog wrapper so callers can configure the level
|
||||||
|
// from a single string ("debug"|"info"|"warn"|"error").
|
||||||
|
package log
|
||||||
|
|
||||||
|
import (
|
||||||
|
"log/slog"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// New returns a slog.Logger writing JSON to stdout at the given level. An
|
||||||
|
// unknown level falls back to info.
|
||||||
|
func New(level string) *slog.Logger {
|
||||||
|
var lvl slog.Level
|
||||||
|
switch strings.ToLower(level) {
|
||||||
|
case "debug":
|
||||||
|
lvl = slog.LevelDebug
|
||||||
|
case "info":
|
||||||
|
lvl = slog.LevelInfo
|
||||||
|
case "warn":
|
||||||
|
lvl = slog.LevelWarn
|
||||||
|
case "error":
|
||||||
|
lvl = slog.LevelError
|
||||||
|
default:
|
||||||
|
lvl = slog.LevelInfo
|
||||||
|
}
|
||||||
|
h := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: lvl})
|
||||||
|
return slog.New(h)
|
||||||
|
}
|
||||||
158
internal/proxy/balancer.go
Normal file
158
internal/proxy/balancer.go
Normal file
|
|
@ -0,0 +1,158 @@
|
||||||
|
// Package proxy implements the Ollama Cloud reverse proxy: account selection
|
||||||
|
// (round-robin with per-account cooldown), the HTTP reverse-proxy handler,
|
||||||
|
// and the pre-stream failover logic.
|
||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Balancer selects the next upstream account using round-robin and skips
|
||||||
|
// accounts currently in cooldown. It is safe for concurrent use.
|
||||||
|
type Balancer struct {
|
||||||
|
accounts []config.Account
|
||||||
|
rr atomic.Uint64 // round-robin counter
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
cooldown map[string]time.Time // account ID -> until
|
||||||
|
last429 map[string]time.Time // account ID -> last 429 time (for display)
|
||||||
|
cooldownDur time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewBalancer builds a Balancer from a list of accounts. The cooldown window
|
||||||
|
// 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
|
||||||
|
accts := make([]config.Account, len(accounts))
|
||||||
|
copy(accts, accounts)
|
||||||
|
return &Balancer{
|
||||||
|
accounts: accts,
|
||||||
|
cooldown: make(map[string]time.Time),
|
||||||
|
last429: make(map[string]time.Time),
|
||||||
|
cooldownDur: cooldown,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Len returns the number of accounts the balancer knows about.
|
||||||
|
func (b *Balancer) Len() int { return len(b.accounts) }
|
||||||
|
|
||||||
|
// Account returns the i-th account (mostly for tests / display).
|
||||||
|
func (b *Balancer) Account(i int) config.Account { return b.accounts[i] }
|
||||||
|
|
||||||
|
// Accounts returns a shallow copy of the account list.
|
||||||
|
func (b *Balancer) Accounts() []config.Account {
|
||||||
|
out := make([]config.Account, len(b.accounts))
|
||||||
|
copy(out, b.accounts)
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// ErrNoAccounts is returned when the balancer has no accounts configured.
|
||||||
|
type ErrNoAccounts struct{}
|
||||||
|
|
||||||
|
func (ErrNoAccounts) Error() string { return "no accounts configured" }
|
||||||
|
|
||||||
|
// ErrAllCooldown is returned when every account is currently in cooldown.
|
||||||
|
type ErrAllCooldown struct {
|
||||||
|
// Until is the earliest time at which any account becomes available again.
|
||||||
|
Until time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e ErrAllCooldown) Error() string {
|
||||||
|
return "all accounts are in cooldown until " + e.Until.Format(time.RFC3339)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Next picks the next available account, skipping accounts in cooldown. It
|
||||||
|
// rotates starting from the round-robin counter so consecutive calls land on
|
||||||
|
// different accounts when possible. Returns ErrNoAccounts or ErrAllCooldown
|
||||||
|
// when nothing is available.
|
||||||
|
func (b *Balancer) Next() (config.Account, error) {
|
||||||
|
if len(b.accounts) == 0 {
|
||||||
|
return config.Account{}, ErrNoAccounts{}
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
var earliest time.Time
|
||||||
|
for i := 0; i < len(b.accounts); i++ {
|
||||||
|
idx := int(b.rr.Add(1)) % len(b.accounts)
|
||||||
|
// 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 earliest.IsZero() || b.cooldownUntil(b.accounts[idx].ID).Before(earliest) {
|
||||||
|
earliest = b.cooldownUntil(b.accounts[idx].ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return config.Account{}, ErrAllCooldown{Until: earliest}
|
||||||
|
}
|
||||||
|
|
||||||
|
// isCooldown reports whether the account is currently rate-limited.
|
||||||
|
func (b *Balancer) isCooldown(id string, now time.Time) bool {
|
||||||
|
b.mu.RLock()
|
||||||
|
defer b.mu.RUnlock()
|
||||||
|
until, ok := b.cooldown[id]
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return now.Before(until)
|
||||||
|
}
|
||||||
|
|
||||||
|
// cooldownUntil returns the cooldown expiry for an account (zero if none).
|
||||||
|
func (b *Balancer) cooldownUntil(id string) time.Time {
|
||||||
|
b.mu.RLock()
|
||||||
|
defer b.mu.RUnlock()
|
||||||
|
return b.cooldown[id]
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarkCooldown puts the given account into cooldown for the configured window
|
||||||
|
// starting from now. Safe to call concurrently; idempotent (extends the window).
|
||||||
|
func (b *Balancer) MarkCooldown(id string) {
|
||||||
|
b.MarkCooldownFor(id, b.cooldownDur)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarkCooldownFor puts the account into cooldown for an explicit duration.
|
||||||
|
func (b *Balancer) MarkCooldownFor(id string, d time.Duration) {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
until := time.Now().Add(d)
|
||||||
|
b.cooldown[id] = until
|
||||||
|
b.last429[id] = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClearCooldown removes the cooldown for an account (e.g. on a successful
|
||||||
|
// request after earlier failures).
|
||||||
|
func (b *Balancer) ClearCooldown(id string) {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
delete(b.cooldown, id)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Status returns a snapshot of per-account cooldown state, suitable for
|
||||||
|
// display in `accounts list`. The map key is the account ID.
|
||||||
|
func (b *Balancer) Status() map[string]AccountStatus {
|
||||||
|
b.mu.RLock()
|
||||||
|
defer b.mu.RUnlock()
|
||||||
|
out := make(map[string]AccountStatus, len(b.accounts))
|
||||||
|
now := time.Now()
|
||||||
|
for _, a := range b.accounts {
|
||||||
|
until := b.cooldown[a.ID]
|
||||||
|
st := AccountStatus{
|
||||||
|
InCooldown: !until.IsZero() && now.Before(until),
|
||||||
|
Until: until,
|
||||||
|
Last429: b.last429[a.ID],
|
||||||
|
}
|
||||||
|
out[a.ID] = st
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// AccountStatus is the per-account cooldown snapshot returned by Status.
|
||||||
|
type AccountStatus struct {
|
||||||
|
InCooldown bool
|
||||||
|
Until time.Time
|
||||||
|
Last429 time.Time
|
||||||
|
}
|
||||||
160
internal/proxy/balancer_test.go
Normal file
160
internal/proxy/balancer_test.go
Normal file
|
|
@ -0,0 +1,160 @@
|
||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
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))}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Next_RoundRobin(t *testing.T) {
|
||||||
|
b := NewBalancer(mkAccts(3), 60*time.Second)
|
||||||
|
seen := make(map[string]int)
|
||||||
|
for i := 0; i < 9; i++ {
|
||||||
|
a, err := b.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next %d: %v", i, err)
|
||||||
|
}
|
||||||
|
seen[a.ID]++
|
||||||
|
}
|
||||||
|
// 9 calls across 3 accounts should hit each exactly 3 times (round-robin).
|
||||||
|
for id, n := range seen {
|
||||||
|
if n != 3 {
|
||||||
|
t.Errorf("account %s got %d hits, want 3", id, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Next_NoAccounts(t *testing.T) {
|
||||||
|
b := NewBalancer(nil, 60*time.Second)
|
||||||
|
_, err := b.Next()
|
||||||
|
if _, ok := err.(ErrNoAccounts); !ok {
|
||||||
|
t.Errorf("Next on empty balancer: %v, want ErrNoAccounts", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Next_SkipsCooldown(t *testing.T) {
|
||||||
|
accts := mkAccts(3)
|
||||||
|
b := NewBalancer(accts, 60*time.Second)
|
||||||
|
// Put the first account in cooldown; the next 3 calls should never pick it.
|
||||||
|
b.MarkCooldown(accts[0].ID)
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
a, err := b.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next %d: %v", i, err)
|
||||||
|
}
|
||||||
|
if a.ID == accts[0].ID {
|
||||||
|
t.Errorf("Next %d returned cooldown account %s", i, a.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Next_AllCooldown(t *testing.T) {
|
||||||
|
accts := mkAccts(2)
|
||||||
|
b := NewBalancer(accts, 60*time.Second)
|
||||||
|
b.MarkCooldown(accts[0].ID)
|
||||||
|
b.MarkCooldown(accts[1].ID)
|
||||||
|
_, err := b.Next()
|
||||||
|
if _, ok := err.(ErrAllCooldown); !ok {
|
||||||
|
t.Errorf("Next with all in cooldown: %v, want ErrAllCooldown", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_ClearCooldown(t *testing.T) {
|
||||||
|
accts := mkAccts(2)
|
||||||
|
b := NewBalancer(accts, 60*time.Second)
|
||||||
|
b.MarkCooldown(accts[0].ID)
|
||||||
|
b.ClearCooldown(accts[0].ID)
|
||||||
|
// After clearing, the first account must be selectable again. Spin until
|
||||||
|
// we see it — round-robin visits each account within len(accts) calls.
|
||||||
|
seen := false
|
||||||
|
for i := 0; i < len(accts)*3; i++ {
|
||||||
|
a, err := b.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next %d: %v", i, err)
|
||||||
|
}
|
||||||
|
if a.ID == accts[0].ID {
|
||||||
|
seen = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !seen {
|
||||||
|
t.Error("cleared account was never returned by Next")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Next_Concurrent(t *testing.T) {
|
||||||
|
b := NewBalancer(mkAccts(3), 60*time.Second)
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
const goroutines = 16
|
||||||
|
const perG = 100
|
||||||
|
results := make(chan string, goroutines*perG)
|
||||||
|
for g := 0; g < goroutines; g++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
for i := 0; i < perG; i++ {
|
||||||
|
a, err := b.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Next: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
results <- a.ID
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
close(results)
|
||||||
|
// Sanity: every result is one of the known account IDs. We don't assert
|
||||||
|
// strict fairness under contention (atomic counter gives best-effort RR).
|
||||||
|
known := map[string]bool{"ida": true, "idb": true, "idc": true}
|
||||||
|
count := 0
|
||||||
|
for id := range results {
|
||||||
|
if !known[id] {
|
||||||
|
t.Errorf("unknown account id returned: %s", id)
|
||||||
|
}
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
if count != goroutines*perG {
|
||||||
|
t.Errorf("got %d results, want %d", count, goroutines*perG)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_Status(t *testing.T) {
|
||||||
|
accts := mkAccts(2)
|
||||||
|
b := NewBalancer(accts, 60*time.Second)
|
||||||
|
b.MarkCooldown(accts[0].ID)
|
||||||
|
st := b.Status()
|
||||||
|
if !st[accts[0].ID].InCooldown {
|
||||||
|
t.Error("account 0 should be in cooldown")
|
||||||
|
}
|
||||||
|
if st[accts[1].ID].InCooldown {
|
||||||
|
t.Error("account 1 should not be in cooldown")
|
||||||
|
}
|
||||||
|
if st[accts[0].ID].Last429.IsZero() {
|
||||||
|
t.Error("account 0 should have Last429 set")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBalancer_CooldownExpires(t *testing.T) {
|
||||||
|
accts := mkAccts(1)
|
||||||
|
b := NewBalancer(accts, 50*time.Millisecond)
|
||||||
|
b.MarkCooldown(accts[0].ID)
|
||||||
|
// Immediately: account is in cooldown.
|
||||||
|
if _, err := b.Next(); err == nil {
|
||||||
|
t.Fatal("Next succeeded immediately after MarkCooldown")
|
||||||
|
}
|
||||||
|
time.Sleep(60 * time.Millisecond)
|
||||||
|
if _, err := b.Next(); err != nil {
|
||||||
|
t.Errorf("Next after cooldown expiry: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
294
internal/proxy/handler.go
Normal file
294
internal/proxy/handler.go
Normal file
|
|
@ -0,0 +1,294 @@
|
||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log/slog"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Handler is the http.Handler that proxies incoming requests to Ollama Cloud
|
||||||
|
// 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
|
||||||
|
Retries int // max attempts per request (= number of accounts to try)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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).
|
||||||
|
func NewHandler(b *Balancer, baseURL string, retries int, log *slog.Logger) *Handler {
|
||||||
|
return &Handler{
|
||||||
|
Balancer: b,
|
||||||
|
BaseURL: strings.TrimRight(baseURL, "/"),
|
||||||
|
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
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Log: log,
|
||||||
|
Retries: retries,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// allowedPaths is the allowlist of upstream paths the proxy will forward.
|
||||||
|
// Anything else returns 404 — we never proxy arbitrary paths.
|
||||||
|
var allowedPaths = map[string]bool{
|
||||||
|
"/api/chat": true,
|
||||||
|
"/api/generate": true,
|
||||||
|
"/api/tags": true,
|
||||||
|
"/api/show": true,
|
||||||
|
"/api/ps": true,
|
||||||
|
"/api/version": true,
|
||||||
|
"/api/delete": true,
|
||||||
|
"/v1/chat/completions": true,
|
||||||
|
"/v1/completions": true,
|
||||||
|
"/v1/models": true,
|
||||||
|
"/v1/embeddings": true,
|
||||||
|
"/v1/files": true,
|
||||||
|
}
|
||||||
|
|
||||||
|
// isAllowed reports whether a path (possibly with a trailing slash or query)
|
||||||
|
// matches one of the allowed upstream endpoints.
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// ServeHTTP proxies a single client request to Ollama Cloud with failover.
|
||||||
|
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !isAllowed(r.URL.Path) {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if h.Balancer.Len() == 0 {
|
||||||
|
writeError(w, http.StatusServiceUnavailable, "no accounts configured")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
maxAttempts := h.Retries
|
||||||
|
if maxAttempts < 1 {
|
||||||
|
maxAttempts = 1
|
||||||
|
}
|
||||||
|
if maxAttempts > h.Balancer.Len() {
|
||||||
|
maxAttempts = h.Balancer.Len()
|
||||||
|
}
|
||||||
|
|
||||||
|
h.proxyWithFailover(w, r, 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) {
|
||||||
|
ctx := r.Context()
|
||||||
|
var lastErr error
|
||||||
|
|
||||||
|
for attempt := 0; attempt < maxAttempts; attempt++ {
|
||||||
|
acct, err := h.Balancer.Next()
|
||||||
|
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()
|
||||||
|
}
|
||||||
|
writeError(w, http.StatusTooManyRequests, msg)
|
||||||
|
return
|
||||||
|
case ErrNoAccounts:
|
||||||
|
writeError(w, http.StatusServiceUnavailable, err.Error())
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
writeError(w, http.StatusServiceUnavailable, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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()
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL, r.Body)
|
||||||
|
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
|
||||||
|
|
||||||
|
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,
|
||||||
|
"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)
|
||||||
|
resp.Body.Close()
|
||||||
|
h.Log.Warn("upstream rejected, failover",
|
||||||
|
"account", acct.Name, "account_id", acct.ID,
|
||||||
|
"path", r.URL.Path, "status", resp.StatusCode,
|
||||||
|
"latency_ms", latency.Milliseconds(), "preview", bodyPreview)
|
||||||
|
if resp.StatusCode == http.StatusTooManyRequests {
|
||||||
|
h.Balancer.MarkCooldown(acct.ID)
|
||||||
|
}
|
||||||
|
lastErr = fmt.Errorf("upstream %d: %s", resp.StatusCode, bodyPreview)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2xx (or other non-retriable): commit and stream.
|
||||||
|
h.copyResponse(w, resp)
|
||||||
|
h.Log.Info("proxied",
|
||||||
|
"account", acct.Name, "account_id", acct.ID,
|
||||||
|
"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).
|
||||||
|
if h.Balancer.Len() > 0 {
|
||||||
|
allCooldown := true
|
||||||
|
status := h.Balancer.Status()
|
||||||
|
for _, a := range h.Balancer.Accounts() {
|
||||||
|
if !status[a.ID].InCooldown {
|
||||||
|
allCooldown = false
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if allCooldown {
|
||||||
|
msg := "all accounts rate-limited"
|
||||||
|
if lastErr != nil {
|
||||||
|
msg = msg + ": " + lastErr.Error()
|
||||||
|
}
|
||||||
|
writeError(w, http.StatusTooManyRequests, msg)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if lastErr != nil {
|
||||||
|
writeError(w, http.StatusBadGateway, "all accounts failed: "+lastErr.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
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.
|
||||||
|
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 {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
if !errors.Is(err, io.EOF) {
|
||||||
|
h.Log.Debug("upstream body read ended", "err", err)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// copyHeaders duplicates src into dst, dropping hop-by-hop headers and any
|
||||||
|
// Authorization that the client may have sent (we always set our own).
|
||||||
|
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
|
||||||
|
}
|
||||||
|
for k, vs := range src {
|
||||||
|
skip := false
|
||||||
|
for _, h := range hopByHop {
|
||||||
|
if strings.EqualFold(k, h) {
|
||||||
|
skip = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if skip {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, v := range vs {
|
||||||
|
dst.Add(k, v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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)
|
||||||
|
m, _ := r.Read(buf)
|
||||||
|
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") {
|
||||||
|
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)
|
||||||
|
_, _ = fmt.Fprintf(w, `{"error":{"message":%q,"type":"ollama_proxy"}}`, msg)
|
||||||
|
}
|
||||||
245
internal/proxy/handler_test.go
Normal file
245
internal/proxy/handler_test.go
Normal file
|
|
@ -0,0 +1,245 @@
|
||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/config"
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/log"
|
||||||
|
)
|
||||||
|
|
||||||
|
// mockUpstream is a configurable test upstream. It responds to /api/chat and
|
||||||
|
// /v1/chat/completions; for other paths it returns 200 with a short body.
|
||||||
|
type mockUpstream struct {
|
||||||
|
status int32 // current status to return
|
||||||
|
body string
|
||||||
|
allowAuth string // if non-empty, require this Bearer token
|
||||||
|
chunks []string
|
||||||
|
requestAuth atomic.Int32 // number of requests seen with each auth value
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockUpstream) handler(w http.ResponseWriter, r *http.Request) {
|
||||||
|
auth := r.Header.Get("Authorization")
|
||||||
|
if m.allowAuth != "" && auth != "Bearer "+m.allowAuth {
|
||||||
|
http.Error(w, "bad auth", http.StatusUnauthorized)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if m.status >= 500 || m.status == 429 {
|
||||||
|
http.Error(w, "rate limited", int(m.status))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if len(m.chunks) > 0 {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Header().Set("Transfer-Encoding", "chunked")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
flusher, _ := w.(http.Flusher)
|
||||||
|
for _, c := range m.chunks {
|
||||||
|
_, _ = io.WriteString(w, c)
|
||||||
|
if flusher != nil {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(int(m.status))
|
||||||
|
_, _ = io.WriteString(w, m.body)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newTestHandler(t *testing.T, upstreamURL string, accounts []config.Account, retries int) *Handler {
|
||||||
|
t.Helper()
|
||||||
|
b := NewBalancer(accounts, 50*time.Millisecond)
|
||||||
|
return NewHandler(b, upstreamURL, retries, log.New("debug"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_ProxiesVersion(t *testing.T) {
|
||||||
|
up := &mockUpstream{status: 200, body: `{"version":"0.1.42"}`}
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != 200 {
|
||||||
|
t.Fatalf("status = %d, want 200", rec.Code)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rec.Body.String(), "0.1.42") {
|
||||||
|
t.Errorf("body = %q, want version", rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_FailoverOn429(t *testing.T) {
|
||||||
|
// Two upstreams: first always 429, second always 200.
|
||||||
|
badUp := &mockUpstream{status: 429, body: "rate limited"}
|
||||||
|
badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
|
||||||
|
goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
|
||||||
|
goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
|
||||||
|
defer badSrv.Close()
|
||||||
|
defer goodSrv.Close()
|
||||||
|
|
||||||
|
// Two accounts, each pointing at a different upstream via per-account BaseURL.
|
||||||
|
accts := []config.Account{
|
||||||
|
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
|
||||||
|
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
|
||||||
|
}
|
||||||
|
h := newTestHandler(t, "https://unused.example.com", 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 (failover to good account)", rec.Code)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rec.Body.String(), `"ok":true`) {
|
||||||
|
t.Errorf("body = %q, want good upstream response", rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_AllAccounts429_Returns429(t *testing.T) {
|
||||||
|
up := &mockUpstream{status: 429, body: "rate limited"}
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
accts := []config.Account{
|
||||||
|
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: srv.URL},
|
||||||
|
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: srv.URL},
|
||||||
|
}
|
||||||
|
h := newTestHandler(t, srv.URL, accts, 2)
|
||||||
|
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusTooManyRequests {
|
||||||
|
t.Fatalf("status = %d, want 429", rec.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_SSEStreaming(t *testing.T) {
|
||||||
|
up := &mockUpstream{
|
||||||
|
status: 200,
|
||||||
|
chunks: []string{"data: {\"a\":1}\n\n", "data: {\"a\":2}\n\n", "data: [DONE]\n\n"},
|
||||||
|
}
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"stream":true}`))
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != 200 {
|
||||||
|
t.Fatalf("status = %d, want 200", rec.Code)
|
||||||
|
}
|
||||||
|
body := rec.Body.String()
|
||||||
|
if !strings.Contains(body, "[DONE]") {
|
||||||
|
t.Errorf("body = %q, want to contain [DONE]", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, "\"a\":1") || !strings.Contains(body, "\"a\":2") {
|
||||||
|
t.Errorf("body = %q, missing chunks", body)
|
||||||
|
}
|
||||||
|
if ct := rec.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
|
||||||
|
t.Errorf("Content-Type = %q, want text/event-stream", ct)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_RejectsUnknownPath(t *testing.T) {
|
||||||
|
up := &mockUpstream{status: 200, body: "x"}
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("GET", "/admin", nil)
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
if rec.Code != http.StatusNotFound {
|
||||||
|
t.Errorf("status = %d, want 404 for unknown path", rec.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_NoAccounts(t *testing.T) {
|
||||||
|
h := newTestHandler(t, "http://unused.example.com", nil, 1)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
if rec.Code != http.StatusServiceUnavailable {
|
||||||
|
t.Errorf("status = %d, want 503", rec.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_StripsClientAuthorization(t *testing.T) {
|
||||||
|
// Upstream echoes the received Authorization header in its body so we can
|
||||||
|
// assert the proxy overrode it with the account's key.
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
_, _ = io.WriteString(w, `{"got_auth":"`+r.Header.Get("Authorization")+`"}`)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
accts := []config.Account{{ID: "idA", Name: "a", APIKey: "sk-account-key-1234", BaseURL: srv.URL}}
|
||||||
|
h := newTestHandler(t, srv.URL, accts, 1)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer client-should-be-stripped")
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if !strings.Contains(rec.Body.String(), "sk-account-key-1234") {
|
||||||
|
t.Errorf("body = %q, want account key forwarded", rec.Body.String())
|
||||||
|
}
|
||||||
|
if strings.Contains(rec.Body.String(), "client-should-be-stripped") {
|
||||||
|
t.Errorf("body = %q, client Authorization leaked to upstream", rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandler_RetriesOn5xx(t *testing.T) {
|
||||||
|
// First account returns 500, second returns 200.
|
||||||
|
badUp := &mockUpstream{status: 500}
|
||||||
|
badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
|
||||||
|
goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
|
||||||
|
goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
|
||||||
|
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, "https://unused.example.com", 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 (failover on 500)", rec.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsAllowed(t *testing.T) {
|
||||||
|
good := []string{
|
||||||
|
"/api/chat", "/api/generate", "/api/tags", "/api/show", "/api/ps",
|
||||||
|
"/api/version", "/api/delete",
|
||||||
|
"/v1/chat/completions", "/v1/completions", "/v1/models",
|
||||||
|
"/v1/embeddings", "/v1/files", "/v1/files/abc-123",
|
||||||
|
}
|
||||||
|
for _, p := range good {
|
||||||
|
if !isAllowed(p) {
|
||||||
|
t.Errorf("isAllowed(%q) = false, want true", p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
bad := []string{"/admin", "/", "/api/unknown", "/v2/foo", "/etc/passwd"}
|
||||||
|
for _, p := range bad {
|
||||||
|
if isAllowed(p) {
|
||||||
|
t.Errorf("isAllowed(%q) = true, want false", p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
28
internal/proxy/retry.go
Normal file
28
internal/proxy/retry.go
Normal file
|
|
@ -0,0 +1,28 @@
|
||||||
|
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.
|
||||||
|
|
||||||
|
// 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.
|
||||||
|
//
|
||||||
|
// 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).
|
||||||
|
//
|
||||||
|
// Once copyResponse has called WriteHeader, no further retries are possible
|
||||||
|
// for this request.
|
||||||
11
main.go
Normal file
11
main.go
Normal file
|
|
@ -0,0 +1,11 @@
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"github.com/Atte149/ollama-proxy/internal/cli"
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
os.Exit(cli.Run(os.Args))
|
||||||
|
}
|
||||||
19
systemd/ollama-proxy.service
Normal file
19
systemd/ollama-proxy.service
Normal file
|
|
@ -0,0 +1,19 @@
|
||||||
|
[Unit]
|
||||||
|
Description=Ollama Cloud multi-account proxy
|
||||||
|
Documentation=https://git.dueattendant149.org/Atte149/ollama-proxy
|
||||||
|
After=network-online.target
|
||||||
|
Wants=network-online.target
|
||||||
|
|
||||||
|
[Service]
|
||||||
|
Type=simple
|
||||||
|
User=dueattendant149
|
||||||
|
Group=dueattendant149
|
||||||
|
ExecStart=/usr/local/bin/ollama-proxy serve --addr 127.0.0.1:11435 --log-level info
|
||||||
|
Restart=always
|
||||||
|
RestartSec=3
|
||||||
|
Environment=OLLAMA_PROXY_LOG_LEVEL=info
|
||||||
|
NoNewPrivileges=true
|
||||||
|
PrivateTmp=true
|
||||||
|
|
||||||
|
[Install]
|
||||||
|
WantedBy=multi-user.target
|
||||||
Loading…
Add table
Add a link
Reference in a new issue