fix: accumulate token usage across retry attempts
Previously only the final attempt's tokenUsage was returned, causing under-reporting of actual model consumption when retries occurred. Now all attempts' token counts are summed in the returned result. Addresses Codex review P2 (token usage) on #37
This commit is contained in:
parent
08cf01c6b4
commit
63e1f7068a
|
|
@ -133,12 +133,20 @@ export async function executeWithRetry(
|
||||||
const backoff = Math.max(1, task.retryBackoff ?? 2)
|
const backoff = Math.max(1, task.retryBackoff ?? 2)
|
||||||
|
|
||||||
let lastError: string = ''
|
let lastError: string = ''
|
||||||
|
// Accumulate token usage across all attempts so billing/observability
|
||||||
|
// reflects the true cost of retries.
|
||||||
|
let totalUsage: TokenUsage = { input_tokens: 0, output_tokens: 0 }
|
||||||
|
|
||||||
for (let attempt = 1; attempt <= maxAttempts; attempt++) {
|
for (let attempt = 1; attempt <= maxAttempts; attempt++) {
|
||||||
try {
|
try {
|
||||||
const result = await run()
|
const result = await run()
|
||||||
|
totalUsage = {
|
||||||
|
input_tokens: totalUsage.input_tokens + result.tokenUsage.input_tokens,
|
||||||
|
output_tokens: totalUsage.output_tokens + result.tokenUsage.output_tokens,
|
||||||
|
}
|
||||||
|
|
||||||
if (result.success) {
|
if (result.success) {
|
||||||
return result
|
return { ...result, tokenUsage: totalUsage }
|
||||||
}
|
}
|
||||||
lastError = result.output
|
lastError = result.output
|
||||||
|
|
||||||
|
|
@ -150,7 +158,7 @@ export async function executeWithRetry(
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
return result
|
return { ...result, tokenUsage: totalUsage }
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
lastError = err instanceof Error ? err.message : String(err)
|
lastError = err instanceof Error ? err.message : String(err)
|
||||||
|
|
||||||
|
|
@ -166,7 +174,7 @@ export async function executeWithRetry(
|
||||||
success: false,
|
success: false,
|
||||||
output: lastError,
|
output: lastError,
|
||||||
messages: [],
|
messages: [],
|
||||||
tokenUsage: { input_tokens: 0, output_tokens: 0 },
|
tokenUsage: totalUsage,
|
||||||
toolCalls: [],
|
toolCalls: [],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -177,7 +185,7 @@ export async function executeWithRetry(
|
||||||
success: false,
|
success: false,
|
||||||
output: lastError,
|
output: lastError,
|
||||||
messages: [],
|
messages: [],
|
||||||
tokenUsage: { input_tokens: 0, output_tokens: 0 },
|
tokenUsage: totalUsage,
|
||||||
toolCalls: [],
|
toolCalls: [],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -276,6 +276,58 @@ describe('executeWithRetry', () => {
|
||||||
expect(mockDelay).toHaveBeenCalledWith(30_000) // capped
|
expect(mockDelay).toHaveBeenCalledWith(30_000) // capped
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('accumulates token usage across retry attempts', async () => {
|
||||||
|
const failResult: AgentRunResult = {
|
||||||
|
...FAILURE_RESULT,
|
||||||
|
tokenUsage: { input_tokens: 100, output_tokens: 50 },
|
||||||
|
}
|
||||||
|
const successResult: AgentRunResult = {
|
||||||
|
...SUCCESS_RESULT,
|
||||||
|
tokenUsage: { input_tokens: 200, output_tokens: 80 },
|
||||||
|
}
|
||||||
|
|
||||||
|
const run = vi.fn()
|
||||||
|
.mockResolvedValueOnce(failResult)
|
||||||
|
.mockResolvedValueOnce(failResult)
|
||||||
|
.mockResolvedValueOnce(successResult)
|
||||||
|
|
||||||
|
const task = createTask({
|
||||||
|
title: 'Token test',
|
||||||
|
description: 'test',
|
||||||
|
maxRetries: 2,
|
||||||
|
retryDelayMs: 10,
|
||||||
|
})
|
||||||
|
|
||||||
|
const result = await executeWithRetry(run, task, undefined, noDelay)
|
||||||
|
|
||||||
|
expect(result.success).toBe(true)
|
||||||
|
// 100+100+200 input, 50+50+80 output
|
||||||
|
expect(result.tokenUsage.input_tokens).toBe(400)
|
||||||
|
expect(result.tokenUsage.output_tokens).toBe(180)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('accumulates token usage even when all retries fail', async () => {
|
||||||
|
const failResult: AgentRunResult = {
|
||||||
|
...FAILURE_RESULT,
|
||||||
|
tokenUsage: { input_tokens: 50, output_tokens: 30 },
|
||||||
|
}
|
||||||
|
|
||||||
|
const run = vi.fn().mockResolvedValue(failResult)
|
||||||
|
|
||||||
|
const task = createTask({
|
||||||
|
title: 'Token fail test',
|
||||||
|
description: 'test',
|
||||||
|
maxRetries: 1,
|
||||||
|
})
|
||||||
|
|
||||||
|
const result = await executeWithRetry(run, task, undefined, noDelay)
|
||||||
|
|
||||||
|
expect(result.success).toBe(false)
|
||||||
|
// 50+50 input, 30+30 output (2 attempts)
|
||||||
|
expect(result.tokenUsage.input_tokens).toBe(100)
|
||||||
|
expect(result.tokenUsage.output_tokens).toBe(60)
|
||||||
|
})
|
||||||
|
|
||||||
it('clamps negative maxRetries to 0 (single attempt)', async () => {
|
it('clamps negative maxRetries to 0 (single attempt)', async () => {
|
||||||
const run = vi.fn().mockRejectedValue(new Error('fail'))
|
const run = vi.fn().mockRejectedValue(new Error('fail'))
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue