mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-20 05:56:56 +00:00
fix(providers): preserve captured usage on stream errors, %w wrapping, defensive transport clone, JSON fallback for ignored stream flag
This commit is contained in:
+28
-7
@@ -6,6 +6,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
// openaiAdapter is the providerAdapter for provider type "openai": OpenAI
|
// openaiAdapter is the providerAdapter for provider type "openai": OpenAI
|
||||||
@@ -101,7 +102,16 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/chat/completions", a.requestHeaders(true), req)
|
hdr := a.requestHeaders(true)
|
||||||
|
if streaming {
|
||||||
|
// Ask the upstream for an SSE response. The Content-Type check below
|
||||||
|
// still decides which forwarding path runs — an upstream that
|
||||||
|
// ignores stream=true may answer with JSON regardless of this
|
||||||
|
// header.
|
||||||
|
hdr.Set("Accept", "text/event-stream")
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/chat/completions", hdr, req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return usage{}, err
|
return usage{}, err
|
||||||
}
|
}
|
||||||
@@ -111,9 +121,13 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
|
|||||||
return usage{}, newProviderHTTPError(resp)
|
return usage{}, newProviderHTTPError(resp)
|
||||||
}
|
}
|
||||||
|
|
||||||
if streaming {
|
if streaming && strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") {
|
||||||
return a.forwardStream(w, resp.Body, clientAskedUsage)
|
return a.forwardStream(w, resp.Body, clientAskedUsage)
|
||||||
}
|
}
|
||||||
|
// Either a non-streaming request, or a streaming one whose upstream
|
||||||
|
// ignored stream=true and answered with a normal JSON body: forward it
|
||||||
|
// through the same verbatim passthrough non-streaming responses use,
|
||||||
|
// instead of emitting an empty SSE stream for a body that was never SSE.
|
||||||
return a.forwardJSON(w, resp)
|
return a.forwardJSON(w, resp)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -151,7 +165,7 @@ type chatUsagePayload struct {
|
|||||||
func (a *openaiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response) (usage, error) {
|
func (a *openaiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response) (usage, error) {
|
||||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
|
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return usage{}, fmt.Errorf("%w: read response body: %v", errUpstream, err) //nolint:errorlint
|
return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
var payload chatUsagePayload
|
var payload chatUsagePayload
|
||||||
@@ -176,7 +190,11 @@ const sseDoneMarker = "[DONE]"
|
|||||||
// the returned usage either way, but only forwarded to the client when
|
// the returned usage either way, but only forwarded to the client when
|
||||||
// clientAskedUsage is true — otherwise it is dropped, since the client
|
// clientAskedUsage is true — otherwise it is dropped, since the client
|
||||||
// never asked for it and chatCompletion only requested it upstream for
|
// never asked for it and chatCompletion only requested it upstream for
|
||||||
// accounting. The "[DONE]" sentinel is always forwarded.
|
// accounting. The "[DONE]" sentinel is always forwarded. On a mid-stream
|
||||||
|
// read error (upstream drops the connection, client context is canceled),
|
||||||
|
// forwardStream still returns whatever usage it had already captured
|
||||||
|
// alongside the error — a usage chunk that arrived before the break must
|
||||||
|
// not be discarded and billed as zero.
|
||||||
func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, clientAskedUsage bool) (usage, error) {
|
func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, clientAskedUsage bool) (usage, error) {
|
||||||
sw := newSSEWriter(w)
|
sw := newSSEWriter(w)
|
||||||
var u usage
|
var u usage
|
||||||
@@ -201,7 +219,10 @@ func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, cli
|
|||||||
return sw.writeData(ev.data)
|
return sw.writeData(ev.data)
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return usage{}, fmt.Errorf("%w: stream: %v", errUpstream, err) //nolint:errorlint
|
// u carries whatever usage was already captured before the stream
|
||||||
|
// broke — a usage chunk that arrived, then a dropped connection,
|
||||||
|
// must not report as zero and get treated as an unbilled request.
|
||||||
|
return u, fmt.Errorf("%w: stream: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
return u, nil
|
return u, nil
|
||||||
}
|
}
|
||||||
@@ -228,12 +249,12 @@ func (a *openaiAdapter) listModels(ctx context.Context) ([]string, error) {
|
|||||||
|
|
||||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
|
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("%w: read response body: %v", errUpstream, err) //nolint:errorlint
|
return nil, fmt.Errorf("%w: read response body: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
var payload modelsPayload
|
var payload modelsPayload
|
||||||
if err := json.Unmarshal(body, &payload); err != nil {
|
if err := json.Unmarshal(body, &payload); err != nil {
|
||||||
return nil, fmt.Errorf("%w: decode models response: %v", errUpstream, err) //nolint:errorlint
|
return nil, fmt.Errorf("%w: decode models response: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
ids := make([]string, len(payload.Data))
|
ids := make([]string, len(payload.Data))
|
||||||
|
|||||||
@@ -247,6 +247,123 @@ func TestOpenAIAdapter_ChatCompletion_Streaming_ClientAskedUsage(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestOpenAIAdapter_ChatCompletion_Streaming_UsageCapturedBeforeAbruptClose
|
||||||
|
// proves that when the upstream sends the usage-only chunk and then the
|
||||||
|
// connection drops (no [DONE], no clean EOF), chatCompletion still returns
|
||||||
|
// the usage it had already captured alongside the error — a dropped
|
||||||
|
// connection after usage arrived must not read as zero usage.
|
||||||
|
func TestOpenAIAdapter_ChatCompletion_Streaming_UsageCapturedBeforeAbruptClose(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
fmt.Fprintf(w, "data: %s\n\n", streamFrames[3]) // usage-only chunk: {prompt:7 completion:9}
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
|
||||||
|
hj, ok := w.(http.Hijacker)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("httptest ResponseWriter does not implement http.Hijacker")
|
||||||
|
}
|
||||||
|
conn, _, err := hj.Hijack()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("hijack: %v", err)
|
||||||
|
}
|
||||||
|
// Abrupt close: no [DONE], no terminating chunk — the client must
|
||||||
|
// see a genuine read error, not a clean EOF.
|
||||||
|
_ = conn.Close()
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
a := newOpenAIAdapter("p1", srv.URL, "sk-test")
|
||||||
|
rec := newRecordingResponseWriter()
|
||||||
|
|
||||||
|
req := map[string]any{"model": "gpt-5", "stream": true}
|
||||||
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("chatCompletion: want error from abrupt close, got nil (usage=%+v)", u)
|
||||||
|
}
|
||||||
|
if !errors.Is(err, errUpstream) {
|
||||||
|
t.Errorf("err = %v, want errors.Is(err, errUpstream) == true", err)
|
||||||
|
}
|
||||||
|
if u.prompt != 7 || u.completion != 9 {
|
||||||
|
t.Errorf("usage = %+v, want {prompt:7 completion:9} (captured before the drop)", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestOpenAIAdapter_ChatCompletion_ContextCanceledMidStream proves that
|
||||||
|
// canceling the caller's context while a streaming response is still being
|
||||||
|
// read surfaces as errors.Is(err, context.Canceled) at chatCompletion's
|
||||||
|
// caller — required so Task 12 can tell a client hang-up apart from a real
|
||||||
|
// upstream failure.
|
||||||
|
func TestOpenAIAdapter_ChatCompletion_ContextCanceledMidStream(t *testing.T) {
|
||||||
|
block := make(chan struct{})
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
fmt.Fprintf(w, "data: %s\n\n", streamFrames[0])
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
<-block // hold the connection open past the context cancellation below
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
defer close(block) // unblock the handler before srv.Close() waits on it
|
||||||
|
|
||||||
|
a := newOpenAIAdapter("p1", srv.URL, "sk-test")
|
||||||
|
rec := newRecordingResponseWriter()
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
go func() {
|
||||||
|
time.Sleep(50 * time.Millisecond)
|
||||||
|
cancel()
|
||||||
|
}()
|
||||||
|
|
||||||
|
_, err := a.chatCompletion(ctx, rec, map[string]any{"model": "gpt-5", "stream": true})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("chatCompletion: want error from context cancellation, got nil")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("err = %v, want errors.Is(err, context.Canceled) == true", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestOpenAIAdapter_ChatCompletion_Streaming_UpstreamIgnoresStreamFallsBackToJSON
|
||||||
|
// proves that when a streaming request's upstream answers with a plain
|
||||||
|
// JSON body instead of SSE (Content-Type not text/event-stream), and sends
|
||||||
|
// "Accept: text/event-stream" on the outgoing request, chatCompletion falls
|
||||||
|
// back to the non-streaming passthrough — the client gets the JSON body
|
||||||
|
// verbatim and usage is still extracted — instead of emitting an empty or
|
||||||
|
// broken SSE stream.
|
||||||
|
func TestOpenAIAdapter_ChatCompletion_Streaming_UpstreamIgnoresStreamFallsBackToJSON(t *testing.T) {
|
||||||
|
const respBody = `{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":3,"completion_tokens":4}}`
|
||||||
|
|
||||||
|
var gotAccept string
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotAccept = r.Header.Get("Accept")
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
_, _ = w.Write([]byte(respBody))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
a := newOpenAIAdapter("p1", srv.URL, "sk-test")
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
|
||||||
|
u, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5", "stream": true})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("chatCompletion: %v", err)
|
||||||
|
}
|
||||||
|
if u.prompt != 3 || u.completion != 4 {
|
||||||
|
t.Errorf("usage = %+v, want {prompt:3 completion:4}", u)
|
||||||
|
}
|
||||||
|
if rec.Body.String() != respBody {
|
||||||
|
t.Errorf("body = %q, want %q (JSON fallback, verbatim passthrough)", rec.Body.String(), respBody)
|
||||||
|
}
|
||||||
|
if got := rec.Header().Get("Content-Type"); got != "application/json" {
|
||||||
|
t.Errorf("Content-Type = %q, want application/json", got)
|
||||||
|
}
|
||||||
|
if gotAccept != "text/event-stream" {
|
||||||
|
t.Errorf("Accept header sent = %q, want %q", gotAccept, "text/event-stream")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// TestOpenAIAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream
|
// TestOpenAIAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream
|
||||||
// response returns a *providerHTTPError carrying the status and body,
|
// response returns a *providerHTTPError carrying the status and body,
|
||||||
// with nothing written to the client — the caller decides envelope vs
|
// with nothing written to the client — the caller decides envelope vs
|
||||||
|
|||||||
+14
-4
@@ -115,7 +115,17 @@ type providerAdapter interface {
|
|||||||
// pool, since every request from one adapter targets the same upstream
|
// pool, since every request from one adapter targets the same upstream
|
||||||
// host.
|
// host.
|
||||||
func newAdapterHTTPClient() *http.Client {
|
func newAdapterHTTPClient() *http.Client {
|
||||||
tr := http.DefaultTransport.(*http.Transport).Clone() //nolint:forcetypeassert // http.DefaultTransport is always *http.Transport
|
tr, ok := http.DefaultTransport.(*http.Transport)
|
||||||
|
if ok {
|
||||||
|
tr = tr.Clone()
|
||||||
|
} else {
|
||||||
|
// Defensive fallback: something in the shared Traefik process
|
||||||
|
// replaced http.DefaultTransport with a type that is not
|
||||||
|
// *http.Transport. A fresh zero-value Transport (Go's own defaults)
|
||||||
|
// is safer than panicking at adapter construction over a type
|
||||||
|
// assertion this package does not control.
|
||||||
|
tr = &http.Transport{}
|
||||||
|
}
|
||||||
tr.MaxIdleConnsPerHost = adapterIdleConnsPerHost
|
tr.MaxIdleConnsPerHost = adapterIdleConnsPerHost
|
||||||
return &http.Client{Transport: tr}
|
return &http.Client{Transport: tr}
|
||||||
}
|
}
|
||||||
@@ -135,14 +145,14 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string,
|
|||||||
if body != nil {
|
if body != nil {
|
||||||
b, err := json.Marshal(body)
|
b, err := json.Marshal(body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("%w: encode request body: %v", errUpstream, err) //nolint:errorlint // errUpstream is a sentinel wrapped for errors.Is, %v is fine for the inner cause
|
return nil, fmt.Errorf("%w: encode request body: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
r = bytes.NewReader(b)
|
r = bytes.NewReader(b)
|
||||||
}
|
}
|
||||||
|
|
||||||
req, err := http.NewRequestWithContext(ctx, method, url, r)
|
req, err := http.NewRequestWithContext(ctx, method, url, r)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("%w: build request: %v", errUpstream, err) //nolint:errorlint
|
return nil, fmt.Errorf("%w: build request: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
for k, vs := range hdr {
|
for k, vs := range hdr {
|
||||||
for _, v := range vs {
|
for _, v := range vs {
|
||||||
@@ -155,7 +165,7 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string,
|
|||||||
|
|
||||||
resp, err := client.Do(req) //nolint:bodyclose // caller closes resp.Body; upstreamJSON hands the response, not its lifecycle, back
|
resp, err := client.Do(req) //nolint:bodyclose // caller closes resp.Body; upstreamJSON hands the response, not its lifecycle, back
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("%w: %v", errUpstream, err) //nolint:errorlint
|
return nil, fmt.Errorf("%w: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
return resp, nil
|
return resp, nil
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user