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

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 02:09:53 +02:00 committed by GitHub
parent a94f35800a
commit 6837ab253d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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'; 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 { function resultError(result: SDKResultMessage): Error | undefined {
if (result.subtype === 'success') { if (result.subtype === 'success') {
return undefined; return undefined;
@ -190,7 +206,7 @@ async function collectResult(params: {
let result: SDKResultMessage | undefined; let result: SDKResultMessage | undefined;
for await (const message of params.query({ prompt: params.prompt, options: params.options })) { for await (const message of params.query({ prompt: params.prompt, options: params.options })) {
assertInitIsolation(message, params.allowedToolIds, params.expectedMcpServerNames); assertInitIsolation(message, params.allowedToolIds, params.expectedMcpServerNames);
if (message.type === 'assistant' && message.parent_tool_use_id === null) { if (countsAsAssistantTurn(message)) {
await params.onAssistantTurn?.(); await params.onAssistantTurn?.();
} }
if (isResult(message)) { 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 () => { it('logs and ignores onStepFinish callback errors', async () => {
const query = vi.fn((_input: any) => const query = vi.fn((_input: any) =>
stream([ stream([

4
uv.lock generated
View file

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