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
245
internal/proxy/handler_test.go
Normal file
245
internal/proxy/handler_test.go
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
package proxy
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Atte149/ollama-proxy/internal/config"
|
||||
"github.com/Atte149/ollama-proxy/internal/log"
|
||||
)
|
||||
|
||||
// mockUpstream is a configurable test upstream. It responds to /api/chat and
|
||||
// /v1/chat/completions; for other paths it returns 200 with a short body.
|
||||
type mockUpstream struct {
|
||||
status int32 // current status to return
|
||||
body string
|
||||
allowAuth string // if non-empty, require this Bearer token
|
||||
chunks []string
|
||||
requestAuth atomic.Int32 // number of requests seen with each auth value
|
||||
}
|
||||
|
||||
func (m *mockUpstream) handler(w http.ResponseWriter, r *http.Request) {
|
||||
auth := r.Header.Get("Authorization")
|
||||
if m.allowAuth != "" && auth != "Bearer "+m.allowAuth {
|
||||
http.Error(w, "bad auth", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if m.status >= 500 || m.status == 429 {
|
||||
http.Error(w, "rate limited", int(m.status))
|
||||
return
|
||||
}
|
||||
if len(m.chunks) > 0 {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Transfer-Encoding", "chunked")
|
||||
w.WriteHeader(200)
|
||||
flusher, _ := w.(http.Flusher)
|
||||
for _, c := range m.chunks {
|
||||
_, _ = io.WriteString(w, c)
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(int(m.status))
|
||||
_, _ = io.WriteString(w, m.body)
|
||||
}
|
||||
|
||||
func newTestHandler(t *testing.T, upstreamURL string, accounts []config.Account, retries int) *Handler {
|
||||
t.Helper()
|
||||
b := NewBalancer(accounts, 50*time.Millisecond)
|
||||
return NewHandler(b, upstreamURL, retries, log.New("debug"))
|
||||
}
|
||||
|
||||
func TestHandler_ProxiesVersion(t *testing.T) {
|
||||
up := &mockUpstream{status: 200, body: `{"version":"0.1.42"}`}
|
||||
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "0.1.42") {
|
||||
t.Errorf("body = %q, want version", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_FailoverOn429(t *testing.T) {
|
||||
// Two upstreams: first always 429, second always 200.
|
||||
badUp := &mockUpstream{status: 429, body: "rate limited"}
|
||||
badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
|
||||
goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
|
||||
goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
|
||||
defer badSrv.Close()
|
||||
defer goodSrv.Close()
|
||||
|
||||
// Two accounts, each pointing at a different upstream via per-account BaseURL.
|
||||
accts := []config.Account{
|
||||
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
|
||||
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
|
||||
}
|
||||
h := newTestHandler(t, "https://unused.example.com", accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200 (failover to good account)", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), `"ok":true`) {
|
||||
t.Errorf("body = %q, want good upstream response", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_AllAccounts429_Returns429(t *testing.T) {
|
||||
up := &mockUpstream{status: 429, body: "rate limited"}
|
||||
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||
defer srv.Close()
|
||||
|
||||
accts := []config.Account{
|
||||
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: srv.URL},
|
||||
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: srv.URL},
|
||||
}
|
||||
h := newTestHandler(t, srv.URL, accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want 429", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_SSEStreaming(t *testing.T) {
|
||||
up := &mockUpstream{
|
||||
status: 200,
|
||||
chunks: []string{"data: {\"a\":1}\n\n", "data: {\"a\":2}\n\n", "data: [DONE]\n\n"},
|
||||
}
|
||||
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"stream":true}`))
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
body := rec.Body.String()
|
||||
if !strings.Contains(body, "[DONE]") {
|
||||
t.Errorf("body = %q, want to contain [DONE]", body)
|
||||
}
|
||||
if !strings.Contains(body, "\"a\":1") || !strings.Contains(body, "\"a\":2") {
|
||||
t.Errorf("body = %q, missing chunks", body)
|
||||
}
|
||||
if ct := rec.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
|
||||
t.Errorf("Content-Type = %q, want text/event-stream", ct)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_RejectsUnknownPath(t *testing.T) {
|
||||
up := &mockUpstream{status: 200, body: "x"}
|
||||
srv := httptest.NewServer(http.HandlerFunc(up.handler))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHandler(t, srv.URL, mkAccts(1), 1)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/admin", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Errorf("status = %d, want 404 for unknown path", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_NoAccounts(t *testing.T) {
|
||||
h := newTestHandler(t, "http://unused.example.com", nil, 1)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusServiceUnavailable {
|
||||
t.Errorf("status = %d, want 503", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_StripsClientAuthorization(t *testing.T) {
|
||||
// Upstream echoes the received Authorization header in its body so we can
|
||||
// assert the proxy overrode it with the account's key.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"got_auth":"`+r.Header.Get("Authorization")+`"}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
accts := []config.Account{{ID: "idA", Name: "a", APIKey: "sk-account-key-1234", BaseURL: srv.URL}}
|
||||
h := newTestHandler(t, srv.URL, accts, 1)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
req.Header.Set("Authorization", "Bearer client-should-be-stripped")
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if !strings.Contains(rec.Body.String(), "sk-account-key-1234") {
|
||||
t.Errorf("body = %q, want account key forwarded", rec.Body.String())
|
||||
}
|
||||
if strings.Contains(rec.Body.String(), "client-should-be-stripped") {
|
||||
t.Errorf("body = %q, client Authorization leaked to upstream", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_RetriesOn5xx(t *testing.T) {
|
||||
// First account returns 500, second returns 200.
|
||||
badUp := &mockUpstream{status: 500}
|
||||
badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
|
||||
goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
|
||||
goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
|
||||
defer badSrv.Close()
|
||||
defer goodSrv.Close()
|
||||
|
||||
accts := []config.Account{
|
||||
{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
|
||||
{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
|
||||
}
|
||||
h := newTestHandler(t, "https://unused.example.com", accts, 2)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/api/version", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("status = %d, want 200 (failover on 500)", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsAllowed(t *testing.T) {
|
||||
good := []string{
|
||||
"/api/chat", "/api/generate", "/api/tags", "/api/show", "/api/ps",
|
||||
"/api/version", "/api/delete",
|
||||
"/v1/chat/completions", "/v1/completions", "/v1/models",
|
||||
"/v1/embeddings", "/v1/files", "/v1/files/abc-123",
|
||||
}
|
||||
for _, p := range good {
|
||||
if !isAllowed(p) {
|
||||
t.Errorf("isAllowed(%q) = false, want true", p)
|
||||
}
|
||||
}
|
||||
bad := []string{"/admin", "/", "/api/unknown", "/v2/foo", "/etc/passwd"}
|
||||
for _, p := range bad {
|
||||
if isAllowed(p) {
|
||||
t.Errorf("isAllowed(%q) = true, want false", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue