From 7a44cdbf821d21a6a2829356ee082699b36949ec Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 19 Aug 2026 02:05:02 +0000 Subject: [PATCH 1/2] test(relay): prove playground stop does not cancel Ollama stream Add a failing regression for #6860: client abort mid-Ollama NDJSON stream must close the upstream body so Ollama stops generating. Co-authored-by: Yuzhong Zhang --- relay/channel/ollama/stream_test.go | 90 +++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/relay/channel/ollama/stream_test.go b/relay/channel/ollama/stream_test.go index 69aff1ce8f..01a30e6af6 100644 --- a/relay/channel/ollama/stream_test.go +++ b/relay/channel/ollama/stream_test.go @@ -1,11 +1,15 @@ package ollama import ( + "context" + "fmt" "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -97,3 +101,89 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) { }) } } + +// TestOllamaStreamHandlerClientCancelClosesUpstream pins the playground-stop +// contract for Ollama NDJSON streams: aborting the client request (Playground +// stop) must close the upstream body so Ollama stops generating, and the +// handler must return without waiting for more tokens. +func TestOllamaStreamHandlerClientCancelClosesUpstream(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pr, pw := io.Pipe() + t.Cleanup(func() { + _ = pr.Close() + _ = pw.Close() + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(ctx) + firstHandled := make(chan struct{}) + c.Writer = ¬ifyAfterWriter{ResponseWriter: c.Writer, needle: "halfway", notify: firstHandled} + + resp := &http.Response{Body: pr} + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "qwen3"}, + } + + done := make(chan struct{}) + go func() { + _, _ = ollamaStreamHandler(c, info, resp) + close(done) + }() + + first := `{"model":"qwen3","created_at":"2026-08-14T12:00:00Z","message":{"role":"assistant","content":"","thinking":"halfway"},"done":false}` + "\n" + _, err := fmt.Fprint(pw, first) + require.NoError(t, err) + + select { + case <-firstHandled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first thinking chunk") + } + + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("handler did not return after client disconnect") + } + + _, err = fmt.Fprint(pw, `{"model":"qwen3","message":{"role":"assistant","content":"more"},"done":false}`+"\n") + require.ErrorIs(t, err, io.ErrClosedPipe, "upstream body should be closed after client disconnect") + + body := recorder.Body.String() + assert.Contains(t, body, "halfway") + assert.NotContains(t, body, `"more"`) +} + +// notifyAfterWriter signals once the streamed payload containing needle has +// been written, so the test can cancel after the first thinking chunk. +type notifyAfterWriter struct { + gin.ResponseWriter + needle string + notify chan struct{} + once sync.Once +} + +func (w *notifyAfterWriter) Write(p []byte) (int, error) { + n, err := w.ResponseWriter.Write(p) + if strings.Contains(string(p), w.needle) { + w.once.Do(func() { close(w.notify) }) + } + return n, err +} + +func (w *notifyAfterWriter) WriteString(s string) (int, error) { + n, err := io.WriteString(w.ResponseWriter, s) + if strings.Contains(s, w.needle) { + w.once.Do(func() { close(w.notify) }) + } + return n, err +} From a2c12a3b3f7d245e77d395ac563b0eafc96421db Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 19 Aug 2026 02:07:04 +0000 Subject: [PATCH 2/2] fix(relay): close Ollama upstream when the client aborts Watch the request context in ollamaStreamHandler and close the upstream body on cancel so Playground stop actually stops Ollama. Co-authored-by: Yuzhong Zhang --- relay/channel/ollama/stream.go | 16 +++++++++++++++- relay/channel/ollama/stream_test.go | 28 ++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index a0d7839f9f..06c3cdd203 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -101,6 +101,17 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } defer service.CloseResponseBodyGracefully(resp) + clientCtx := c.Request.Context() + stopWatch := make(chan struct{}) + defer close(stopWatch) + go func() { + select { + case <-clientCtx.Done(): + service.CloseResponseBodyGracefully(resp) + case <-stopWatch: + } + }() + helper.SetEventStreamHeaders(c) scanner := helper.NewStreamScanner(resp.Body) usage := &dto.Usage{} @@ -114,6 +125,9 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } for scanner.Scan() { + if clientCtx.Err() != nil { + return usage, nil + } line := scanner.Text() line = strings.TrimSpace(line) if line == "" { @@ -200,7 +214,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http helper.Done(c) break } - if err := scanner.Err(); err != nil && err != io.EOF { + if err := scanner.Err(); err != nil && err != io.EOF && clientCtx.Err() == nil { logger.LogError(c, "ollama stream scan error: "+err.Error()) } return usage, nil diff --git a/relay/channel/ollama/stream_test.go b/relay/channel/ollama/stream_test.go index 01a30e6af6..d2aee7f392 100644 --- a/relay/channel/ollama/stream_test.go +++ b/relay/channel/ollama/stream_test.go @@ -102,6 +102,34 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) { } } +func TestOllamaStreamHandlerCompletesThinkingAndContent(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + raw := strings.Join([]string{ + `{"model":"qwen3","created_at":"2026-08-14T12:00:00Z","message":{"role":"assistant","content":"","thinking":"plan"},"done":false}`, + `{"model":"qwen3","created_at":"2026-08-14T12:00:01Z","message":{"role":"assistant","content":"hello"},"done":false}`, + `{"model":"qwen3","created_at":"2026-08-14T12:00:02Z","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":3,"eval_count":2}`, + }, "\n") + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + usage, apiErr := ollamaStreamHandler(c, &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "qwen3"}, + }, &http.Response{Body: io.NopCloser(strings.NewReader(raw))}) + require.Nil(t, apiErr) + require.NotNil(t, usage) + assert.Equal(t, 5, usage.TotalTokens) + + body := recorder.Body.String() + assert.Contains(t, body, "plan") + assert.Contains(t, body, "hello") + assert.Contains(t, body, "[DONE]") +} + // TestOllamaStreamHandlerClientCancelClosesUpstream pins the playground-stop // contract for Ollama NDJSON streams: aborting the client request (Playground // stop) must close the upstream body so Ollama stops generating, and the