fix(cost-estimation): fix zero output tokens - #913
Conversation
![terragon-labs[bot]](https://github.githubassets.com/images/gravatars/gravatar-user-420.png){kind=avatar}
- Add support for toolResults in calculateCosts to count tokens from tool calls - Update chat.ts to pass toolResults in completions - Add tests for toolResults handling and edge cases in costs.spec.ts - Handle null and zero token counts gracefully to ensure accurate cost calculation This enhancement allows more accurate cost estimation by including tokens from tool function calls in the calculation. Co-authored-by: terragon-labs[bot] <terragon-labs[bot]@users.noreply.github.com>
❌ Preview Environment deleted from BunnyshellAvailable commands (reply to this comment):
|
WalkthroughPropagates tool call results (streamingToolCalls/toolResults) into provider request payloads for streaming and non‑streaming chat paths, and extends cost calculation to include toolResults when estimating completion tokens. Adds tests covering zero/missing tokens and toolResults-driven estimation. No public signatures removed; Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant C as Client
participant GW as Gateway Chat Handler
participant P as Provider API
participant Cost as Costs.calculateCosts
C->>GW: Send chat request
rect rgba(230,245,255,0.5)
note right of GW: Collect tool call results (streamingToolCalls or final toolResults)
GW->>P: Provider request { messages, ..., toolResults? }
P-->>GW: Response { completion?/stream, tool calls? }
end
alt Streaming or Non-streaming finalization
GW->>Cost: calculateCosts({ prompt, completion?, toolResults? })
Cost-->>GW: { inputCost, outputCost, totalCost, estimatedCost }
end
GW-->>C: Final response (+ costs)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
github-actions
Bot
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
apps/gateway/src/lib/costs.ts (1)
124-138: Nullish check: don’t treat 0 prompt tokens as “missing”Use
== nullso a valid 0 doesn’t early‑return and drop request costs.-// If we don't have prompt tokens, we can't calculate any costs -if (!calculatedPromptTokens) { +// If prompt tokens are absent, we can't calculate costs +if (calculatedPromptTokens == null) {
🧹 Nitpick comments (6)
apps/gateway/src/lib/costs.ts (2)
93-121: Avoid double-stringifying args; add separators to reduce token-estimation biasConcatenation without delimiters glues tokens; JSON.stringify over an already‑string args adds quotes/escapes. Build parts, handle string vs object, then join with a delimiter.
- if (!completionTokens && fullOutput) { - let completionText = ""; - - // Include main completion content - if (fullOutput.completion) { - completionText += fullOutput.completion; - } - - // Include tool results if available - if (fullOutput.toolResults && Array.isArray(fullOutput.toolResults)) { - for (const toolResult of fullOutput.toolResults) { - if (toolResult.function?.name) { - completionText += toolResult.function.name; - } - if (toolResult.function?.arguments) { - completionText += JSON.stringify(toolResult.function.arguments); - } - } - } - - if (completionText) { - try { - calculatedCompletionTokens = encode(completionText).length; - } catch (error) { - // If encoding fails, leave as null - logger.error(`Failed to encode completion text: ${error}`); - } - } - } + if (!completionTokens && fullOutput) { + const parts: string[] = []; + if (fullOutput.completion) parts.push(fullOutput.completion); + if (Array.isArray(fullOutput.toolResults)) { + for (const tr of fullOutput.toolResults) { + const name = tr.function?.name; + if (name) parts.push(`tool:${name}`); + const args = tr.function?.arguments; + if (typeof args === "string") { + parts.push(`args:${args}`); + } else if (args != null) { + parts.push(`args:${JSON.stringify(args)}`); + } + } + } + if (parts.length) { + try { + calculatedCompletionTokens = encode(parts.join(" ")).length; + } catch (error) { + logger.error(`Failed to encode completion text: ${error}`); + } + } + }
140-144: Use nullish check for completion fallbackThis avoids unnecessary reassignment and is consistent with the prompt check.
-// Set completion tokens to 0 if not available (but still calculate input costs) -if (!calculatedCompletionTokens) { - calculatedCompletionTokens = 0; -} +// Set completion tokens to 0 if not available (but still calculate input costs) +if (calculatedCompletionTokens == null) { + calculatedCompletionTokens = 0; +}apps/gateway/src/lib/costs.spec.ts (4)
175-184: Zero-output tokens case covered; tighten assertions for stabilityMake precision explicit on close‑to checks and use exact zero where applicable to avoid flaky comparisons.
- expect(result.inputCost).toBeCloseTo(0.001); // 100 * 0.00001 - expect(result.outputCost).toBeCloseTo(0); // 0 * 0.00003 - expect(result.totalCost).toBeCloseTo(0.001); // 0.001 + 0 + expect(result.inputCost).toBeCloseTo(0.001, 6); // 100 * 0.00001 + expect(result.outputCost).toBe(0); // 0 * 0.00003 + expect(result.totalCost).toBeCloseTo(0.001, 6); // 0.001 + 0
197-226: Assert the incremental effect of toolResults vs. no toolsStrengthen this by comparing against the same input without toolResults to prove tokens/costs actually increase due to tools.
@@ it("should include tool results in completion token estimation", () => { @@ expect(result.estimatedCost).toBe(true); }); + + it("should increase completion tokens when toolResults are present vs none", () => { + const base = { + prompt: "What's the weather like?", + completion: "", + }; + const withTools = calculateCosts("gpt-4", "openai", null, null, null, { + ...base, + toolResults: [ + { id: "call_1", function: { name: "get_weather", arguments: '{"location":"SF"}' } }, + ], + }); + const withoutTools = calculateCosts("gpt-4", "openai", null, null, null, { + ...base, + toolResults: [], + }); + expect(withTools.completionTokens).toBeGreaterThan(withoutTools.completionTokens); + expect(withTools.outputCost!).toBeGreaterThan(withoutTools.outputCost!); + });
205-214: Cover object-form arguments and avoid double-stringify in implementationThese fixtures pass arguments as strings. In calculateCosts the code JSON.stringifys arguments, which can double‑quote strings and inflate tokens. Add a test with arguments as an object, and (in costs.ts) only stringify when not already a string.
Add a test:
+ it("should handle tool results where arguments are objects", () => { + const result = calculateCosts("gpt-4", "openai", null, null, null, { + prompt: "q", + completion: "", + toolResults: [ + { id: "call_obj", function: { name: "get_weather", arguments: { location: "SF" } } }, + ], + }); + expect(result.completionTokens).toBeGreaterThan(0); + expect(result.estimatedCost).toBe(true); + });Update implementation (apps/gateway/src/lib/costs.ts):
- if (toolResult.function?.arguments) { - completionText += JSON.stringify(toolResult.function.arguments); - } + if (toolResult.function?.arguments !== undefined) { + const args = toolResult.function.arguments; + completionText += typeof args === "string" ? args : JSON.stringify(args); + }
228-249: Add a case with only name or only argumentsThis test covers “missing function” and “empty function” but not “only name” or “only arguments”. Add one of each to harden behavior.
@@ - toolResults: [ + toolResults: [ { id: "call_1" }, // Missing function data { id: "call_2", function: {} }, // Missing name and arguments + { id: "call_2b", function: { name: "ping" } }, // only name + { id: "call_2c", function: { arguments: { foo: "bar" } } }, // only arguments { id: "call_3", function: { name: "get_weather", arguments: '{"location": "Paris"}', }, }, ],
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
apps/gateway/src/chat/chat.ts(2 hunks)apps/gateway/src/lib/costs.spec.ts(1 hunks)apps/gateway/src/lib/costs.ts(3 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
{apps/api,apps/gateway,apps/ui,apps/docs,packages}/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Always use top-level import; never use require() or dynamic imports (e.g., import(), next/dynamic)
Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/costs.ts
{apps/api,apps/gateway,packages/db}/**/*.ts
📄 CodeRabbit inference engine (AGENTS.md)
.findMany() or db().query.
{apps/api,apps/gateway,packages/db}/**/*.ts: Use Drizzle ORM with the latest object syntax for database access
For reads, use db().query..findFirst() Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/costs.ts**/*.spec.ts
📄 CodeRabbit inference engine (AGENTS.md)
Place unit tests in files named *.spec.ts
Unit test files must be named with the .spec.ts suffix
Files:
apps/gateway/src/lib/costs.spec.ts**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.{ts,tsx}: Never useanyoras anyin this TypeScript project unless absolutely necessary
Always use top-levelimport; never userequireor dynamic importsFiles:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/costs.ts{apps/api,apps/gateway}/src/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
For reads, use db().query.
.findMany() or db().query.
.findFirst() Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/costs.ts🧬 Code graph analysis (3)
apps/gateway/src/lib/costs.spec.ts (1)
apps/gateway/src/lib/costs.ts (1)
calculateCosts(20-198)apps/gateway/src/chat/chat.ts (1)
packages/db/src/types.ts (1)
toolResults(45-45)apps/gateway/src/lib/costs.ts (1)
packages/logger/src/index.ts (2)
error(147-154)logger(175-175)⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: e2e-shards (3)
- GitHub Check: e2e-shards (2)
- GitHub Check: e2e-shards (5)
- GitHub Check: e2e-shards (4)
- GitHub Check: e2e-shards (1)
- GitHub Check: build / run
- GitHub Check: test / run
- GitHub Check: generate / run
- GitHub Check: lint / run
- GitHub Check: autofix
🔇 Additional comments (3)
apps/gateway/src/lib/costs.ts (1)
171-185: Clamp uncachedPromptTokens to prevent negative costsClamp uncachedPromptTokens to 0 so it can't go negative when cachedTokens > calculatedPromptTokens.
File: apps/gateway/src/lib/costs.ts (lines 171–185)
const uncachedPromptTokens = Math.max( 0, cachedTokens ? calculatedPromptTokens - cachedTokens : calculatedPromptTokens, );Quick repo check for the pattern "cachedTokens:\s*\d+\s*||\s*null" in apps/gateway/src/chat returned 0 matches.
apps/gateway/src/chat/chat.ts (1)
2654-2657: Passing toolResults into cost estimation — good addition
- Both streaming and non‑streaming paths already pass toolResults into calculateCosts (apps/gateway/src/chat/chat.ts — calls at ~2647 and ~3034).
- calculateCosts currently types fullOutput.toolResults as any[] (apps/gateway/src/lib/costs.ts); change that to a proper ToolResult[] and ensure the element shape matches the tool objects you build in chat.ts (at minimum: { id, function?: { name?: string; arguments?: string | object } }).
apps/gateway/src/lib/costs.spec.ts (1)
186-195: LGTM — assertions align with new defaulting behaviorCorrectly verifies that completionTokens default to 0 and estimatedCost remains false without fullOutput.
| prompt: messages.map((m) => m.content).join("\n"), | ||
| completion: content, | ||
| toolResults: toolResults, | ||
| }, |
There was a problem hiding this comment.
Non‑streaming path parity — good; consider redacting secrets before logging
You’re also attaching toolResults to logs elsewhere; tool args may contain secrets/PII. Redact keys like token, authorization, apiKey, password before persisting.
Example helper (place in a shared util and use before insertLog/calculateCosts if desired):
const SENSITIVE_KEYS = /^(authorization|api[-_]?key|token|access[_-]?token|secret|password)$/i;
function redactDeep(value: unknown): unknown {
if (Array.isArray(value)) return value.map(redactDeep);
if (value && typeof value === "object") {
return Object.fromEntries(
Object.entries(value as Record<string, unknown>).map(([k, v]) => [
k,
SENSITIVE_KEYS.test(k) ? "***" : redactDeep(v),
]),
);
}
return value;
}
// usage: toolResults?.map(tr => ({ ...tr, function: tr.function ? { ...tr.function, arguments: typeof tr.function.arguments === 'string' ? tr.function.arguments : redactDeep(tr.function.arguments) } : tr.function }))🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 3041 to 3044, toolResults are
being attached to logs and may contain secrets/PII; before persisting or calling
insertLog/calculateCosts, deep-redact sensitive keys (authorization, apiKey,
token, access_token, secret, password, etc.) from toolResults and any nested
objects or arrays. Add or reuse a shared util (redactDeep) that recursively
traverses arrays/objects, replaces matching keys with a placeholder (e.g.,
"***"), and apply it to toolResults and to tool.function.arguments (handling
string vs structured args) so logs never store raw secrets. Ensure the redaction
runs in the non-streaming and streaming paths just prior to any log or storage
call.
| toolResults?: any[]; | ||
| }, |
There was a problem hiding this comment.
Replace any[] with a proper ToolResult type
Project guideline forbids any. Define a minimal ToolResult shape and use it here.
interface ChatMessage {
role: "user" | "system" | "assistant" | undefined;
content: string;
name?: string;
}
+
+// Minimal tool result shape (align with OpenAI tool call)
+interface ToolFunctionCall {
+ name?: string;
+ arguments?: string | Record<string, unknown>;
+}
+interface ToolResult {
+ id?: string;
+ function?: ToolFunctionCall;
+}
@@
completion?: string;
- toolResults?: any[];
+ toolResults?: ToolResult[];📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| toolResults?: any[]; | |
| }, | |
| interface ChatMessage { | |
| role: "user" | "system" | "assistant" | undefined; | |
| content: string; | |
| name?: string; | |
| } | |
| // Minimal tool result shape (align with OpenAI tool call) | |
| interface ToolFunctionCall { | |
| name?: string; | |
| arguments?: string | Record<string, unknown>; | |
| } | |
| interface ToolResult { | |
| id?: string; | |
| function?: ToolFunctionCall; | |
| } | |
| completion?: string; | |
| toolResults?: ToolResult[]; | |
| }, |
🤖 Prompt for AI Agents
In apps/gateway/src/lib/costs.ts around lines 30-31, the toolResults property is
typed as any[] which violates project guidelines; define a minimal ToolResult
interface (e.g., { name: string; durationMs?: number; costCents?: number; meta?:
Record<string, unknown> }) either in this file or a shared types file, export it
if needed, and replace toolResults?: any[] with toolResults?: ToolResult[];
update any code consuming this property to match the new shape and add
imports/exports as appropriate.
…on tests - Added 'type: function' to toolResults entries in costs.spec.ts to improve test clarity and type safety. - Updated costs.ts to use ToolCall type for toolResults parameter, enhancing type correctness in cost calculations. Co-authored-by: terragon-labs[bot] <terragon-labs[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
apps/gateway/src/lib/costs.spec.ts (1)
235-236: Verify type assertion usage against coding guidelinesThe test uses
as anytype assertions which may violate the project's TypeScript guidelines that discourageanyusage. Consider using more specific typing or test utilities.Based on the search results, I can see that TypeScript provides
Partial<T>utility type to facilitate common type transformations, which is a better approach than usingas anyfor testing with incomplete objects. Here's my suggestion:Replace the
as anytype assertions withPartial<ToolCall>to maintain type safety:toolResults: [ - { id: "call_1", type: "function" } as any, // Missing function data - { id: "call_2", type: "function", function: {} as any }, // Missing name and arguments + { id: "call_1", type: "function" } as Partial<ToolCall>, // Missing function data + { id: "call_2", type: "function", function: {} } as Partial<ToolCall> & { function: {} }, // Missing name and arguments { id: "call_3", type: "function", function: { name: "get_weather", arguments: '{"location": "Paris"}', }, }, ],This approach provides better type safety while still allowing partial objects for testing incomplete tool data scenarios.
apps/gateway/src/lib/costs.ts (1)
113-113: Don't JSON.stringify toolResult.function.arguments blindly — append raw if string, otherwise safe-stringifytool arguments are generally stored/used as JSON strings in the codebase (e.g. apps/gateway/src/anthropic/anthropic.ts, chat.ts and e2e tests), so JSON.stringify here will double-encode strings and may throw for non-serializable/circular objects. Replace the single JSON.stringify call with a short guard: if typeof arguments === "string" append it as-is; otherwise try JSON.stringify(...) and fall back to String(...) on error.
Location: apps/gateway/src/lib/costs.ts:113
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
apps/gateway/src/lib/costs.spec.ts(1 hunks)apps/gateway/src/lib/costs.ts(4 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
{apps/api,apps/gateway,apps/ui,apps/docs,packages}/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Always use top-level import; never use require() or dynamic imports (e.g., import(), next/dynamic)
Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/lib/costs.ts
{apps/api,apps/gateway,packages/db}/**/*.ts
📄 CodeRabbit inference engine (AGENTS.md)
.findMany() or db().query.
{apps/api,apps/gateway,packages/db}/**/*.ts: Use Drizzle ORM with the latest object syntax for database access
For reads, use db().query..findFirst() Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/lib/costs.ts**/*.spec.ts
📄 CodeRabbit inference engine (AGENTS.md)
Place unit tests in files named *.spec.ts
Unit test files must be named with the .spec.ts suffix
Files:
apps/gateway/src/lib/costs.spec.ts**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.{ts,tsx}: Never useanyoras anyin this TypeScript project unless absolutely necessary
Always use top-levelimport; never userequireor dynamic importsFiles:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/lib/costs.ts{apps/api,apps/gateway}/src/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
For reads, use db().query.
.findMany() or db().query.
.findFirst() Files:
apps/gateway/src/lib/costs.spec.tsapps/gateway/src/lib/costs.ts🧬 Code graph analysis (2)
apps/gateway/src/lib/costs.spec.ts (1)
apps/gateway/src/lib/costs.ts (1)
calculateCosts(25-203)apps/gateway/src/lib/costs.ts (2)
packages/models/src/types.ts (1)
ToolCall(54-61)packages/logger/src/index.ts (2)
error(147-154)logger(175-175)⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: e2e-shards (2)
- GitHub Check: e2e-shards (4)
- GitHub Check: e2e-shards (1)
- GitHub Check: e2e-shards (5)
- GitHub Check: e2e-shards (3)
- GitHub Check: build / run
- GitHub Check: test / run
- GitHub Check: generate / run
- GitHub Check: lint / run
- GitHub Check: autofix
🔇 Additional comments (8)
apps/gateway/src/lib/costs.ts (4)
4-9: LGTM: Proper import of ToolCall typeThe addition of ToolCall import from
@llmgateway/modelsaligns with the codebase structure and provides proper typing for the new toolResults feature.
35-35: Address the existing type safety issueA previous review identified that using
any[]violates project guidelines. The current code now properly typestoolResultsasToolCall[], which resolves the previous concern about type safety.
98-127: LGTM: Robust tool results integration with proper error handlingThe implementation correctly includes tool results in completion token estimation by:
- Concatenating function names and JSON-stringified arguments
- Gracefully handling missing or partial function data with optional chaining
- Including proper error logging for encoding failures
- Maintaining existing completion text logic
The approach ensures accurate cost estimation when tool calls are involved.
129-148: LGTM: Improved cost calculation logic for zero output tokensThe changes correctly address the core issue by:
- Only requiring prompt tokens to proceed with cost calculations (Line 130)
- Defaulting completion tokens to 0 when missing, allowing input costs to be calculated (Lines 146-148)
- Maintaining the existing behavior for cases where prompt tokens are unavailable
This ensures input costs are calculated even when output tokens are zero or null, which was the main objective of this PR.
apps/gateway/src/lib/costs.spec.ts (4)
175-184: LGTM: Comprehensive test for zero output tokensThe test correctly verifies that input costs are calculated when output tokens are zero, confirming the fix for the main issue addressed in this PR.
186-195: LGTM: Essential test for null completion tokens with existing prompt tokensThis test validates the key scenario where completion tokens are missing but prompt tokens exist, ensuring input costs are still calculated with completion tokens defaulted to 0.
197-228: LGTM: Thorough test coverage for tool results integrationThe test effectively validates:
- Tool function names and arguments are included in token estimation
- Multiple tool calls are properly handled
- Token counts and costs reflect tool call content
- Estimated cost flag is correctly set
230-252: LGTM: Robust error handling test for partial tool dataThe test properly verifies graceful handling of incomplete tool result data, including:
- Missing function data entirely
- Empty function objects
- Mixed valid/invalid tool results
This ensures the cost calculation remains stable even with malformed tool data.
1 participant
Summary
Changes
Cost Calculation Logic
calculateCoststo:Testing
costs.spec.tsto cover:Chat Completion Updates
toolResultsto completion data for accurate cost trackingTest plan
🌿 Generated by Terry
ℹ️ Tag @terragon-labs to ask questions and address PR feedback
📎 Task: https://www.terragonlabs.com/task/a6d0c0ea-10ca-4c74-9aea-69bcedf18ebf
Summary by CodeRabbit