From 6837ab253d6173d3270a8fdc08cff066d1137d1b Mon Sep 17 00:00:00 2001
From: Andrey Avtomonov <andreybavt@gmail.com>
Date: Thu, 28 May 2026 02:09:53 +0200
Subject: [PATCH] fix(cli): align ingest step counter with SDK num_turns (#225)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../src/context/llm/claude-code-runtime.ts    | 18 +++++-
 .../context/llm/claude-code-runtime.test.ts   | 58 +++++++++++++++++++
 uv.lock                                       |  4 +-
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/packages/cli/src/context/llm/claude-code-runtime.ts b/packages/cli/src/context/llm/claude-code-runtime.ts
index 0eb3eadb..22055c28 100644
--- a/packages/cli/src/context/llm/claude-code-runtime.ts
+++ b/packages/cli/src/context/llm/claude-code-runtime.ts
@@ -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)) {
diff --git a/packages/cli/test/context/llm/claude-code-runtime.test.ts b/packages/cli/test/context/llm/claude-code-runtime.test.ts
index 205c74e0..706d5d55 100644
--- a/packages/cli/test/context/llm/claude-code-runtime.test.ts
+++ b/packages/cli/test/context/llm/claude-code-runtime.test.ts
@@ -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([
diff --git a/uv.lock b/uv.lock
index 7c2c368f..25a8fab6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -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" },