mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-21 05:57:53 +00:00
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:
+22
-9
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user