diff --git a/provider_openai.go b/provider_openai.go index 9103be1..106126f 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "net/http" + "strings" ) // 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 { return usage{}, err } @@ -111,9 +121,13 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite 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) } + // 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) } @@ -151,7 +165,7 @@ type chatUsagePayload struct { func (a *openaiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response) (usage, error) { body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) 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 @@ -176,7 +190,11 @@ const sseDoneMarker = "[DONE]" // the returned usage either way, but only forwarded to the client when // clientAskedUsage is true — otherwise it is dropped, since the client // 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) { sw := newSSEWriter(w) var u usage @@ -201,7 +219,10 @@ func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, cli return sw.writeData(ev.data) }) 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 } @@ -228,12 +249,12 @@ func (a *openaiAdapter) listModels(ctx context.Context) ([]string, error) { body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) 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 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)) diff --git a/provider_openai_test.go b/provider_openai_test.go index d371cf6..a48ee22 100644 --- a/provider_openai_test.go +++ b/provider_openai_test.go @@ -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 // response returns a *providerHTTPError carrying the status and body, // with nothing written to the client — the caller decides envelope vs diff --git a/providers.go b/providers.go index dbd354c..207528f 100644 --- a/providers.go +++ b/providers.go @@ -115,7 +115,17 @@ type providerAdapter interface { // pool, since every request from one adapter targets the same upstream // host. 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 return &http.Client{Transport: tr} } @@ -135,14 +145,14 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string, if body != nil { b, err := json.Marshal(body) 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) } req, err := http.NewRequestWithContext(ctx, method, url, r) 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 _, 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 if err != nil { - return nil, fmt.Errorf("%w: %v", errUpstream, err) //nolint:errorlint + return nil, fmt.Errorf("%w: %w", errUpstream, err) } return resp, nil }