fix(cli): align ingest step counter with SDK num_turns

The Claude Code runtime counted every SDKAssistantMessage with
parent_tool_use_id === null as a step, but the SDK emits extra messages
within a single num_turns round-trip — `stop_reason: 'pause_turn'`
continuations and errored partials it retries internally. The local
counter then outran maxTurns and the ingest HUD rendered confusing
ratios like `step 69/40`.

Filter both cases in collectResult so stepIndex tracks num_turns and
stays bounded by the work-unit stepBudget.
This commit is contained in:
Andrey Avtomonov 2026-05-28 01:55:39 +02:00
parent a94f35800a
commit 20454bcb8f
3 changed files with 77 additions and 3 deletions

View file

@ -58,6 +58,22 @@ function isResult(message: SDKMessage): message is SDKResultMessage {
return message.type === 'result';
}
// Skip emissions the SDK does not count toward `num_turns`: `pause_turn` continuations and
// errored partials (e.g. `max_output_tokens`) it retries internally. Without this, the
// runtime's step counter outruns `maxTurns` and the HUD renders e.g. `step 69/40`.
function countsAsAssistantTurn(message: SDKMessage): boolean {
if (message.type !== 'assistant' || message.parent_tool_use_id !== null) {
return false;
}
if (message.error !== undefined) {
return false;
}
if (message.message.stop_reason === 'pause_turn') {
return false;
}
return true;
}
function resultError(result: SDKResultMessage): Error | undefined {
if (result.subtype === 'success') {
return undefined;
@ -190,7 +206,7 @@ async function collectResult(params: {
let result: SDKResultMessage | undefined;
for await (const message of params.query({ prompt: params.prompt, options: params.options })) {
assertInitIsolation(message, params.allowedToolIds, params.expectedMcpServerNames);
if (message.type === 'assistant' && message.parent_tool_use_id === null) {
if (countsAsAssistantTurn(message)) {
await params.onAssistantTurn?.();
}
if (isResult(message)) {

View file

@ -415,6 +415,64 @@ describe('ClaudeCodeKtxLlmRuntime', () => {
);
});
it('counts only assistant turns the SDK counts toward num_turns', async () => {
const assistantMessage = (
overrides: Partial<Extract<SDKMessage, { type: 'assistant' }>> & { uuid: string },
): SDKMessage =>
({
type: 'assistant',
message: { role: 'assistant', content: [], stop_reason: 'end_turn' },
parent_tool_use_id: null,
session_id: 'session-id',
...overrides,
}) as unknown as SDKMessage;
const query = vi.fn((_input: any) =>
stream([
initMessage(),
assistantMessage({
uuid: '00000000-0000-4000-8000-0000000000a1',
error: 'max_output_tokens',
}),
assistantMessage({
uuid: '00000000-0000-4000-8000-0000000000a2',
message: { role: 'assistant', content: [], stop_reason: 'pause_turn' } as never,
}),
assistantMessage({ uuid: '00000000-0000-4000-8000-0000000000a3' }),
{
type: 'assistant',
message: { role: 'assistant', content: [], stop_reason: 'end_turn' },
parent_tool_use_id: 'tool-use-1',
uuid: '00000000-0000-4000-8000-0000000000a4',
session_id: 'session-id',
} as unknown as SDKMessage,
resultMessage({ subtype: 'success', terminal_reason: 'completed' }),
]),
);
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
});
const onStepFinish = vi.fn();
await expect(
runtime.runAgentLoop({
modelRole: 'default',
systemPrompt: 'system',
userPrompt: 'user',
toolSet: {},
stepBudget: 40,
telemetryTags: { operationName: 'test' },
onStepFinish,
}),
).resolves.toEqual({ stopReason: 'natural' });
expect(onStepFinish).toHaveBeenCalledTimes(1);
expect(onStepFinish).toHaveBeenCalledWith({ stepIndex: 1, stepBudget: 40 });
});
it('logs and ignores onStepFinish callback errors', async () => {
const query = vi.fn((_input: any) =>
stream([

4
uv.lock generated
View file

@ -458,7 +458,7 @@ wheels = [
[[package]]
name = "ktx-daemon"
version = "0.5.0"
version = "0.6.0"
source = { editable = "python/ktx-daemon" }
dependencies = [
{ name = "fastapi" },
@ -515,7 +515,7 @@ dev = [
[[package]]
name = "ktx-sl"
version = "0.5.0"
version = "0.6.0"
source = { editable = "python/ktx-sl" }
dependencies = [
{ name = "pydantic" },