Merge a2c12a3b3f into 66031a09d9

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

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,117 @@ 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
+// 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 = &notifyAfterWriter{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
+}