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
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)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue