From 3ed473bdfe54e578a889fa28ed09789708ac150f Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 10:55:39 +0100 Subject: [PATCH] fix(routes): bill partial stream usage on error, pin openai alias stripping Review fixes for Task 12: 1. (CRITICAL) runUnified now calls g.limiter.account BEFORE handling an adapter error, not after returning on it. All three adapters' forwardStream deliberately return captured usage alongside a mid-stream error, so a client that aborts right after the usage frame arrives was getting tokens delivered but never billed - repeated aborts bypassed every token/cost budget. A non-streaming error path always carries zero usage by contract, so accounting it is a no-op (account skips writes once both total tokens and cost are zero). The zero-usage-stream logf and the zero-usage-non-stream estimate now only fire when callErr is nil, so a genuine error no longer gets a fabricated token estimate. New test: TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage (hijack+close after the usage chunk; asserts the limiter's counter shows the captured tokens). Verified this test fails against the pre-fix ordering and the two new alias tests below fail when their guarded delete() call is removed, before restoring and re-running clean. 2. Added TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody and TestHandleEmbeddings_OpenAI_AliasKeyStrippedFromUpstreamBody: the openai-type adapter marshals req verbatim as the upstream wire body, so it is the one path where a leftover gatewayAliasKey entry actually reaches a real provider. The only existing alias-leak assertion covered anthropic, which rebuilds the body field-by-field and was never at risk. 3. Reworded the stale comment on limiter.evaluateScope (limits.go): it described a nil-limits scope's req counters as a normal, expected case. Ruling (e), implemented in routes_unified.go's buildLimitScopes, means the unified route never builds a nil-limits scope in the first place - the nil check remains a defensive no-op, not a path any current caller exercises. Gates: gofmt, go vet, fieldalignment, golangci-lint (incl. gosec) all clean; go test ./... -race -count=1 passes (full suite, including all Task 12 tests). --- limits.go | 11 ++- routes_unified.go | 31 ++++--- routes_unified_test.go | 179 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 209 insertions(+), 12 deletions(-) diff --git a/limits.go b/limits.go index 65f2b1a..f80c9db 100644 --- a/limits.go +++ b/limits.go @@ -499,9 +499,14 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation { // evaluateScope checks one scope's already-set limits against its // just-incremented request counts and its accumulated token/cost counts. -// A nil limits field means nothing to evaluate — the scope's req counters -// still incremented above, feeding other scopes' aggregates and future -// observability, but this scope itself never blocks a request. +// A nil limits field means nothing to evaluate for this scope. In +// principle a nil-limits scope can still reach here with its req counters +// already incremented by checkAndCount, since that increment runs before +// any evaluation; in practice the unified route's buildLimitScopes +// (routes_unified.go, ruling e) never builds one — it omits a user or +// group scope from the slice entirely whenever that scope's own limits +// are nil — so this check is a defensive no-op today, not a path any +// current caller exercises. func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now time.Time) *limitViolation { if sc.limits == nil { return nil diff --git a/routes_unified.go b/routes_unified.go index 7011a6d..852fcd4 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -75,9 +75,12 @@ func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *us // runUnified is the shared chat/embeddings pipeline: decode the request, // resolve its model, enforce per-user and per-group limits, invoke the -// adapter via call, then account the resulting usage. w is wrapped in its -// own statusTrackingWriter so a mid-stream adapter error (headers already -// sent) can be told apart from one that failed before any write. +// adapter via call, then account the resulting usage — even when call +// itself returned an error, so usage captured before a mid-stream failure +// still gets billed — before translating that error into a response. w is +// wrapped in its own statusTrackingWriter so a mid-stream adapter error +// (headers already sent) can be told apart from one that failed before any +// write. func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, call adapterCall) { sw := &statusTrackingWriter{ResponseWriter: w} @@ -116,12 +119,18 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr req[gatewayAliasKey] = requestedModel result, callErr := call(adapter, r.Context(), sw, req) - if callErr != nil { - g.handleAdapterError(sw, callErr, adapter.name()) - return - } - if result.total() == 0 { + // Usage is accounted before the error branch below runs, not after: + // every adapter that can fail mid-stream (forwardStream in each of the + // three provider files) still returns whatever usage it had already + // captured alongside the error — a usage chunk that arrived just before + // a dropped connection must still be billed, or a client that aborts + // right after that frame arrives could repeat the trick to dodge every + // token/cost budget. A non-streaming failure (providerHTTPError, + // translateError, or a connection failure before any write) always + // carries zero usage by contract, so accounting it here is a no-op — + // account skips every write once both total tokens and cost are zero. + if callErr == nil && result.total() == 0 { if streaming { g.logf("unified route: zero usage reported for a streaming response from model %q; accounting the request only", canonical) } else { @@ -129,8 +138,12 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr result.estimated = true } } - g.limiter.account(scopes, result, unifiedCostMicros(canonical, upstreamModel, result, g.cfg.Pricing)) + + if callErr != nil { + g.handleAdapterError(sw, callErr, adapter.name()) + return + } } // buildLimitScopes returns the limitScope slice runUnified passes to the diff --git a/routes_unified_test.go b/routes_unified_test.go index c1b4d98..d5ece92 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -373,6 +373,101 @@ func TestHandleEmbeddings_HappyPath(t *testing.T) { } } +// TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody asserts the +// gatewayAliasKey convention field never reaches a real openai-type +// provider on POST /v1/chat/completions: unlike anthropic/gemini, this +// adapter marshals req verbatim as the upstream wire body, so a leftover +// alias entry would arrive as an unrecognized request field. This test +// fails if provider_openai.go's chatCompletion ever drops its +// delete(req, gatewayAliasKey) call. +func TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody(t *testing.T) { + const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` + + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + const requestedID = "openai/gpt-test" + body := map[string]any{"model": requestedID, "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + sentUpstream := decodeJSONBody(t, gotBody) + if _, present := sentUpstream[gatewayAliasKey]; present { + t.Errorf("gatewayAliasKey %q leaked into the openai chat upstream request body: %v", gatewayAliasKey, sentUpstream) + } + if sentUpstream["model"] != "gpt-test" { + t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "gpt-test") + } +} + +// TestHandleEmbeddings_OpenAI_AliasKeyStrippedFromUpstreamBody mirrors +// TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody for POST +// /v1/embeddings: this test fails if provider_openai.go's embeddings ever +// drops its delete(req, gatewayAliasKey) call. +func TestHandleEmbeddings_OpenAI_AliasKeyStrippedFromUpstreamBody(t *testing.T) { + const respBody = `{"object":"list","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":1,"total_tokens":1}}` + + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + const requestedID = "openai/text-embedding-3" + body := map[string]any{"model": requestedID, "input": "hello"} + req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + sentUpstream := decodeJSONBody(t, gotBody) + if _, present := sentUpstream[gatewayAliasKey]; present { + t.Errorf("gatewayAliasKey %q leaked into the openai embeddings upstream request body: %v", gatewayAliasKey, sentUpstream) + } + if sentUpstream["model"] != "text-embedding-3" { + t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "text-embedding-3") + } +} + // TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse // requests a provider-prefixed model id ("anthropic/claude-x"): the // upstream call must use the bare upstream id, but the translated @@ -556,3 +651,87 @@ func TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope(t *testing.T) { t.Errorf("want the mid-stream error logged, got %q", logBuf.String()) } } + +// TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage is the +// regression test for the critical accounting gap: an adapter that +// captures a usage frame before the connection drops must still have that +// usage billed, even though the request as a whole fails. The fake +// upstream sends one content chunk, then the terminal usage-only chunk +// (openaiAdapter.forwardStream always captures this into its returned +// usage, whether or not the client asked to see it), then hijacks and +// closes the connection without a clean chunked terminator — forcing +// readSSE to return a genuine error after the usage frame was already +// parsed. Without ruling (the CRITICAL fix), that captured usage would be +// silently discarded instead of reaching the limiter, letting a client +// dodge every token/cost budget by aborting right after the usage chunk +// arrives. +func TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage(t *testing.T) { + const contentChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n" + const usageChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3}}\n\n" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + _, _ = w.Write([]byte(contentChunk)) + fl.Flush() + _, _ = w.Write([]byte(usageChunk)) + fl.Flush() + + hj, ok := w.(http.Hijacker) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Hijack") + } + conn, _, hjErr := hj.Hijack() + if hjErr == nil { + _ = conn.Close() + } + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + origStderr := os.Stderr + pr, pw, _ := os.Pipe() + os.Stderr = pw + defer func() { os.Stderr = origStderr }() + + body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + _ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test + os.Stderr = origStderr + var logBuf bytes.Buffer + if _, err := io.Copy(&logBuf, pr); err != nil { + t.Fatalf("io.Copy: %v", err) + } + if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") { + t.Errorf("want the mid-stream error logged, got %q", logBuf.String()) + } + + tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 10 { + t.Errorf("user token/day counter = %d (ok=%v), want 10 (the usage chunk captured before the drop, 7 prompt + 3 completion)", tok, ok) + } +}