mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
Michael Suchacz
ad100452d4
## Problem CODAGT-906: chats using OpenAI-compatible backends (e.g. poolside) through the AI Bridge persist token usage inflated 105x-640x, which falsely triggers automatic chat compaction on every turn. The chat-completions streaming interceptor summed usage across every SSE chunk of one upstream stream and rewrote each relayed usage-bearing chunk with that running sum. Spec-compliant OpenAI emits usage once (final chunk with `stream_options.include_usage`), so the sum equals the final value. vLLM-style backends emit cumulative usage snapshots on every chunk, so the relayed final usage becomes roughly `N_chunks x prompt_tokens` (e.g. 417,012 persisted for a ~6,000-token context). chatd persists that value per assistant message and its compaction trigger reads it as context occupancy. ## Fix Track the latest usage-bearing chunk's raw usage (last-wins) in the stream processor, updating only when a chunk actually carries usage so a trailing usage-less chunk cannot zero it. `marshalChunk` relays that value and `recordTokenUsage` records the same value, unifying relayed and recorded usage. Last-wins is correct for both shapes: a single final usage chunk, and cumulative snapshots where each snapshot already includes all prior tokens. Per-iteration semantics are unchanged: each tool-loop iteration has its own processor, and the final iteration's usage is what the client sees. ## Tests - `TestStreamProcessorUsage` (internal): cumulative snapshots with a trailing usage-less chunk, and the spec-compliant final-only shape; asserts relayed and recorded usage equal the last snapshot. - New txtar fixture `streaming_cumulative_usage_injected_tool.txtar` with per-chunk cumulative usage plus an injected tool call; asserts client-visible final usage through the full interceptor. - Red-green verified: with the fix reverted, the internal test reports zeroed usage (trailing chunk overwrite) and the fixture test reports 18000 summed prompt tokens instead of 6000. The blocking (non-streaming) path deliberately keeps its cross-iteration summation for external clients and is untouched. Remote dogfood UAT validated chat streaming, tool calls, plausible usage numbers, and zero spurious compactions. > Mux acted on Mike's behalf to author this change.
These fixtures were created by adding logging middleware to API calls to view the raw requests/responses.
...
opts = append(opts, option.WithMiddleware(LoggingMiddleware))
...
func LoggingMiddleware(req *http.Request, next option.MiddlewareNext) (res *http.Response, err error) {
reqOut, _ := httputil.DumpRequest(req, true)
// Forward the request to the next handler
res, err = next(req)
fmt.Printf("[req] %s\n", reqOut)
// Handle stuff after the request
if err != nil {
return res, err
}
respOut, _ := httputil.DumpResponse(res, true)
fmt.Printf("[resp] %s\n", respOut)
return res, err
}