fix: add cache write accounting to OpenAI interceptors by pawbana · Pull Request #28567 · coder/coder · GitHub
Skip to content

fix: add cache write accounting to OpenAI interceptors - #28567

Merged
pawbana merged 4 commits into
mainfrom
pawel/aigov-630-account-for-cache-write-tokens-in-ai-gateway-costs
Aug 27, 2026
Merged

fix: add cache write accounting to OpenAI interceptors#28567
pawbana merged 4 commits into
mainfrom
pawel/aigov-630-account-for-cache-write-tokens-in-ai-gateway-costs

Conversation

@pawbana

@pawbana pawbana commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Adds cache write token accounting to OpenAI interceptors.
This should automatically be picked up by const control.

@linear-code

linear-code Bot commented Aug 25, 2026

Copy link
Copy Markdown

Comment thread aibridge/fixtures/openai/chatcompletions/simple.txtar
@pawbana
pawbana requested a review from ssncferreira August 26, 2026 13:54
@pawbana pawbana changed the title feat: add cache write accounting to OpenAI interceptors fix: add cache write accounting to OpenAI interceptors Aug 26, 2026
"rejected_prediction_tokens": 7
}
},
"service_tier": "default",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah interesting, we already have service_tier in the fixtures 👀

lastErr = xerrors.Errorf("marshal chunk: %w", err)
break
}
if err := events.Send(ctx, payload); err != nil {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure if this is correct 🤔 so we relay the chunks back to the client, and in each iteration we keep accumulating it?
For instance, assume 3 iterations:

  • Iteration 1: previousUsage = 0, prc.lastUsage = u1 => cumulativeUsage: u1 => token usage u1 send to the client
  • Iteration 2: previousUsage = u1, prc.lastUsage = u2 => cumulativeUsage: u1+u2 => token usage u1+u2 send to the client
  • Iteration 3: previousUsage = u1+u2, prc.lastUsage = u3 => cumulativeUsage: u1+u2+u3 => token usage u1+u2+u3 send to the client

So basically the client would receive something like:

data: {... "usage": u1}
data: {... "usage": u1+u2}
data: {... "usage": u1+u2+u3}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is WAI.
It used to work like that (or worked if it stream would only contain 1 usage chunk) until ad10045

Linked commit fixed issue where some OpenAI compatible providers return multiple usages for one request which resulted in old solution counting all usages from all chunks from all streams which was too much but also dropped the sum of all inner loop requests feature.

@ssncferreira ssncferreira left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM 🤞 just a question about a case that I'm not sure can happen

lastUsage := processor.lastUsage
// Preserve the existing recording gate, but accumulate every valid
// snapshot so later client-visible usage cannot move backwards.
if lastUsage.CompletionTokens > 0 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see this was already before 👀 but could it happen that CompletionTokens == 0 but cache read/write > 0?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test is kinda related to the above comment, it is only cumulative because the previous usage has a completion_tokens > 0. If all chunks have completion_tokens = 0, there would be no accumulation, which would also affect cost control, because we don't store RecordTokenUsage 🤔

@pawbana
pawbana force-pushed the pawel/aigov-630-account-for-cache-write-tokens-in-ai-gateway-costs branch from 87bd945 to ea6b520 Compare August 27, 2026 09:45

@ssncferreira ssncferreira left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM 🤞

…recording token usage in chat completion interceptor and improve TestStreamProcessorZeroCompletionUsage test to accutally test something
@pawbana
pawbana force-pushed the pawel/aigov-630-account-for-cache-write-tokens-in-ai-gateway-costs branch from ea6b520 to 6162b6a Compare August 27, 2026 10:37
@pawbana
pawbana merged commit 3ca2ba4 into main Aug 27, 2026
28 checks passed
@pawbana
pawbana deleted the pawel/aigov-630-account-for-cache-write-tokens-in-ai-gateway-costs branch August 27, 2026 10:49
@github-actions github-actions Bot locked and limited conversation to collaborators Aug 27, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants