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).
This commit is contained in:
2026-08-20 10:55:39 +01:00
parent dcfcfb7e48
commit 3ed473bdfe
3 changed files with 209 additions and 12 deletions
+22 -9
View File
@@ -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