diff --git a/README.md b/README.md index a9f85e3..2d0d467 100644 --- a/README.md +++ b/README.md @@ -89,14 +89,18 @@ reasoning encrypted, so there only the token counts show up. auth: bearer `sk-lanv-...` from `key add`. fresh data dir with zero keys runs open (single-user localhost); the moment one key exists, all `/v1/*` -require it. one stable `x-opencode-session` is derived per api key: -minting a fresh session per request burns free quota into 429s. +require it. on a public endpoint set `requireKey` in the config (or +`LANVELLO_REQUIRE_KEY=1`) so an empty data dir refuses instead of serving +open. keys.json is re-read when it changes, so `key add` / `key revoke` +take effect without a restart. one stable `x-opencode-session` is derived +per api key: minting a fresh session per request burns free quota into 429s. ## limits -429 retires the lane exit until `retry-after` and rotates country, request is -retried on another lane. 502/503/504 retry on another lane. per-request -proof goes to `$dataDir/proof.jsonl`. +429 retires the lane exit until `retry-after`; with tor lanes the process is +respawned in the background for a fresh exit and the lane returns once it is +up, while the request is retried on another lane. 502/503/504 retry on +another lane. per-request proof goes to `$dataDir/proof.jsonl`. a lane stays busy for the whole answer, not just while the request is in flight, so N lanes means N concurrent answers; the rest wait in the queue up diff --git a/cmd/lanvello/main.go b/cmd/lanvello/main.go index efcdd26..e7c1098 100644 --- a/cmd/lanvello/main.go +++ b/cmd/lanvello/main.go @@ -116,6 +116,11 @@ func cmdServe(args []string) { // tor-only: direct egress is gone. without tor lanes there is no service. if len(cfg.Socks) == 0 { tor.Ensure(m, cfg.DataDir, 52001, 52301, 90*time.Second) + // A 429 rotates the lane: respawn its tor in the background so the + // next request lands on a fresh exit. + m.OnRotate = func(l *lanes.Lane) { + tor.Respawn(m, l, cfg.DataDir, 52001, 52301, 90*time.Second) + } } if n := healthySocks(m); n == 0 { fmt.Fprintln(os.Stderr, "no tor lanes up (tor binary missing or bootstrap failed) -- refusing to serve direct") @@ -132,6 +137,7 @@ func cmdServe(args []string) { up.WantCountry[pref] = cc } srv := server.New(ks, m, up, cfg.DataDir) + srv.RequireKey = cfg.RequireKey fmt.Printf("lanvello %s listening on http://%s lanes=%d tor-only, free, no login\n", version, cfg.Listen, cfg.Lanes) if err := http.ListenAndServe(cfg.Listen, srv.Handler()); err != nil { fmt.Fprintln(os.Stderr, err) diff --git a/internal/config/config.go b/internal/config/config.go index f0f1951..e193015 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -5,6 +5,7 @@ import ( "encoding/json" "os" "path/filepath" + "strings" ) type Config struct { @@ -21,6 +22,9 @@ type Config struct { // CatalogBudgetS bounds the lane wait for /v1/models, which must stay // snappy even when every lane is busy with a long answer. CatalogBudgetS float64 `json:"catalogBudgetS"` + // RequireKey refuses keyless /v1 requests even before the first key is + // issued: a public deployment must never fall back to open mode. + RequireKey bool `json:"requireKey"` } func DefaultDataDir() string { @@ -44,9 +48,18 @@ func Defaults() Config { Countries: []string{"us", "de", "nl", "fr", "ro", "gb", "ca", "se", "pl", "ch"}, WaitBudgetS: 90, CatalogBudgetS: 20, + RequireKey: envBool("LANVELLO_REQUIRE_KEY"), } } +func envBool(k string) bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(k))) { + case "1", "true", "yes", "on": + return true + } + return false +} + func envOr(k, d string) string { if v := os.Getenv(k); v != "" { return v @@ -100,6 +113,9 @@ func LoadFile(path string, base Config) (Config, error) { if f.CatalogBudgetS > 0 { base.CatalogBudgetS = f.CatalogBudgetS } + if f.RequireKey { + base.RequireKey = true + } if len(f.ModelLanes) > 0 { if base.ModelLanes == nil { base.ModelLanes = map[string]string{} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d89fc35..119fba9 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -73,3 +73,17 @@ func TestLoadFileRejectsGarbage(t *testing.T) { t.Fatal("garbage config must be an error, not silently defaulted") } } + +func TestLoadFileRequireKey(t *testing.T) { + p := filepath.Join(t.TempDir(), "c.json") + if err := os.WriteFile(p, []byte(`{"requireKey":true}`), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFile(p, Defaults()) + if err != nil { + t.Fatal(err) + } + if !cfg.RequireKey { + t.Fatalf("requireKey lost: %+v", cfg) + } +} diff --git a/internal/keys/keys.go b/internal/keys/keys.go index 5fde3ea..8fe7ad9 100644 --- a/internal/keys/keys.go +++ b/internal/keys/keys.go @@ -26,6 +26,11 @@ type Store struct { path string mu sync.Mutex en []Entry + // mtime/size of the keys.json the in-memory list came from. A change on + // disk (lanvello key add/revoke runs in another process) is picked up on + // the next request instead of requiring a restart. + mtime int64 + size int64 } func Path(dataDir string) string { return filepath.Join(dataDir, "keys.json") } @@ -48,9 +53,42 @@ func Open(dataDir string) (*Store, error) { if err := json.Unmarshal(b, &s.en); err != nil { return nil, err } + s.stamp() return s, nil } +// stamp records the current file identity so maybeReload can tell later +// writes apart. Best-effort: a missing file simply never matches. +func (s *Store) stamp() { + if st, err := os.Stat(s.path); err == nil { + s.mtime, s.size = st.ModTime().UnixNano(), st.Size() + } +} + +// maybeReload re-reads keys.json when another process changed it. Callers +// hold s.mu. Best-effort: a torn or unparsable read keeps the in-memory list. +func (s *Store) maybeReload() { + st, err := os.Stat(s.path) + if err != nil { + return + } + if st.ModTime().UnixNano() == s.mtime && st.Size() == s.size { + return + } + b, err := os.ReadFile(s.path) + if err != nil { + return + } + var en []Entry + if len(b) > 0 { + if err := json.Unmarshal(b, &en); err != nil { + return + } + } + s.en = en + s.mtime, s.size = st.ModTime().UnixNano(), st.Size() +} + func (s *Store) save() error { tmp := s.path + ".tmp" b, err := json.MarshalIndent(s.en, "", " ") @@ -60,7 +98,11 @@ func (s *Store) save() error { if err := os.WriteFile(tmp, b, 0o600); err != nil { return err } - return os.Rename(tmp, s.path) + if err := os.Rename(tmp, s.path); err != nil { + return err + } + s.stamp() + return nil } func hashOf(secret string) string { @@ -72,6 +114,7 @@ func hashOf(secret string) string { func (s *Store) Add(name string) (string, Entry, error) { s.mu.Lock() defer s.mu.Unlock() + s.maybeReload() var raw [24]byte if _, err := rand.Read(raw[:]); err != nil { return "", Entry{}, err @@ -100,6 +143,7 @@ func (s *Store) Add(name string) (string, Entry, error) { func (s *Store) List() []Entry { s.mu.Lock() defer s.mu.Unlock() + s.maybeReload() cp := make([]Entry, len(s.en)) copy(cp, s.en) return cp @@ -108,6 +152,7 @@ func (s *Store) List() []Entry { func (s *Store) Revoke(name string) bool { s.mu.Lock() defer s.mu.Unlock() + s.maybeReload() found := false out := s.en[:0] for _, x := range s.en { @@ -129,6 +174,7 @@ func (s *Store) Revoke(name string) bool { func (s *Store) Verify(secret string, allowEmpty bool) bool { s.mu.Lock() defer s.mu.Unlock() + s.maybeReload() if secret == "" { return false } diff --git a/internal/keys/keys_test.go b/internal/keys/keys_test.go index dca2d9e..d294794 100644 --- a/internal/keys/keys_test.go +++ b/internal/keys/keys_test.go @@ -29,3 +29,37 @@ func TestAddVerifyRevoke(t *testing.T) { t.Fatal("revoked still verifies") } } + +// The running server must see keys added or revoked by the CLI (`lanvello +// key add`, a separate process) without a restart. +func TestStorePicksUpExternalChanges(t *testing.T) { + dir := t.TempDir() + srv, err := keys.Open(dir) + if err != nil { + t.Fatal(err) + } + cli, err := keys.Open(dir) + if err != nil { + t.Fatal(err) + } + sec, _, err := cli.Add("atractio") + if err != nil { + t.Fatal(err) + } + if !srv.Verify(sec, false) { + t.Fatal("a key added by another process must verify without a restart") + } + var names []string + for _, e := range srv.List() { + names = append(names, e.Name) + } + if len(names) != 1 || names[0] != "atractio" { + t.Fatalf("list=%v", names) + } + if !cli.Revoke("atractio") { + t.Fatal("revoke failed") + } + if srv.Verify(sec, false) { + t.Fatal("a key revoked by another process must stop verifying") + } +} diff --git a/internal/lanes/lanes.go b/internal/lanes/lanes.go index f954ccd..601b67f 100644 --- a/internal/lanes/lanes.go +++ b/internal/lanes/lanes.go @@ -26,6 +26,7 @@ type Lane struct { Score int Healthy bool Healing bool + Rotating bool LastUsed time.Time LastReal time.Time cmd *exec.Cmd @@ -52,6 +53,10 @@ type Manager struct { fallback []string noTor bool seq int + // OnRotate, when set (tor-spawned lanes), respawns the lane's process so + // a 429 lands on a fresh exit. The lane stays retired until RotateDone + // reports the new exit. + OnRotate func(*Lane) } func New(dataDir string, countries, fallback []string, socks []string, noTor bool, count int) *Manager { @@ -182,7 +187,7 @@ func (m *Manager) pick(exclude map[int]bool, cc string) *Lane { if exclude != nil && exclude[l.Index] { continue } - if !l.Healthy || l.Healing { + if !l.Healthy || l.Healing || l.Rotating { continue } if cc != "" && l.Country != cc { @@ -265,10 +270,12 @@ func (m *Manager) NoteLimited(l *Lane, retryAfter time.Duration) time.Time { return l.LimitedTill } -// Rotate moves lane to the best country with spare capacity (simple score sort). +// Rotate retires the lane's exit after a 429 and asks OnRotate (tor respawn, +// when the lanes are tor-spawned) for a fresh one in the best-scoring +// country. The retry-after limit is kept until the new exit is up — clearing +// it here handed the same limited exit straight back to the next request. func (m *Manager) Rotate(l *Lane) string { m.mu.Lock() - defer m.mu.Unlock() pool := append(append([]string{}, m.countries...), m.fallback...) seen := map[string]bool{} var ranked []string @@ -291,12 +298,51 @@ func (m *Manager) Rotate(l *Lane) string { continue } l.Country = c - l.LimitedTill = time.Time{} - return c + break + } + hook := m.OnRotate + if hook != nil && !l.Rotating { + l.Rotating = true + go hook(l) + } + cc := l.Country + m.mu.Unlock() + return cc +} + +// RotateDone reports a lane's fresh exit (a respawn finished): the lane is +// usable again and its retry-after limit is lifted. +func (m *Manager) RotateDone(l *Lane, socks, ip string) { + m.mu.Lock() + defer m.mu.Unlock() + if socks != "" { + l.SocksAddr = socks + l.Healthy = true } - return l.Country + l.ExitIP = ip + l.Rotating = false + l.LimitedTill = time.Time{} } +// RotateFailed clears the rotating flag after a failed respawn. The lane +// keeps its retry-after limit (with a cooldown floor) so the next request +// does not immediately try to rotate it again. +func (m *Manager) RotateFailed(l *Lane) { + m.mu.Lock() + defer m.mu.Unlock() + l.Rotating = false + if floor := time.Now().Add(5 * time.Minute); l.LimitedTill.Before(floor) { + l.LimitedTill = floor + } +} + +// SetProcess records the tor process backing this lane so rotation can +// replace it; nil for lanes that are not tor-spawned. +func (l *Lane) SetProcess(cmd *exec.Cmd) { l.cmd = cmd } + +// Process returns the tor process backing this lane, if any. +func (l *Lane) Process() *exec.Cmd { return l.cmd } + func (m *Manager) Emit(p Proof) { m.mu.Lock() path := m.proofPath diff --git a/internal/lanes/lanes_test.go b/internal/lanes/lanes_test.go index 5f263e1..6e313ba 100644 --- a/internal/lanes/lanes_test.go +++ b/internal/lanes/lanes_test.go @@ -2,6 +2,7 @@ package lanes_test import ( "testing" + "time" "lanvello/internal/lanes" ) @@ -44,3 +45,48 @@ func TestLimitedLanesDeprioritized(t *testing.T) { } m.Release(got) } + +// A 429 rotation must not hand the same limited exit back: the retry-after +// limit stays until the respawned exit is up. +func TestRotateKeepsRetryAfter(t *testing.T) { + m := lanes.New(t.TempDir(), []string{"de", "nl"}, nil, []string{"127.0.0.1:1111", "127.0.0.1:2222"}, true, 2) + ls := m.Lanes() + m.NoteLimited(ls[0], 10*time.Minute) + if ls[0].LimitedTill.IsZero() { + t.Fatal("limit not set") + } + m.Rotate(ls[0]) + if ls[0].LimitedTill.IsZero() { + t.Fatal("rotation cleared the retry-after limit") + } + got := m.Pick(nil) + if got == nil || got.Index == ls[0].Index { + t.Fatalf("limited lane picked after rotate: %+v", got) + } + m.Release(got) +} + +func TestRotateHookAndDone(t *testing.T) { + m := lanes.New(t.TempDir(), []string{"de", "nl"}, nil, []string{"127.0.0.1:1111", "127.0.0.1:2222"}, true, 2) + ls := m.Lanes() + called := make(chan int, 1) + m.OnRotate = func(l *lanes.Lane) { called <- l.Index } + m.NoteLimited(ls[0], 0) + m.Rotate(ls[0]) + select { + case idx := <-called: + if idx != ls[0].Index { + t.Fatalf("hook lane=%d", idx) + } + case <-time.After(2 * time.Second): + t.Fatal("OnRotate hook never fired") + } + // while rotating, the lane is out of the pool + if got := m.Pick(map[int]bool{ls[1].Index: true}); got != nil { + t.Fatalf("rotating lane picked: %+v", got) + } + m.RotateDone(ls[0], "127.0.0.1:3333", "203.0.113.7") + if ls[0].Rotating || !ls[0].LimitedTill.IsZero() || ls[0].SocksAddr != "127.0.0.1:3333" { + t.Fatalf("RotateDone not applied: %+v", ls[0]) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 42983a7..73a14e4 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -27,6 +27,10 @@ type Server struct { Lane *lanes.Manager Up *upstream.Client Caps *caps.Store + // RequireKey refuses keyless /v1 requests even before the first key is + // issued (config `requireKey` / LANVELLO_REQUIRE_KEY): a public + // deployment must never fall back to open mode. + RequireKey bool catMu sync.Mutex catCache []modelEnt @@ -77,7 +81,7 @@ func (s *Server) identity(r *http.Request) string { func (s *Server) authed(r *http.Request) bool { if len(s.Keys.List()) == 0 { - return true // open mode: no keys issued yet + return !s.RequireKey // open mode: no keys issued yet } sec := keys.BearerOf(r.Header.Get("Authorization")) if sec == "" { @@ -260,11 +264,15 @@ func shortID(id string) string { } func readBody(w http.ResponseWriter, r *http.Request) ([]byte, bool) { - raw, err := io.ReadAll(io.LimitReader(r.Body, maxBody)) + raw, err := io.ReadAll(io.LimitReader(r.Body, maxBody+1)) if err != nil { http.Error(w, "read body", 400) return nil, false } + if int64(len(raw)) > maxBody { + http.Error(w, "body too large", http.StatusRequestEntityTooLarge) + return nil, false + } return raw, true } @@ -427,8 +435,12 @@ func (s *Server) responses(w http.ResponseWriter, r *http.Request) { if !s.known(model) { model = shortID(defaultModel) } + // Read the client's intent BEFORE fingerprintResponses forces stream=true + // on this same map — otherwise every stream:false request took the relay + // branch below and got raw sse back. + wantStream, _ := body["stream"].(bool) fwd := s.fingerprintResponses(body, model) - if stream, _ := body["stream"].(bool); stream { + if wantStream { b, err := json.Marshal(fwd) if err != nil { http.Error(w, "encode", 500) @@ -446,6 +458,17 @@ func (s *Server) responses(w http.ResponseWriter, r *http.Request) { return } status, hdr, respBody := s.Up.Raw(model, s.identity(r), b, r.Header) + // A stream:false client asked for one response object; upstream still + // answers sse (the fingerprint forces streaming), so fold the terminal + // event back into the response JSON instead of leaking sse bytes. + if status == 200 && isSSE(hdr, respBody) { + if agg, ok := translate.AggregateResponsesObject(respBody); ok { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(200) + _, _ = w.Write(agg) + return + } + } pass(w, status, hdr, respBody) } diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 694b020..6d39869 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -553,6 +553,34 @@ func TestResponsesNativeStreamRelayed(t *testing.T) { } } +// A stream:false native responses client gets one response object, not sse. +func TestResponsesNonStreamAggregates(t *testing.T) { + h := newHarness(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/zen/v1/responses" { + t.Errorf("path=%s", r.URL.Path) + } + sseWriter(w, responsesTextStream("Hello")...) + }) + rec := h.do(t, "POST", "/v1/responses", map[string]any{ + "model": "muse-spark-1.3-contributor-free", "stream": false, + "input": "hi", + }) + if rec.Code != 200 { + t.Fatalf("code=%d body=%s", rec.Code, rec.Body.String()) + } + body := rec.Body.String() + if strings.Contains(body, "event:") || strings.Contains(body, "data:") { + t.Fatalf("sse leaked to a stream:false client:\n%s", body) + } + var v map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &v); err != nil { + t.Fatal(body) + } + if v["id"] != "resp_abc" { + t.Fatalf("id=%v", v["id"]) + } +} + func TestUpstream429RotatesBeforeFirstByte(t *testing.T) { var n int lm := lanes.New(t.TempDir(), []string{"us", "de"}, nil, nil, true, 2) @@ -597,10 +625,10 @@ func TestBodyLimit(t *testing.T) { req.Header.Set("Authorization", "Bearer "+h.key) rec := httptest.NewRecorder() h.s.Handler().ServeHTTP(rec, req) - // 200MB+1 of json is not a valid short request: whatever happens, the - // process must survive and answer with a status, not hang or crash. - if rec.Code == 0 { - t.Fatal("no status written") + // 200MB of json is over the 128MB cap: it must be refused with a real + // status, not silently truncated. + if rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("oversized body must be refused with 413, got %d", rec.Code) } } @@ -625,6 +653,34 @@ func TestBadKeyRejected(t *testing.T) { } } +// With no keys issued yet the API defaults to open mode (localhost dev); a +// public deployment sets requireKey and an empty store must refuse instead. +func TestRequireKeyClosesOpenMode(t *testing.T) { + h := newHarness(t, func(w http.ResponseWriter, r *http.Request) {}) + dir := t.TempDir() + ks, _ := keys.Open(dir) + lm := lanes.New(dir, []string{"de"}, nil, nil, true, 1) + cl := upstream.NewClient(lm) + cl.BaseURL = h.srv.URL + cl.WaitBudget = 2 * time.Second + s := server.New(ks, lm, cl, "") + + req := httptest.NewRequest("GET", "/v1/models", nil) + rec := httptest.NewRecorder() + s.Handler().ServeHTTP(rec, req) + if rec.Code != 200 { + t.Fatalf("default open mode must still answer: %d", rec.Code) + } + + s.RequireKey = true + rec = httptest.NewRecorder() + req = httptest.NewRequest("GET", "/v1/models", nil) + s.Handler().ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("requireKey must refuse keyless requests: %d", rec.Code) + } +} + // streaming must not depend on the response recorder: real http.Flusher path. func TestStreamFlushesOnRealServer(t *testing.T) { release := make(chan struct{}) diff --git a/internal/tor/tor.go b/internal/tor/tor.go index 53f23df..8454ff6 100644 --- a/internal/tor/tor.go +++ b/internal/tor/tor.go @@ -332,6 +332,7 @@ func Ensure(m *lanes.Manager, dataDir string, base, ctrlBase int, timeout time.D if err := cmd.Start(); err != nil { continue } + l.SetProcess(cmd) addr := fmt.Sprintf("127.0.0.1:%d", socks) if waitPort(addr, timeout) { l.SocksAddr = addr @@ -340,3 +341,40 @@ func Ensure(m *lanes.Manager, dataDir string, base, ctrlBase int, timeout time.D } } } + +// Respawn restarts one lane's tor process with its current (rotated) country +// and hands the fresh exit back to the manager. Called from the manager's +// OnRotate hook when a 429 retires the lane. +func Respawn(m *lanes.Manager, l *lanes.Lane, dataDir string, base, ctrlBase int, timeout time.Duration) { + bin := FindTor() + if bin == "" { + m.RotateFailed(l) + return + } + if old := l.Process(); old != nil && old.Process != nil { + _ = old.Process.Kill() + _, _ = old.Process.Wait() + } + socks := base + l.Index - 1 + ctrl := ctrlBase + l.Index - 1 + dir := filepath.Join(dataDir, "lanes", fmt.Sprintf("tor-%d", l.Index)) + if err := writeTorrc(bin, dir, socks, ctrl, l.Country); err != nil { + m.RotateFailed(l) + return + } + env := append(os.Environ(), "LD_LIBRARY_PATH="+filepath.Dir(bin)) + cmd := exec.Command(bin, "-f", filepath.Join(dir, "torrc")) + cmd.Env = env + if err := cmd.Start(); err != nil { + m.RotateFailed(l) + return + } + l.SetProcess(cmd) + addr := fmt.Sprintf("127.0.0.1:%d", socks) + if !waitPort(addr, timeout) { + _ = cmd.Process.Kill() + m.RotateFailed(l) + return + } + m.RotateDone(l, addr, exitIP(addr)) +} diff --git a/internal/translate/translate.go b/internal/translate/translate.go index 39006b8..ea10d22 100644 --- a/internal/translate/translate.go +++ b/internal/translate/translate.go @@ -553,6 +553,25 @@ func AggregateResponses(sse []byte, model string) []byte { return mustJSON(resp) } +// AggregateResponsesObject folds a native responses sse stream into the one +// response object a stream:false client expects: the terminal event carries +// the full response. ok=false when the stream never completed. +func AggregateResponsesObject(sse []byte) ([]byte, bool) { + var last map[string]any + forEachPayload(sse, func(ev map[string]any) { + switch t, _ := ev["type"].(string); t { + case "response.completed", "response.failed", "response.incomplete": + if r, ok := ev["response"].(map[string]any); ok { + last = r + } + } + }) + if last == nil { + return nil, false + } + return mustJSON(last), true +} + // AggregateOpenAI collects openai sse chunks into one chat completion. func AggregateOpenAI(sse []byte, model string) []byte { var sb, reason strings.Builder diff --git a/internal/upstream/loop_test.go b/internal/upstream/loop_test.go index 9be1ba8..06d5070 100644 --- a/internal/upstream/loop_test.go +++ b/internal/upstream/loop_test.go @@ -3,6 +3,7 @@ package upstream import ( "bufio" "fmt" + "io" "net" "net/http" "net/http/httptest" @@ -231,3 +232,87 @@ func TestIdleBodyFailsStalledRead(t *testing.T) { t.Fatalf("idle guard too slow: %s", time.Since(start)) } } + +// A stream that never starts (no lane up, upstream unreachable) must answer +// with a real status. It used to return without writing anything, which +// net/http turned into an empty 200 that hid the failure. +func TestStreamAnswers502WhenNothingStarts(t *testing.T) { + m := lanes.New(t.TempDir(), []string{"us"}, nil, nil, true, 1) + c := NewClient(m) + c.BaseURL = "http://127.0.0.1:1" // nothing listens here + c.WaitBudget = 300 * time.Millisecond + hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if _, err := c.StreamEvents("fledge-alpha-free", "k", []byte(`{"messages":[]}`), nil, w, + func(payload string, ev map[string]any) ([]string, bool) { return nil, false }); err == nil { + t.Errorf("a dead upstream must surface an error") + } + })) + defer hs.Close() + resp, err := http.Get(hs.URL) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusBadGateway { + t.Fatalf("status=%d, want 502 (an empty 200 hides the failure)", resp.StatusCode) + } + b, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(b), "error") { + t.Fatalf("body=%q", b) + } +} + +// A far-end that is plainly down (503 every time) must not spin until the +// wait budget: after a couple of full passes the real answer goes back. +func TestRawBailsOnPersistent503(t *testing.T) { + var hits int32 + c, _ := testClient(t, 3, func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(503) + fmt.Fprint(w, `{"error":{"type":"server_error","message":"Endpoint is unavailable."}}`) + }) + c.WaitBudget = 60 * time.Second // would spin for a minute without the bound + start := time.Now() + status, _, body := c.Raw("fledge-alpha-free", "k", []byte(`{"messages":[]}`), nil) + if status != 503 { + t.Fatalf("status=%d, want 503", status) + } + if !strings.Contains(string(body), "Endpoint is unavailable") { + t.Fatalf("real upstream error lost: %q", string(body)) + } + if n := atomic.LoadInt32(&hits); n > 8 { + t.Fatalf("retried %d times, the bound is not working", n) + } + if d := time.Since(start); d > 15*time.Second { + t.Fatalf("spun %s instead of bailing", d) + } +} + +// Same for the streaming path: the drained upstream answer must be relayed +// to the client instead of an empty 200. +func TestStreamRelaysPersistent503(t *testing.T) { + c, _ := testClient(t, 3, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(503) + fmt.Fprint(w, `{"error":{"type":"server_error","message":"Endpoint is unavailable."}}`) + }) + c.WaitBudget = 60 * time.Second + hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = c.StreamEvents("fledge-alpha-free", "k", []byte(`{"messages":[]}`), nil, w, + func(payload string, ev map[string]any) ([]string, bool) { return nil, false }) + })) + defer hs.Close() + resp, err := http.Get(hs.URL) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != 503 { + t.Fatalf("status=%d, want 503", resp.StatusCode) + } + b, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(b), "Endpoint is unavailable") { + t.Fatalf("body=%q", b) + } +} diff --git a/internal/upstream/upstream.go b/internal/upstream/upstream.go index a2ef03f..73e6fb6 100644 --- a/internal/upstream/upstream.go +++ b/internal/upstream/upstream.go @@ -457,15 +457,23 @@ func (e *StreamError) Error() string { return e.Err.Error() } // another lane before a single byte is written downstream. func (c *Client) stream(model, identity string, fwd []byte, in http.Header, w http.ResponseWriter, body func(sink *streamSink, resp *http.Response) error) (int, error) { var sink *streamSink + committed := false + // The last retryable answer is kept so a loop that gives up (a plainly + // down endpoint) relays the real status instead of an empty 200. + var lastSt int + var lastHdr http.Header + var lastBody []byte status, err := c.loop(model, identity, "POST", endpointFor(model), fwd, in, func(resp *http.Response, st int, hdr http.Header) (bool, error) { if retryable(st) { - _, _ = drain(resp) + lastBody, _ = drain(resp) + lastSt, lastHdr = st, hdr return false, nil } if st != 200 { b, _ := drain(resp) relay(w, st, hdr, b) + committed = true return true, nil } w.Header().Set("Content-Type", "text/event-stream") @@ -473,10 +481,28 @@ func (c *Client) stream(model, identity string, fwd []byte, in http.Header, w ht w.Header().Set("Connection", "keep-alive") w.Header().Set("X-Accel-Buffering", "no") w.WriteHeader(200) + committed = true sink = &streamSink{w: w} sink.flush() return true, body(sink, resp) }) + if err == nil && !committed && lastSt != 0 { + // The retries were spent on a far-end failure: hand the drained + // upstream answer back instead of nothing. + relay(w, lastSt, lastHdr, lastBody) + return lastSt, nil + } + if err != nil && !committed { + // Nothing was written downstream yet (no lane up, upstream down): + // answer with a real status. Returning without a body let net/http + // send an empty 200, hiding the failure from the client. + st := http.StatusBadGateway + if se, ok := err.(*StreamError); ok && se.Status != 0 { + st = se.Status + } + relay(w, st, jsonHeader(), []byte(`{"error":`+quote(err.Error())+`}`)) + return st, err + } if err != nil { if se, ok := err.(*StreamError); ok { return se.Status, se.Err @@ -542,6 +568,7 @@ func (c *Client) loopBudget(model, identity, method, path string, fwd []byte, in reqID := deriveRequestID(session, lastUserText(probe)) tried := map[int]bool{} + farFails := 0 deadline := time.Now().Add(budget) var lastErr string wantCC := c.countryFor(model) @@ -583,6 +610,13 @@ func (c *Client) loopBudget(model, identity, method, path string, fwd []byte, in if stop { return st, cerr } + // Far-end failures earn another exit, but a model whose + // endpoint is plainly down must not spin until the deadline: + // after two full passes the real status goes to the caller. + farFails++ + if farFails >= 2*len(c.Lanes.Lanes()) { + return st, cerr + } continue case st == 500: // Chat-path 500s are transient far-end failures; rotate once.