mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 14:22:28 +08:00
fix: use totalUsage with all token types for accurate quota tracking
The onFinish callback's 'usage' only contains the final step's tokens, which underreports usage for multi-step tool calls (like diagram generation). Changed to 'totalUsage' which provides cumulative counts across all steps. Include all 4 token types for accurate counting: 1. inputTokens - non-cached input tokens 2. outputTokens - generated output tokens 3. cachedInputTokens - tokens read from prompt cache 4. inputTokenDetails.cacheWriteTokens - tokens written to cache Tested locally: - Request 1 (cache write): 334 + 62 + 0 + 6671 = 7,067 tokens - Request 2 (cache read): 334 + 184 + 6551 + 120 = 7,189 tokens - DynamoDB total: 14,256 ✓
This commit is contained in:
@@ -542,19 +542,24 @@ ${userInputText}
|
||||
userId,
|
||||
}),
|
||||
}),
|
||||
onFinish: ({ text, usage }) => {
|
||||
onFinish: ({ text, totalUsage }) => {
|
||||
// AI SDK 6 telemetry auto-reports token usage on its spans
|
||||
setTraceOutput(text)
|
||||
|
||||
// Record token usage for server-side quota tracking (if enabled)
|
||||
// Use totalUsage (cumulative across all steps) instead of usage (final step only)
|
||||
// Include all 4 token types: input, output, cache read, cache write
|
||||
if (
|
||||
isQuotaEnabled() &&
|
||||
!hasOwnApiKey &&
|
||||
userId !== "anonymous" &&
|
||||
usage
|
||||
totalUsage
|
||||
) {
|
||||
const totalTokens =
|
||||
(usage.inputTokens || 0) + (usage.outputTokens || 0)
|
||||
(totalUsage.inputTokens || 0) +
|
||||
(totalUsage.outputTokens || 0) +
|
||||
(totalUsage.cachedInputTokens || 0) +
|
||||
(totalUsage.inputTokenDetails?.cacheWriteTokens || 0)
|
||||
recordTokenUsage(userId, totalTokens)
|
||||
}
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user