diff --git a/internal/gateway/anthropic.go b/internal/gateway/anthropic.go index dfcef5cd..074ecfa4 100644 --- a/internal/gateway/anthropic.go +++ b/internal/gateway/anthropic.go @@ -691,6 +691,12 @@ func (e *anthropicEncoder) event(ev Event) { e.col.add(ev) } +// keepalive is Anthropic's own: a ping event, which its SDKs skip and +// Claude Code's stream watchdog counts. +func (e *anthropicEncoder) keepalive() { + e.w.event("ping", map[string]any{"type": "ping"}) +} + func (e *anthropicEncoder) finish() { if !e.started { e.start(Event{}) diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index b45b5ae3..29e40cdf 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -807,6 +807,10 @@ func (e *chatEncoder) event(ev Event) { e.col.add(ev) } +// keepalive is an SSE comment, as OpenAI-compatible servers keep a Chat +// Completions stream alive; its readers skip one. +func (e *chatEncoder) keepalive() { e.w.comment("keepalive") } + func (e *chatEncoder) finish() { if !e.started { e.start(Event{}) diff --git a/internal/gateway/claude_subscription.go b/internal/gateway/claude_subscription.go index e94088c5..7834144b 100644 --- a/internal/gateway/claude_subscription.go +++ b/internal/gateway/claude_subscription.go @@ -1478,7 +1478,8 @@ func relay(w http.ResponseWriter, r *http.Request, from provider.Protocol, name } return writeError(w, from, code, name+": "+msg), msg } - enc := encoder(from, newSSEWriter(w), req) + sw := newSSEWriter(w) + enc := encoder(from, sw, req) var failed, said, stop string see := func(ev Event) { switch ev.Kind { @@ -1497,9 +1498,12 @@ func relay(w http.ResponseWriter, r *http.Request, from provider.Protocol, name for _, ev := range head { see(ev) } - for ev := range events { + // a client waiting on a reply that goes on is kept from its idle + // timeout while none of it comes (#436) + relayEvents(events, sw, enc, func(ev Event) bool { see(ev) - } + return true + }) // before the reply's last event, which the agent may answer at once ended(said, stop, failed == "" && r.Context().Err() == nil) // an error event already ended the reply in the client's protocol diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 54adfe46..ee29a914 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -2288,9 +2288,10 @@ func (s *Server) translate(w http.ResponseWriter, r *http.Request, p provider.Pr return writeError(w, from, 502, msg), msg } if stream { - enc := encoder(from, newSSEWriter(w), request) + sw := newSSEWriter(w) + enc := encoder(from, sw, request) var failed string - serr := readSSE(rd, func(_, data string) error { + serr := readSSEAlive(rd, func(_, data string) error { return dec(data, func(ev Event) { switch ev.Kind { case KError: @@ -2301,6 +2302,12 @@ func (s *Server) translate(w http.ResponseWriter, r *http.Request, p provider.Pr } enc.event(ev) }) + }, func() { + // the provider's keepalives aren't events to translate: while + // it is heard from, the client hears from magpie (#436) + if failed == "" && sw.quiet() >= keepaliveGap { + enc.keepalive() + } }) if serr != nil && failed == "" { // the upstream died mid-reply: say so in the client's own @@ -2429,6 +2436,44 @@ func decoder(proto provider.Protocol) func(data string, emit func(Event)) error type streamEncoder interface { event(Event) finish() + // keepalive tells the client the reply goes on, as its protocol does + // with no answer to give: one its idle timeout counts (#436) + keepalive() +} + +// keepaliveGap is how long a translated reply's client may hear nothing +// while its provider is heard from (a keepalive, an event that has nothing +// for the client) before it is sent a keepalive of its own. +var keepaliveGap = time.Second + +// keepaliveEvery is how often a relayed reply that has gone quiet is kept +// alive, and keepaliveLongest how long it is kept so with no event: a +// reply stuck longer is left for the client's own idle timeout to end. +var keepaliveEvery, keepaliveLongest = 15 * time.Second, 5 * time.Minute + +// relayEvents hands see each of events until they end or see says stop, +// keeping the client of sw alive while none comes (keepaliveEvery, for up +// to keepaliveLongest since the last). +func relayEvents(events <-chan Event, sw *sseWriter, enc streamEncoder, see func(Event) bool) { + tick := time.NewTicker(keepaliveEvery) + defer tick.Stop() + last := time.Now() + for { + select { + case ev, ok := <-events: + if !ok { + return + } + last = time.Now() + if !see(ev) { + return + } + case <-tick.C: + if time.Since(last) < keepaliveLongest && sw.quiet() >= keepaliveEvery/2 { + enc.keepalive() + } + } + } } func encoder(proto provider.Protocol, w *sseWriter, r *Request) streamEncoder { diff --git a/internal/gateway/gemini.go b/internal/gateway/gemini.go index a8f4930b..0665f306 100644 --- a/internal/gateway/gemini.go +++ b/internal/gateway/gemini.go @@ -420,6 +420,11 @@ func (e *geminiEncoder) chunk(parts []map[string]any, finish string, usage map[s e.w.event("", msg) } +// keepalive sends nothing: Google's SDK takes a stream's "data: " lines +// one after another from the front and gets stuck on anything else, a +// comment too, and there is no event without content to send instead. +func (e *geminiEncoder) keepalive() {} + func (e *geminiEncoder) start(ev Event) { if e.started { return diff --git a/internal/gateway/keepalive_test.go b/internal/gateway/keepalive_test.go new file mode 100644 index 00000000..75c4a90f --- /dev/null +++ b/internal/gateway/keepalive_test.go @@ -0,0 +1,245 @@ +package gateway + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// keepFast makes the keepalives of a test's streams come at its pace. +func keepFast(t *testing.T, gap, every, longest time.Duration) { + t.Helper() + g, e, l := keepaliveGap, keepaliveEvery, keepaliveLongest + keepaliveGap, keepaliveEvery, keepaliveLongest = gap, every, longest + t.Cleanup(func() { keepaliveGap, keepaliveEvery, keepaliveLongest = g, e, l }) +} + +// TestTranslatedStreamKeepsClientAlive: a provider that says "progress", +// then only ": keepalive" every 20ms for half a second, then finishes, +// translated for a client whose idle timeout is 150ms (#436, congee949's +// repro). The client must hear something it counts all along: Codex's +// timeout counts events, never a comment, so a Responses or Anthropic +// client is kept alive with events, a Chat one with comments. +func TestTranslatedStreamKeepsClientAlive(t *testing.T) { + keepFast(t, 10*time.Millisecond, keepaliveEvery, keepaliveLongest) + for _, c := range []struct { + name, path, body, done, ping string + upChat bool + comments bool // the client counts a comment as life + }{ + {"chat_to_responses", "/v1/responses", `{"model":"fake/m1","stream":true,"input":"hi"}`, + `"type":"response.completed"`, `"type":"response.in_progress"`, true, false}, + {"chat_to_anthropic", "/v1/messages", `{"model":"fake/m1","stream":true,"max_tokens":64,"messages":[{"role":"user","content":"hi"}]}`, + `"type":"message_stop"`, `"type":"ping"`, true, false}, + {"responses_to_chat", "/v1/chat/completions", `{"model":"fake/m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`, + `"finish_reason":"stop"`, ": keepalive", false, true}, + } { + t.Run(c.name, func(t *testing.T) { + fresh(t) + var upPings atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.Copy(io.Discard, r.Body) + w.Header().Set("Content-Type", "text/event-stream") + f := w.(http.Flusher) + if c.upChat { + io.WriteString(w, "data: "+`{"id":"c1","model":"m1","choices":[{"index":0,"delta":{"content":"progress"},"finish_reason":null}]}`+"\n\n") + } else { + io.WriteString(w, "event: response.created\ndata: "+`{"type":"response.created","response":{"id":"r1","model":"m1"}}`+"\n\n") + io.WriteString(w, "event: response.output_text.delta\ndata: "+`{"type":"response.output_text.delta","delta":"progress"}`+"\n\n") + } + f.Flush() + tick := time.NewTicker(20 * time.Millisecond) + defer tick.Stop() + end := time.NewTimer(500 * time.Millisecond) + defer end.Stop() + for { + select { + case <-r.Context().Done(): + return + case <-tick.C: + if _, err := io.WriteString(w, ": keepalive\n\n"); err != nil { + return + } + upPings.Add(1) + f.Flush() + case <-end.C: + if c.upChat { + io.WriteString(w, "data: "+`{"id":"c1","model":"m1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n") + } else { + io.WriteString(w, "event: response.completed\ndata: "+`{"type":"response.completed","response":{"id":"r1","model":"m1","status":"completed","output":[]}}`+"\n\n") + } + f.Flush() + return + } + } + })) + defer up.Close() + p := provider.Provider{ID: "fake", Name: "Fake", Key: "fixture", Models: []string{"m1"}} + if c.upChat { + p.Chat = up.URL + "/v1" + } else { + p.Responses = up.URL + "/v1" + } + if err := provider.Save(p); err != nil { + t.Fatal(err) + } + server := httptest.NewServer(New().Handler()) + defer server.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, "POST", server.URL+c.path, strings.NewReader(c.body)) + req.Header.Set("Content-Type", "application/json") + res, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer res.Body.Close() + lines := make(chan string, 128) + go func() { + defer close(lines) + sc := bufio.NewScanner(res.Body) + for sc.Scan() { + lines <- sc.Text() + } + }() + idle := time.NewTimer(150 * time.Millisecond) + defer idle.Stop() + timedOut, completed, pings := false, false, 0 + read: + for { + select { + case line, ok := <-lines: + if !ok { + break read + } + if strings.Contains(line, c.ping) { + pings++ + } + if strings.Contains(line, c.done) { + completed = true + } + // what the client's idle timeout counts + if strings.HasPrefix(line, "data:") || (c.comments && strings.HasPrefix(line, ":")) { + idle.Reset(150 * time.Millisecond) + } + case <-idle.C: + timedOut = true + break read + } + } + cancel() + t.Logf("upstream_keepalives=%d downstream_keepalives=%d idle_timeout=%v completed=%v", upPings.Load(), pings, timedOut, completed) + if timedOut || !completed || pings < 3 { + t.Fatal("the client wasn't kept alive while the provider was") + } + }) + } +} + +// TestTranslatedSilenceStillTimesOut: a provider that says nothing at all +// isn't covered for; the client's idle timeout still ends the wait. +func TestTranslatedSilenceStillTimesOut(t *testing.T) { + keepFast(t, 10*time.Millisecond, keepaliveEvery, keepaliveLongest) + fresh(t) + release := make(chan struct{}) + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.Copy(io.Discard, r.Body) + w.Header().Set("Content-Type", "text/event-stream") + io.WriteString(w, "data: "+`{"id":"c1","model":"m1","choices":[{"index":0,"delta":{"content":"progress"},"finish_reason":null}]}`+"\n\n") + w.(http.Flusher).Flush() + select { + case <-r.Context().Done(): + case <-release: + } + })) + defer up.Close() + defer close(release) + if err := provider.Save(provider.Provider{ID: "fake", Name: "Fake", Key: "fixture", Models: []string{"m1"}, Chat: up.URL + "/v1"}); err != nil { + t.Fatal(err) + } + server := httptest.NewServer(New().Handler()) + defer server.Close() + ctx, cancel := context.WithTimeout(context.Background(), 400*time.Millisecond) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, "POST", server.URL+"/v1/responses", strings.NewReader(`{"model":"fake/m1","stream":true,"input":"hi"}`)) + res, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer res.Body.Close() + b, _ := io.ReadAll(res.Body) + if n := strings.Count(string(b), `"type":"response.in_progress"`); n > 1 { + t.Fatalf("a silent provider was covered for: %d in_progress events\n%s", n, b) + } +} + +// TestRelayKeepsClientAlive: a relayed reply (the Claude subscription's +// bridge, Kiro, Cursor, …) that goes quiet keeps its client alive every +// keepaliveEvery, but only for keepaliveLongest since its last event. +func TestRelayKeepsClientAlive(t *testing.T) { + keepFast(t, keepaliveGap, 20*time.Millisecond, 200*time.Millisecond) + events := make(chan Event) + go func() { + events <- Event{Kind: KStart, MsgID: "m1"} + events <- Event{Kind: KText, Text: "progress"} + time.Sleep(150 * time.Millisecond) // kept alive + events <- Event{Kind: KText, Text: " more"} + time.Sleep(600 * time.Millisecond) // kept alive for 200ms only + events <- Event{Kind: KStop, Stop: "end_turn"} + close(events) + }() + rec := httptest.NewRecorder() + req := &Request{Model: "m1", Stream: true} + r := httptest.NewRequest("POST", "/v1/messages", nil) + code, failed := relay(rec, r, provider.Anthropic, "Bridge", req, events, &Usage{}, func() {}, func(string, string, bool) {}) + if code != 200 || failed != "" { + t.Fatalf("relay: %d %q", code, failed) + } + body := rec.Body.String() + parts := strings.SplitN(body, "more", 2) + if len(parts) != 2 || !strings.Contains(body, "message_stop") { + t.Fatalf("reply incomplete:\n%s", body) + } + before, after := strings.Count(parts[0], `"type":"ping"`), strings.Count(parts[1], `"type":"ping"`) + t.Logf("pings in a 150ms gap: %d, in a 600ms gap (200ms kept): %d", before, after) + if before < 3 { + t.Fatalf("a quiet reply wasn't kept alive: %d pings in 150ms", before) + } + if after < 3 || after > 12 { + t.Fatalf("a reply quiet past keepaliveLongest was kept alive on: %d pings", after) + } +} + +// TestTraceKeepsInflight: a request still going stays in the trace while +// more than traceKeep come and finish after it (#436). +func TestTraceKeepsInflight(t *testing.T) { + s := New() + pending := s.trace.begin(Route{Model: "pending", Time: time.Now()}) + for i := 0; i < traceKeep+1; i++ { + r := s.trace.begin(Route{Model: "later", Time: time.Now()}) + s.trace.update(r, func(r *Route) { r.Done, r.Status = true, 200 }) + } + state := s.Trace(context.Background(), 0, 0) + found := false + for _, r := range state.Routes { + found = found || r.ID == pending.ID + } + if !found || len(state.Routes) != traceKeep { + t.Fatalf("pending visible %v, %d routes kept", found, len(state.Routes)) + } + // more going than twice traceKeep: the oldest go after all + for i := 0; i < 2*traceKeep; i++ { + s.trace.begin(Route{Model: "going", Time: time.Now()}) + } + if n := len(s.Trace(context.Background(), 0, 0).Routes); n != 2*traceKeep { + t.Fatalf("%d routes kept, want %d", n, 2*traceKeep) + } +} diff --git a/internal/gateway/qoder.go b/internal/gateway/qoder.go index 1e9f3f85..9ef2b7cd 100644 --- a/internal/gateway/qoder.go +++ b/internal/gateway/qoder.go @@ -120,22 +120,28 @@ func relayStatus(w http.ResponseWriter, from provider.Protocol, name string, req if len(head) == 0 || head[len(head)-1].Kind == KStart || head[len(head)-1].Kind == KUsage { return fail(Event{Text: name + " ended without an answer"}) } - enc := encoder(from, newSSEWriter(w), req) + sw := newSSEWriter(w) + enc := encoder(from, sw, req) + var failed string see := func(ev Event) bool { if ev.Kind == KStart || ev.Kind == KUsage { usage.add(ev.Usage) } enc.event(ev) - return ev.Kind != KError + if ev.Kind == KError { + failed = ev.Text + return false + } + return true } for _, ev := range head { see(ev) } - for ev := range events { - if !see(ev) { - abort() - return 200, ev.Text - } + // kept from the client's idle timeout while none comes (#436) + relayEvents(events, sw, enc, see) + if failed != "" { + abort() + return 200, failed } enc.finish() return 200, "" diff --git a/internal/gateway/responses.go b/internal/gateway/responses.go index 687ec7e0..ee3a20a6 100644 --- a/internal/gateway/responses.go +++ b/internal/gateway/responses.go @@ -827,6 +827,16 @@ func (e *responsesEncoder) event(ev Event) { e.col.add(ev) } +// keepalive is response.in_progress again: Codex's idle timeout counts +// events only, an SSE comment never reaching it, and skips this one. +func (e *responsesEncoder) keepalive() { + if !e.started { + e.start(Event{}) // response.created and response.in_progress + return + } + e.send("response.in_progress", map[string]any{"response": e.response("in_progress", nil)}) +} + func (e *responsesEncoder) finish() { if !e.started { e.start(Event{}) diff --git a/internal/gateway/sse.go b/internal/gateway/sse.go index 892e5bd9..442f1062 100644 --- a/internal/gateway/sse.go +++ b/internal/gateway/sse.go @@ -7,11 +7,19 @@ import ( "io" "net/http" "strings" + "time" ) // readSSE walks a server-sent event stream, calling fn with each event's // name and data. Data lines of one event are joined with newlines. func readSSE(r io.Reader, fn func(event, data string) error) error { + return readSSEAlive(r, fn, nil) +} + +// readSSEAlive is readSSE that also calls alive, when not nil, after each +// event and each comment: every sign that the stream is alive, a +// keepalive (": keepalive") as much as an answer. +func readSSEAlive(r io.Reader, fn func(event, data string) error, alive func()) error { sc := bufio.NewScanner(r) sc.Buffer(make([]byte, 0, 64<<10), 32<<20) var event string @@ -23,6 +31,9 @@ func readSSE(r io.Reader, fn func(event, data string) error) error { } err := fn(event, strings.Join(data, "\n")) event, data = "", nil + if err == nil && alive != nil { + alive() + } return err } for sc.Scan() { @@ -33,6 +44,9 @@ func readSSE(r io.Reader, fn func(event, data string) error) error { return err } case strings.HasPrefix(line, ":"): + if alive != nil { + alive() + } case strings.HasPrefix(line, "event:"): event = strings.TrimSpace(line[6:]) case strings.HasPrefix(line, "data:"): @@ -50,6 +64,7 @@ type sseWriter struct { w http.ResponseWriter f http.Flusher begun bool + wrote time.Time // when the client was last written to } func newSSEWriter(w http.ResponseWriter) *sseWriter { @@ -69,6 +84,7 @@ func (s *sseWriter) begin() { h.Set("X-Accel-Buffering", "no") s.w.WriteHeader(http.StatusOK) s.flush() + s.wrote = time.Now() } func (s *sseWriter) flush() { @@ -77,6 +93,24 @@ func (s *sseWriter) flush() { } } +// quiet is how long the client has had nothing from s; a stream not yet +// begun has been quiet for ever. +func (s *sseWriter) quiet() time.Duration { + if !s.begun { + return 1<<63 - 1 + } + return time.Since(s.wrote) +} + +// comment writes an SSE comment, which every event stream reader skips: +// no event, only the news that the stream is alive. +func (s *sseWriter) comment(text string) { + s.begin() + io.WriteString(s.w, ": "+text+"\n\n") + s.flush() + s.wrote = time.Now() +} + // event writes one event; a JSON-marshallable value or a raw string. func (s *sseWriter) event(name string, v any) { s.begin() @@ -101,4 +135,5 @@ func (s *sseWriter) event(name string, v any) { b.WriteString("\n\n") s.w.Write(b.Bytes()) s.flush() + s.wrote = time.Now() } diff --git a/internal/gateway/trace.go b/internal/gateway/trace.go index 413ceb4c..42ea6c1a 100644 --- a/internal/gateway/trace.go +++ b/internal/gateway/trace.go @@ -10,6 +10,7 @@ package gateway import ( "context" + "slices" "strings" "sync" "time" @@ -299,14 +300,29 @@ func (t *trace) begin(r Route) *Route { } rp := &r t.routes = append(t.routes, rp) - if len(t.routes) > traceKeep { - t.routes = t.routes[len(t.routes)-traceKeep:] - } + t.trim() t.changed() rp.Seq = t.seq return rp } +// trim keeps traceKeep routes, the oldest finished ones going first: a +// request still going isn't dropped for those that came after it (#436), +// as the history has only finished ones, unless more than twice traceKeep +// are going at once. +func (t *trace) trim() { + for len(t.routes) > traceKeep { + i := slices.IndexFunc(t.routes, func(r *Route) bool { return r.Done }) + if i < 0 { + if len(t.routes) <= 2*traceKeep { + return + } + i = 0 + } + t.routes = slices.Delete(t.routes, i, i+1) + } +} + // update changes a route under the lock. func (t *trace) update(r *Route, f func(r *Route)) { if r == nil {