1

internal/stream/engine.go

@@ -71,15 +71,30 @@ func ConsumeSSE(cfg ConsumeConfig, hooks ConsumeHooks) {
 			hooks.OnFinalize(reason, scannerErr)
 		}
 	}
+	contextDone := func() bool {
+		if cfg.Context.Err() == nil {
+			return false
+		}
+		if hooks.OnContextDone != nil {
+			hooks.OnContextDone()
+		}
+		return true
+	}
 
 	for {
+		if contextDone() {
+			return
+		}
 		select {
 		case <-cfg.Context.Done():
-			if hooks.OnContextDone != nil {
-				hooks.OnContextDone()
+			if contextDone() {
+				return
 			}
 			return
 		case <-tickCh(ticker):
+			if contextDone() {
+				return
+			}
 			if !hasContent {
 				keepaliveCount++
 				if cfg.MaxKeepAliveNoInput > 0 && keepaliveCount >= cfg.MaxKeepAliveNoInput {
@@ -95,6 +110,9 @@ func ConsumeSSE(cfg ConsumeConfig, hooks ConsumeHooks) {
 				hooks.OnKeepAlive()
 			}
 		case parsed, ok := <-parsedLines:
+			if contextDone() {
+				return
+			}
 			if !ok {
 				finalize(StopReasonUpstreamCompleted, <-done)
 				return

internal/stream/engine_test.go

@@ -0,0 +1,47 @@
+package stream
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"ds2api/internal/sse"
+)
+
+func TestConsumeSSEPrefersContextCancellationOverReadyParsedLines(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	var finalized bool
+	var contextDone bool
+	var parsedCalled bool
+
+	ConsumeSSE(ConsumeConfig{
+		Context:           ctx,
+		Body:              strings.NewReader("data: {\"p\":\"response/content\",\"v\":\"hello\"}\n\ndata: [DONE]\n"),
+		ThinkingEnabled:   false,
+		InitialType:       "text",
+		KeepAliveInterval: 0,
+	}, ConsumeHooks{
+		OnParsed: func(_ sse.LineResult) ParsedDecision {
+			parsedCalled = true
+			return ParsedDecision{}
+		},
+		OnFinalize: func(_ StopReason, _ error) {
+			finalized = true
+		},
+		OnContextDone: func() {
+			contextDone = true
+		},
+	})
+
+	if !contextDone {
+		t.Fatal("expected OnContextDone to run for an already-cancelled context")
+	}
+	if finalized {
+		t.Fatal("expected OnFinalize not to run after context cancellation wins")
+	}
+	if parsedCalled {
+		t.Fatal("expected parsed lines not to be processed after context cancellation wins")
+	}
+}