Skip to content

Commit d5c9ad0

Browse files
authored
fix(ai): emit per-flush deltas to keep ai.toolCalls linear step count (#322)
1 parent 4f51eb6 commit d5c9ad0

4 files changed

Lines changed: 369 additions & 719 deletions

File tree

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import { bench, describe } from 'vitest'
2+
import type { LanguageModelV3 } from '@ai-sdk/provider'
3+
import { createAILogger } from '../../src/ai/index'
4+
import { createLogger } from '../../src/logger'
5+
import type { RequestLogger } from '../../src/types'
6+
7+
function makeMockResult(toolName: string) {
8+
return {
9+
content: [{ type: 'tool-call', toolCallId: `tc-${toolName}`, toolName, args: '{}' }],
10+
finishReason: { unified: 'tool-calls', raw: undefined },
11+
usage: {
12+
inputTokens: { total: 100, cacheRead: undefined, cacheWrite: undefined },
13+
outputTokens: { total: 20, reasoning: undefined },
14+
},
15+
response: { modelId: 'claude-sonnet-4.6' },
16+
}
17+
}
18+
19+
function makeWrappedModel(steps: number) {
20+
const calls: ReturnType<typeof makeMockResult>[] = []
21+
for (let i = 0; i < steps; i++) calls.push(makeMockResult(`tool-${i}`))
22+
let i = 0
23+
return {
24+
specificationVersion: 'v3',
25+
provider: 'anthropic',
26+
modelId: 'claude-sonnet-4.6',
27+
defaultObjectGenerationMode: 'json',
28+
doGenerate: () => Promise.resolve(calls[i++]),
29+
doStream: () => Promise.resolve({ stream: new ReadableStream() }),
30+
} as unknown as LanguageModelV3
31+
}
32+
33+
async function runScenario(steps: number, withSubscriber: boolean): Promise<void> {
34+
const log = createLogger()
35+
const ai = createAILogger(log as unknown as RequestLogger)
36+
if (withSubscriber) ai.onUpdate(() => { /* noop */ })
37+
const wrapped = ai.wrap(makeWrappedModel(steps))
38+
for (let i = 0; i < steps; i++) {
39+
await wrapped.doGenerate({} as any)
40+
}
41+
log.emit({ _forceKeep: true } as any)
42+
}
43+
44+
describe('multi-step agent flush throughput', () => {
45+
bench('6 steps, no subscriber', async () => {
46+
await runScenario(6, false)
47+
})
48+
49+
bench('6 steps, with subscriber', async () => {
50+
await runScenario(6, true)
51+
})
52+
53+
bench('50 steps, no subscriber', async () => {
54+
await runScenario(50, false)
55+
})
56+
57+
bench('50 steps, with subscriber', async () => {
58+
await runScenario(50, true)
59+
})
60+
})

0 commit comments

Comments
 (0)