lukaszraczylo/traefik-llmgateway
1
0
Fork 0
mirror of https://github.com/lukaszraczylo/traefik-llmgateway.git synced 2026-10-12 04:54:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(routes): bill partial stream usage on error, pin openai alias stripping

Browse Source 
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:
lukaszraczylo
2026-08-20 10:55:39 +01:00
parent dcfcfb7e48
commit 3ed473bdfe
3 changed files with 209 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
limits.go
+8 -3
View File
@@ -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