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
241
internal/cli/integration_test.go
Normal file
241
internal/cli/integration_test.go
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// findBinary locates the ollama-proxy executable built by `go test`. The test
|
||||
// binary is in the package dir; the project binary is one level up. We try both
|
||||
// and fall back to building one on the fly.
|
||||
func findBinary(t *testing.T) string {
|
||||
t.Helper()
|
||||
candidates := []string{
|
||||
filepath.Join("..", "..", "ollama-proxy"),
|
||||
"ollama-proxy",
|
||||
}
|
||||
for _, c := range candidates {
|
||||
if abs, err := filepath.Abs(c); err == nil {
|
||||
if fi, err := os.Stat(abs); err == nil && !fi.IsDir() {
|
||||
return abs
|
||||
}
|
||||
}
|
||||
}
|
||||
// Build it now into a temp path.
|
||||
out := filepath.Join(t.TempDir(), "ollama-proxy")
|
||||
cmd := exec.Command("go", "build", "-o", out, ".")
|
||||
cmd.Dir = filepath.Join("..", "..")
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
t.Fatalf("go build: %v\n%s", err, out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// freePort picks an unused TCP port for the proxy to listen on.
|
||||
func freePort(t *testing.T) string {
|
||||
t.Helper()
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer l.Close()
|
||||
return l.Addr().String()
|
||||
}
|
||||
|
||||
// startProxy starts the binary as a subprocess and returns a cancel function
|
||||
// that stops it.
|
||||
func startProxy(t *testing.T, bin, xdgDir, addr, upstream string) func() {
|
||||
t.Helper()
|
||||
cmd := exec.Command(bin, "serve",
|
||||
"--addr", addr,
|
||||
"--base-url", upstream,
|
||||
"--cooldown", "1s",
|
||||
"--retries", "3",
|
||||
"--log-level", "debug",
|
||||
)
|
||||
cmd.Env = append(os.Environ(),
|
||||
"XDG_CONFIG_HOME="+xdgDir,
|
||||
"OLLAMA_PROXY_LOG_LEVEL=debug",
|
||||
)
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start proxy: %v", err)
|
||||
}
|
||||
// Wait for the port to accept connections.
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
c, err := net.Dial("tcp", addr)
|
||||
if err == nil {
|
||||
c.Close()
|
||||
break
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
return func() {
|
||||
_ = cmd.Process.Signal(os.Interrupt)
|
||||
done := make(chan struct{})
|
||||
go func() { _ = cmd.Wait(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
_ = cmd.Process.Kill()
|
||||
<-done
|
||||
}
|
||||
t.Logf("proxy stdout: %s", stdout.String())
|
||||
t.Logf("proxy stderr: %s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestEndToEnd_ProxyForwardsToUpstream builds the binary and verifies the full
|
||||
// stack: account JSON → balancer → handler → real HTTP server → upstream.
|
||||
func TestEndToEnd_ProxyForwardsToUpstream(t *testing.T) {
|
||||
// Upstream that echoes the auth and path.
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = fmt.Fprintf(w, `{"path":%q,"auth":%q}`, r.URL.Path, r.Header.Get("Authorization"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
bin := findBinary(t)
|
||||
xdg := t.TempDir()
|
||||
// Write accounts.json directly via the binary's `accounts add`.
|
||||
for _, a := range []struct{ key, name string }{
|
||||
{"sk-keyAAAAAAAA", "a1"},
|
||||
{"sk-keyBBBBBBBB", "a2"},
|
||||
} {
|
||||
cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
|
||||
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
|
||||
}
|
||||
}
|
||||
|
||||
addr := freePort(t)
|
||||
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||
defer stop()
|
||||
|
||||
resp, err := http.Get("http://" + addr + "/api/version")
|
||||
if err != nil {
|
||||
t.Fatalf("GET /api/version: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("status = %d, body = %s", resp.StatusCode, body)
|
||||
}
|
||||
if !strings.Contains(string(body), "sk-key") {
|
||||
t.Errorf("body = %s, want account key forwarded", body)
|
||||
}
|
||||
if !strings.Contains(string(body), "/api/version") {
|
||||
t.Errorf("body = %s, want path echoed", body)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEndToEnd_SSEStreaming verifies SSE passthrough through the real binary.
|
||||
func TestEndToEnd_SSEStreaming(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Transfer-Encoding", "chunked")
|
||||
w.WriteHeader(200)
|
||||
flusher, _ := w.(http.Flusher)
|
||||
for i := 0; i < 3; i++ {
|
||||
_, _ = io.WriteString(w, "data: {\"i\":"+string(rune('0'+i))+"}\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
bin := findBinary(t)
|
||||
xdg := t.TempDir()
|
||||
cmd := exec.Command(bin, "accounts", "add", "sk-testkey", "--name", "a1")
|
||||
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
t.Fatalf("accounts add: %v\n%s", err, out)
|
||||
}
|
||||
|
||||
addr := freePort(t)
|
||||
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||
defer stop()
|
||||
|
||||
resp, err := http.Post("http://"+addr+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
|
||||
if err != nil {
|
||||
t.Fatalf("POST: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if !strings.Contains(string(body), "[DONE]") {
|
||||
t.Errorf("body = %s, want [DONE]", body)
|
||||
}
|
||||
if !strings.Contains(string(body), `"i":0`) || !strings.Contains(string(body), `"i":1`) {
|
||||
t.Errorf("body = %s, missing chunks", body)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEndToEnd_FailoverOn429 verifies the binary fails over between accounts.
|
||||
func TestEndToEnd_FailoverOn429(t *testing.T) {
|
||||
// Upstream that 429s for the first account key and 200s for the second.
|
||||
var firstSeen string
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
auth := r.Header.Get("Authorization")
|
||||
if firstSeen == "" {
|
||||
firstSeen = auth
|
||||
}
|
||||
if auth == "Bearer sk-bad" {
|
||||
http.Error(w, "rate limited", 429)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
_, _ = io.WriteString(w, `{"ok":true}`)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
bin := findBinary(t)
|
||||
xdg := t.TempDir()
|
||||
for _, a := range []struct{ key, name string }{
|
||||
{"sk-bad", "bad"},
|
||||
{"sk-good", "good"},
|
||||
} {
|
||||
cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
|
||||
cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
|
||||
}
|
||||
}
|
||||
|
||||
addr := freePort(t)
|
||||
stop := startProxy(t, bin, xdg, addr, upstream.URL)
|
||||
defer stop()
|
||||
|
||||
resp, err := http.Get("http://" + addr + "/api/version")
|
||||
if err != nil {
|
||||
t.Fatalf("GET: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("status = %d, body = %s (want 200 after failover)", resp.StatusCode, body)
|
||||
}
|
||||
if !strings.Contains(string(body), `"ok":true`) {
|
||||
t.Errorf("body = %s, want good response", body)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue