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 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 }