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).
2026-10-14 05:18:54 +00:00 · 2026-08-20 10:55:39 +01:00
parent dcfcfb7e48
commit 3ed473bdfe
3 changed files with 209 additions and 12 deletions
@@ -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
@@ -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
@@ -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)
+	}
+}