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