From 5a4bbefe722ff5f7bd09eb31ed2ae10f459625dd Mon Sep 17 00:00:00 2001
From: Zexi Li <zexi.li@icloud.com>
Date: Tue, 23 Jun 2026 14:09:37 +0800
Subject: [PATCH] fix(aiproxy): defer stream context cancel until chunks are
 consumed (#25058)

---
 pkg/aiproxy/handlers/chat_completions.go | 19 ++++++++++++++++++-
 pkg/aiproxy/handlers/completions.go      |  7 ++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pkg/aiproxy/handlers/chat_completions.go b/pkg/aiproxy/handlers/chat_completions.go
index 80725c0778..3ab365ab4c 100644
--- a/pkg/aiproxy/handlers/chat_completions.go
+++ b/pkg/aiproxy/handlers/chat_completions.go
@@ -82,6 +82,19 @@ func flushIf(w http.ResponseWriter) {
 	}
 }
 
+// streamChunksWithCancel forwards chunks until ch closes, then releases reqCtx.
+func streamChunksWithCancel(ch <-chan upstream.StreamChunk, cancel context.CancelFunc) <-chan upstream.StreamChunk {
+	out := make(chan upstream.StreamChunk, 16)
+	go func() {
+		defer cancel()
+		defer close(out)
+		for chunk := range ch {
+			out <- chunk
+		}
+	}()
+	return out
+}
+
 // chatCompletionsHandler implements OpenAI-compatible POST /openai/v1/chat/completions.
 // Auth is the ai_virtual_key only (Authorization: Bearer <vk> or X-Ai-Virtual-Key).
 // Upstream is resolved: ai_virtual_key -> project ai_routing -> ai_routing_model -> ai_key (by catalog model_key).
@@ -287,7 +300,11 @@ func chatCompletionStreamWithKeyFailover(
 		}
 		reqCtx, cancel := context.WithTimeout(ctx, timeout)
 		ch, uerr := providerStreamChunks(reqCtx, up, upReq, prov)
-		cancel()
+		if uerr != nil {
+			cancel()
+		} else {
+			ch = streamChunksWithCancel(ch, cancel)
+		}
 		if uerr == nil {
 			return ch, nil
 		}
diff --git a/pkg/aiproxy/handlers/completions.go b/pkg/aiproxy/handlers/completions.go
index a1087b20f2..5806375788 100644
--- a/pkg/aiproxy/handlers/completions.go
+++ b/pkg/aiproxy/handlers/completions.go
@@ -218,9 +218,14 @@ func completionsStreamWithKeyFailover(
 		if compProv.OpenAICompletionsStreamPassthrough() {
 			ch, uerr = upstream.ChatCompletionStream(reqCtx, upReq)
 		} else {
+			cancel()
 			return nil, &upstream.Error{StatusCode: http.StatusBadRequest, Message: "streaming completions not supported for provider"}
 		}
-		cancel()
+		if uerr != nil {
+			cancel()
+		} else {
+			ch = streamChunksWithCancel(ch, cancel)
+		}
 		if uerr == nil {
 			return ch, nil
 		}