From 3963eede703b76717224ed81e26f79b03fd9edc0 Mon Sep 17 00:00:00 2001 From: Atte149 Date: Fri, 19 Jun 2026 13:58:13 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20ollama-proxy=20=E2=80=94=20multi-accoun?= =?UTF-8?q?t=20Ollama=20Cloud=20reverse=20proxy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .gitignore | 4 + AGENTS.md | 91 ++++++++++ Makefile | 28 +++ README.md | 140 +++++++++++++++ go.mod | 7 + go.sum | 4 + internal/cli/accounts.go | 223 +++++++++++++++++++++++ internal/cli/integration_test.go | 241 +++++++++++++++++++++++++ internal/cli/now.go | 7 + internal/cli/root.go | 59 +++++++ internal/cli/serve.go | 109 ++++++++++++ internal/config/accounts.go | 218 +++++++++++++++++++++++ internal/config/accounts_test.go | 265 ++++++++++++++++++++++++++++ internal/config/server.go | 91 ++++++++++ internal/config/server_test.go | 103 +++++++++++ internal/log/logger.go | 29 +++ internal/proxy/balancer.go | 158 +++++++++++++++++ internal/proxy/balancer_test.go | 160 +++++++++++++++++ internal/proxy/handler.go | 294 +++++++++++++++++++++++++++++++ internal/proxy/handler_test.go | 245 ++++++++++++++++++++++++++ internal/proxy/retry.go | 28 +++ main.go | 11 ++ systemd/ollama-proxy.service | 19 ++ 23 files changed, 2534 insertions(+) create mode 100644 .gitignore create mode 100644 AGENTS.md create mode 100644 Makefile create mode 100644 README.md create mode 100644 go.mod create mode 100644 go.sum create mode 100644 internal/cli/accounts.go create mode 100644 internal/cli/integration_test.go create mode 100644 internal/cli/now.go create mode 100644 internal/cli/root.go create mode 100644 internal/cli/serve.go create mode 100644 internal/config/accounts.go create mode 100644 internal/config/accounts_test.go create mode 100644 internal/config/server.go create mode 100644 internal/config/server_test.go create mode 100644 internal/log/logger.go create mode 100644 internal/proxy/balancer.go create mode 100644 internal/proxy/balancer_test.go create mode 100644 internal/proxy/handler.go create mode 100644 internal/proxy/handler_test.go create mode 100644 internal/proxy/retry.go create mode 100644 main.go create mode 100644 systemd/ollama-proxy.service diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..976f261 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +ollama-proxy +*.test +*.out +/.scratch/ \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..6c36060 --- /dev/null +++ b/AGENTS.md @@ -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 ` (промпт) или `... add --name ` +- **Перезапустить после правки аккаунтов:** `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` \ No newline at end of file diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..390fc7a --- /dev/null +++ b/Makefile @@ -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 \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..3b07927 --- /dev/null +++ b/README.md @@ -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 . + +## 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. \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..aa251fe --- /dev/null +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..6208220 --- /dev/null +++ b/go.sum @@ -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= diff --git a/internal/cli/accounts.go b/internal/cli/accounts.go new file mode 100644 index 0000000..ca7da11 --- /dev/null +++ b/internal/cli/accounts.go @@ -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 ] + ollama-proxy accounts list + ollama-proxy accounts remove + ollama-proxy accounts set-base-url +`) +} + +// accountsAdd implements `accounts add [API_KEY] [--name ]`. +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 ]") + 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 ") + 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 ") + 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:] +} diff --git a/internal/cli/integration_test.go b/internal/cli/integration_test.go new file mode 100644 index 0000000..fe31df7 --- /dev/null +++ b/internal/cli/integration_test.go @@ -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) + } +} diff --git a/internal/cli/now.go b/internal/cli/now.go new file mode 100644 index 0000000..45a214b --- /dev/null +++ b/internal/cli/now.go @@ -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() } diff --git a/internal/cli/root.go b/internal/cli/root.go new file mode 100644 index 0000000..59f2376 --- /dev/null +++ b/internal/cli/root.go @@ -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 Manage Ollama Cloud API keys + ollama-proxy version Print version + +Accounts subcommands: + accounts add [API_KEY] [--name ] Add an account (prompts for key if not given) + accounts list List configured accounts + accounts remove Remove an account + accounts set-base-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) +`) +} diff --git a/internal/cli/serve.go b/internal/cli/serve.go new file mode 100644 index 0000000..895cf95 --- /dev/null +++ b/internal/cli/serve.go @@ -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 diff --git a/internal/config/accounts.go b/internal/config/accounts.go new file mode 100644 index 0000000..d214d82 --- /dev/null +++ b/internal/config/accounts.go @@ -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 +} diff --git a/internal/config/accounts_test.go b/internal/config/accounts_test.go new file mode 100644 index 0000000..342902e --- /dev/null +++ b/internal/config/accounts_test.go @@ -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 +} diff --git a/internal/config/server.go b/internal/config/server.go new file mode 100644 index 0000000..457f15e --- /dev/null +++ b/internal/config/server.go @@ -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 +} diff --git a/internal/config/server_test.go b/internal/config/server_test.go new file mode 100644 index 0000000..3d5e69e --- /dev/null +++ b/internal/config/server_test.go @@ -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) + } +} diff --git a/internal/log/logger.go b/internal/log/logger.go new file mode 100644 index 0000000..d7c3af2 --- /dev/null +++ b/internal/log/logger.go @@ -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) +} diff --git a/internal/proxy/balancer.go b/internal/proxy/balancer.go new file mode 100644 index 0000000..4f8b9c2 --- /dev/null +++ b/internal/proxy/balancer.go @@ -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 +} diff --git a/internal/proxy/balancer_test.go b/internal/proxy/balancer_test.go new file mode 100644 index 0000000..57af6f7 --- /dev/null +++ b/internal/proxy/balancer_test.go @@ -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) + } +} diff --git a/internal/proxy/handler.go b/internal/proxy/handler.go new file mode 100644 index 0000000..12f0941 --- /dev/null +++ b/internal/proxy/handler.go @@ -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/ 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) +} diff --git a/internal/proxy/handler_test.go b/internal/proxy/handler_test.go new file mode 100644 index 0000000..349613a --- /dev/null +++ b/internal/proxy/handler_test.go @@ -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) + } + } +} diff --git a/internal/proxy/retry.go b/internal/proxy/retry.go new file mode 100644 index 0000000..e5c361d --- /dev/null +++ b/internal/proxy/retry.go @@ -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. diff --git a/main.go b/main.go new file mode 100644 index 0000000..05243ad --- /dev/null +++ b/main.go @@ -0,0 +1,11 @@ +package main + +import ( + "os" + + "github.com/Atte149/ollama-proxy/internal/cli" +) + +func main() { + os.Exit(cli.Run(os.Args)) +} diff --git a/systemd/ollama-proxy.service b/systemd/ollama-proxy.service new file mode 100644 index 0000000..434038f --- /dev/null +++ b/systemd/ollama-proxy.service @@ -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 \ No newline at end of file