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
+179
View File
@@ -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)
}
}