mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-07 07:55:13 +02:00
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:
parent
a94f35800a
commit
6837ab253d
3 changed files with 77 additions and 3 deletions
|
|
@ -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)) {
|
||||
|
|
|
|||
|
|
@ -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
4
uv.lock
generated
|
|
@ -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" },
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue