fix(aiproxy): defer stream context cancel until chunks are consumed (#25059)

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
 		}

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
 		}