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';
|
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)) {
|
||||||
|
|
|
||||||
|
|
@ -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
4
uv.lock
generated
|
|
@ -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" },
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue