From f9fb0ecdc66daeac265e81a196ede5edebf3f30e Mon Sep 17 00:00:00 2001 From: Rick Gao Date: Thu, 23 Jul 2026 14:52:54 -0700 Subject: [PATCH 1/4] feat(trace): trace ordinary tool calls --- src/callLifecycle.ts | 105 +++++++ src/hookHandler.ts | 49 ++++ src/session.ts | 57 +++- tests/daemon-idle-inflight.test.ts | 30 +- tests/tool-lifecycle.test.ts | 437 +++++++++++++++++++++++++++++ 5 files changed, 669 insertions(+), 9 deletions(-) create mode 100644 src/callLifecycle.ts create mode 100644 tests/tool-lifecycle.test.ts diff --git a/src/callLifecycle.ts b/src/callLifecycle.ts new file mode 100644 index 0000000..a68a9da --- /dev/null +++ b/src/callLifecycle.ts @@ -0,0 +1,105 @@ +// SPDX-FileCopyrightText: 2026 CoreWeave, Inc. +// SPDX-License-Identifier: MIT +// SPDX-PackageName: weave-claude-code + +import type { Tool } from 'weave'; +import { ATTR, jsonStr, toolDisplayName } from './genaiSpans.js'; +import type { TurnTrace } from './session.js'; + +export type ToolOutcome = + | { kind: 'success'; value: unknown } + | { kind: 'failure'; error: unknown }; + +export type ToolDescriptor = { + toolUseId: string; + name: string; + input: Record; +}; + +export type TracedTool = { + span: Tool; + parent: TurnTrace; + toolUseId: string; +}; + +/** Owns the exact hook identities for ordinary tools in one session. */ +export class ToolLifecycle { + private readonly openById = new Map(); + private readonly tombstones = new Set(); + + start(parent: TurnTrace, tool: ToolDescriptor): TracedTool | undefined { + if (this.openById.has(tool.toolUseId) || this.tombstones.has(tool.toolUseId)) { + return undefined; + } + + const span = parent.span.startTool({ + name: tool.name, + args: jsonStr(tool.input), + toolCallId: tool.toolUseId, + }); + span.setAttributes({ + [ATTR.WEAVE_DISPLAY_NAME]: toolDisplayName(tool.name, tool.input), + }); + + const traced = { span, parent, toolUseId: tool.toolUseId }; + parent.children.add(traced); + this.openById.set(tool.toolUseId, traced); + return traced; + } + + /** A terminal hook may be the first hook observed after a daemon restart. */ + finishOrRecover( + parent: () => TurnTrace, + tool: ToolDescriptor, + outcome: ToolOutcome, + ): boolean { + if (this.tombstones.has(tool.toolUseId)) return false; + const traced = this.openById.get(tool.toolUseId) ?? this.start(parent(), tool); + if (!traced) return false; + + if (outcome.kind === 'success') { + traced.span.result = jsonStr(outcome.value); + traced.span.end(); + } else { + const error = String(outcome.error); + traced.span.result = error; + traced.span.setAttributes({ [ATTR.ERROR_TYPE]: errorType(outcome.error) }); + traced.span.end({ error: new Error(error) }); + } + this.complete(tool.toolUseId, traced); + return true; + } + + /** End every unfinished child before its owning turn. */ + finalizeChildren(parent: TurnTrace, reason: string): string[] { + const closed: string[] = []; + for (const traced of [...parent.children].reverse()) { + traced.span.setAttributes({ [ATTR.WEAVE_ORPHAN_REASON]: reason }); + traced.span.end({ error: new Error(`call did not complete (${reason})`) }); + this.complete(traced.toolUseId, traced); + closed.push(traced.toolUseId); + } + return closed; + } + + hasOpenTools(): boolean { + return this.openById.size > 0; + } + + private complete(toolUseId: string, traced: TracedTool): void { + traced.parent.children.delete(traced); + this.openById.delete(toolUseId); + this.tombstones.add(toolUseId); + } +} + +function errorType(error: unknown): string { + if (typeof error === 'string') { + return error.trim().match(/^[A-Z][A-Za-z_]*Error/)?.[0] ?? 'tool_error'; + } + if (error && typeof error === 'object' && 'type' in error) { + const type = (error as Record)['type']; + if (typeof type === 'string' && type) return type; + } + return 'tool_error'; +} diff --git a/src/hookHandler.ts b/src/hookHandler.ts index 95309fd..18fede2 100644 --- a/src/hookHandler.ts +++ b/src/hookHandler.ts @@ -7,6 +7,7 @@ import * as path from 'path'; import type { HookInput, InstructionsLoadedHookInput, + PreToolUseHookInput, PreCompactHookInput, SessionEndHookInput, SessionStartHookInput, @@ -14,6 +15,7 @@ import type { UserPromptSubmitHookInput, } from '@anthropic-ai/claude-agent-sdk'; import * as weave from 'weave'; +import type { ToolDescriptor, ToolOutcome } from './callLifecycle.js'; import type { CompactionAttrs } from './genaiSpans.js'; import { snippet } from './genaiSpans.js'; import { Session } from './session.js'; @@ -21,6 +23,12 @@ import { TranscriptFile } from './transcriptFile.js'; type TraceLog = (level: 'DEBUG' | 'INFO' | 'ERROR', message: string) => void; +type HookInputFor = Extract< + HookInput, + { hook_event_name: Event } +>; +type PostToolResultHookInput = HookInputFor<'PostToolUse' | 'PostToolUseFailure'>; + function parseHookInput(payload: unknown): HookInput | undefined { if (!payload || typeof payload !== 'object') return undefined; const input = payload as Record; @@ -86,6 +94,13 @@ export class HookHandler { case 'UserPromptSubmit': await this.handleUserPromptSubmit(sessionId, input); return; + case 'PreToolUse': + await this.handlePreToolUse(sessionId, input); + return; + case 'PostToolUse': + case 'PostToolUseFailure': + await this.handlePostToolResult(sessionId, input); + return; case 'PreCompact': this.handlePreCompact(sessionId, input); return; @@ -254,6 +269,40 @@ export class HookHandler { } } + private async handlePreToolUse( + sessionId: string, + input: PreToolUseHookInput, + ): Promise { + if (input.tool_name === 'Agent' || input.agent_id) return; + const session = await this.getOrReconstructSession(sessionId, input); + if (!session) return; + session.startTool(input.prompt_id, this.toolDescriptor(input)); + } + + private async handlePostToolResult( + sessionId: string, + input: PostToolResultHookInput, + ): Promise { + if (input.tool_name === 'Agent' || input.agent_id) return; + const session = await this.getOrReconstructSession(sessionId, input); + if (!session) return; + + const outcome: ToolOutcome = input.hook_event_name === 'PostToolUse' + ? { kind: 'success', value: input.tool_response } + : { kind: 'failure', error: input.error }; + session.finishTool(input.prompt_id, this.toolDescriptor(input), outcome); + } + + private toolDescriptor( + input: Pick, + ): ToolDescriptor { + return { + toolUseId: input.tool_use_id, + name: input.tool_name, + input: input.tool_input as Record, + }; + } + private async handleStop(sessionId: string, input: StopHookInput): Promise { const session = await this.getOrReconstructSession(sessionId, input); if (!session) return; diff --git a/src/session.ts b/src/session.ts index 179a8d7..5eb281e 100644 --- a/src/session.ts +++ b/src/session.ts @@ -15,6 +15,12 @@ import { setCompactionAttrs, } from './genaiSpans.js'; import type { CompactionAttrs } from './genaiSpans.js'; +import { ToolLifecycle } from './callLifecycle.js'; +import type { + ToolDescriptor, + ToolOutcome, + TracedTool, +} from './callLifecycle.js'; import { assistantResponses, extractAssistantTextBlocks, @@ -27,12 +33,14 @@ import { TranscriptFile, readFirstTranscriptLine } from './transcriptFile.js'; type TraceLog = (level: 'DEBUG' | 'INFO' | 'ERROR', message: string) => void; -type TurnTrace = { +export type TurnTrace = { span: weave.Turn; promptId?: string; userText?: string; /** A Stop snapshot is quiescent but remains reopenable because hooks block. */ phase: 'active' | 'stopped'; + /** Calls are owned by their parent span; hook ids are only lookup indexes. */ + children: Set; /** Number of provider responses already present when this prompt began. */ responseOffset: number; /** Frozen when a newer prompt starts, preventing cross-turn replay. */ @@ -70,6 +78,7 @@ export class Session { /** File path → latest loaded contents, preserving first-load order. */ private readonly systemInstructions = new Map(); private readonly turns = new Set(); + private readonly tools = new ToolLifecycle(); private currentTurn?: TurnTrace; private pendingCompaction?: CompactionAttrs; @@ -116,6 +125,27 @@ export class Session { this.systemInstructions.set(filePath, content); } + startTool(promptId: string | undefined, tool: ToolDescriptor): boolean { + const turn = this.ensureTurn(promptId); + const traced = this.tools.start(turn, tool); + if (traced) turn.phase = 'active'; + return Boolean(traced); + } + + finishTool( + promptId: string | undefined, + tool: ToolDescriptor, + outcome: ToolOutcome, + ): boolean { + const finished = this.tools.finishOrRecover( + () => this.ensureTurn(promptId), + tool, + outcome, + ); + if (finished) this.finalizeIdleSupersededTurns(); + return finished; + } + submitPrompt( promptId: string | undefined, prompt: string, @@ -131,7 +161,9 @@ export class Session { parseSessionFd(this.transcript.getFd()) ?? { turns: [] }, ).length; responseOffsetFloor = previous.responseLimit; - this.finalizeTurn(previous, 'superseded_by_next_prompt'); + if (promptId === undefined || previous.children.size === 0) { + this.finalizeTurn(previous, 'superseded_by_next_prompt'); + } } const turn = this.startTurn({ @@ -188,13 +220,14 @@ export class Session { const turnCount = this.turns.size; for (const turn of [...this.turns]) { this.recordFinalTurnOutput(turn, orphanReason, parsed); - this.endTurn(turn); + this.closeTurn(turn, orphanReason); } return turnCount; } hasInFlightWork(): boolean { - return [...this.turns].some(turn => turn.phase === 'active'); + return this.tools.hasOpenTools() + || [...this.turns].some(turn => turn.phase === 'active'); } close(): void { @@ -246,6 +279,7 @@ export class Session { promptId: options.promptId, userText: cursor.userText, phase: 'active', + children: new Set(), responseOffset: cursor.responseOffset, seenResponses: new Set(), }; @@ -353,15 +387,26 @@ export class Session { private finalizeTurn(turn: TurnTrace, orphanReason: string): void { this.recordFinalTurnOutput(turn, orphanReason, this.parseTranscript()); - this.endTurn(turn); + this.closeTurn(turn, orphanReason); } - private endTurn(turn: TurnTrace): void { + private closeTurn(turn: TurnTrace, orphanReason: string): void { + for (const toolUseId of this.tools.finalizeChildren(turn, orphanReason)) { + this.log('DEBUG', `Closed pending tool: ${toolUseId}`); + } turn.span.end(); this.turns.delete(turn); if (this.currentTurn === turn) this.currentTurn = undefined; } + private finalizeIdleSupersededTurns(): void { + for (const turn of [...this.turns]) { + if (turn.responseLimit !== undefined && turn.children.size === 0) { + this.finalizeTurn(turn, 'superseded_by_next_prompt'); + } + } + } + private parseTranscript(): ParsedSession | null { try { return parseSessionFd(this.transcript.getFd()); diff --git a/tests/daemon-idle-inflight.test.ts b/tests/daemon-idle-inflight.test.ts index 7b15d63..1848159 100644 --- a/tests/daemon-idle-inflight.test.ts +++ b/tests/daemon-idle-inflight.test.ts @@ -2,9 +2,8 @@ // SPDX-License-Identifier: MIT // SPDX-PackageName: weave-claude-code -// Active root work pins the daemon across its idle window. A blockable Stop -// keeps that root reopenable but makes it quiescent; later call-state slices -// extend the same predicate for tools and subagents. +// Active root or tool work pins the daemon across its idle window. A blockable +// Stop keeps its root reopenable but makes it quiescent when no call is open. import { test } from 'node:test'; import assert from 'node:assert/strict'; @@ -68,3 +67,28 @@ test('daemon idles out once a stopped turn is quiescent', async () => { await d.stop(); } }); + +test('an open tool keeps a stopped turn alive', async () => { + const d = await startTestDaemon({ env: { WEAVE_INACTIVITY_MS: '1000' } }); + try { + const sessionId = 'inflight-tool'; + const transcript = writeTranscript(d.home, sessionId); + await d.send({ hook_event_name: 'SessionStart', session_id: sessionId, transcript_path: transcript }); + await d.send({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + transcript_path: transcript, prompt: 'a background tool', + }); + await d.send({ + hook_event_name: 'PreToolUse', session_id: sessionId, + transcript_path: transcript, tool_use_id: 'long-read', + tool_name: 'Read', tool_input: { file_path: '/tmp/slow' }, + }); + await d.send({ hook_event_name: 'Stop', session_id: sessionId, transcript_path: transcript }); + + const stayedUp = await d.waitForLog(/work in flight — staying up/, 3000); + assert.ok(stayedUp, `daemon should hold open for the tool; log was:\n${d.readLog()}`); + assert.equal(d.hasExited(), false); + } finally { + await d.stop(); + } +}); diff --git a/tests/tool-lifecycle.test.ts b/tests/tool-lifecycle.test.ts new file mode 100644 index 0000000..4640400 --- /dev/null +++ b/tests/tool-lifecycle.test.ts @@ -0,0 +1,437 @@ +// SPDX-FileCopyrightText: 2026 CoreWeave, Inc. +// SPDX-License-Identifier: MIT +// SPDX-PackageName: weave-claude-code + +import { test, type TestContext } from 'node:test'; +import assert from 'node:assert/strict'; +import * as fs from 'node:fs'; +import * as os from 'node:os'; +import * as path from 'node:path'; +import type { ReadableSpan } from '@opentelemetry/sdk-trace-base'; +import { ATTR } from '../src/genaiSpans.ts'; +import { + flushWeave, + initWeaveInMemory, + makeGenaiDaemon, + spanParentId, +} from './helpers.ts'; + +function makeTranscript(t: TestContext, sessionId: string, prompt: string) { + const dir = fs.mkdtempSync(path.join(os.homedir(), '.weave-tool-lifecycle-')); + const file = path.join(dir, `${sessionId}.jsonl`); + const append = (entry: Record) => { + fs.appendFileSync(file, JSON.stringify(entry) + '\n'); + }; + fs.writeFileSync(file, ''); + append({ type: 'user', message: { role: 'user', content: prompt } }); + t.after(() => fs.rmSync(dir, { recursive: true, force: true })); + return { + file, + appendPrompt(text: string) { + append({ type: 'user', message: { role: 'user', content: text } }); + }, + appendResponse(id: string, text: string) { + append({ + type: 'assistant', + message: { + role: 'assistant', + id, + model: 'claude-opus-4-8', + usage: { input_tokens: 10, output_tokens: 5 }, + content: [{ type: 'text', text }], + }, + }); + }, + }; +} + +function toolSpans(spans: ReadableSpan[]): ReadableSpan[] { + return spans.filter(span => span.attributes[ATTR.OPERATION_NAME] === 'execute_tool'); +} + +function turnSpans(spans: ReadableSpan[]): ReadableSpan[] { + return spans.filter(span => span.attributes[ATTR.OPERATION_NAME] === 'invoke_agent'); +} + +test('ordinary tool calls are traced once while Agent-owned calls remain deferred', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'ordinary-tool'; + const transcript = makeTranscript(t, sessionId, 'read it'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, prompt: 'read it', + }); + const tool = { + session_id: sessionId, + tool_use_id: 'read-1', + tool_name: 'Read', + tool_input: { file_path: '/tmp/input.txt' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...tool, tool_response: 'contents' }); + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...tool, tool_response: 'duplicate' }); + + const agent = { + session_id: sessionId, + tool_use_id: 'agent-1', + tool_name: 'Agent', + tool_input: { subagent_type: 'Explore', prompt: 'inspect it' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...agent }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...agent, tool_response: 'done' }); + const child = { + session_id: sessionId, + agent_id: 'untraced-agent', + tool_use_id: 'child-read', + tool_name: 'Read', + tool_input: { file_path: '/tmp/child.txt' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...child }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...child, tool_response: 'child' }); + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const tools = toolSpans(spans); + const [turn] = turnSpans(spans); + assert.equal(tools.length, 1); + assert.ok(turn); + assert.equal(tools[0].attributes['gen_ai.tool.call.id'], 'read-1'); + assert.equal(tools[0].attributes['gen_ai.tool.call.result'], 'contents'); + assert.equal(tools[0].attributes[ATTR.WEAVE_DISPLAY_NAME], 'Read: /tmp/input.txt'); + assert.equal(spanParentId(tools[0]), turn.spanContext().spanId); +}); + +test('PostToolUseFailure records the tool result and error type', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'failed-tool'; + const transcript = makeTranscript(t, sessionId, 'run it'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ hook_event_name: 'UserPromptSubmit', session_id: sessionId, prompt: 'run it' }); + const tool = { + session_id: sessionId, + tool_use_id: 'bash-1', + tool_name: 'Bash', + tool_input: { command: 'exit 1' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ + hook_event_name: 'PostToolUseFailure', ...tool, error: 'CommandError: exit 1', + }); + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const [span] = toolSpans(exporter.getFinishedSpans()); + assert.ok(span); + assert.equal(span.attributes['gen_ai.tool.call.result'], 'CommandError: exit 1'); + assert.equal(span.attributes[ATTR.ERROR_TYPE], 'CommandError'); + assert.equal(span.status.code, 2); +}); + +test('a restart-first terminal hook recovers one exact tool and turn', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'terminal-tool-restart'; + const transcript = makeTranscript(t, sessionId, 'read after restart'); + const daemon = makeGenaiDaemon(); + const tool = { + session_id: sessionId, + transcript_path: transcript.file, + cwd: '/x', + tool_use_id: 'recovered-read', + tool_name: 'Read', + tool_input: { file_path: '/tmp/recovered.txt' }, + }; + + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...tool, tool_response: 'contents' }); + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...tool, tool_response: 'duplicate' }); + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const tools = toolSpans(spans); + const turns = turnSpans(spans); + assert.equal(tools.length, 1); + assert.equal(turns.length, 1); + assert.equal(spanParentId(tools[0]), turns[0].spanContext().spanId); + assert.equal( + turns[0].attributes[ATTR.INPUT_MESSAGES], + JSON.stringify([{ role: 'user', parts: [{ type: 'text', content: 'read after restart' }] }]), + ); +}); + +test('restart-first tool results preserve SessionEnd prompt identity', async (t) => { + const scenarios: Array<{ + name: string; + eventPrompt?: string; + endPrompt?: string; + sameRoot: boolean; + }> = [ + { name: 'same explicit prompt', eventPrompt: 'prompt-a', endPrompt: 'prompt-a', sameRoot: true }, + { name: 'different explicit prompts', eventPrompt: 'prompt-a', endPrompt: 'prompt-b', sameRoot: false }, + { name: 'legacy result then explicit end', endPrompt: 'prompt-b', sameRoot: true }, + { name: 'explicit result then legacy end', eventPrompt: 'prompt-a', sameRoot: false }, + ]; + + for (const scenario of scenarios) { + await t.test(scenario.name, async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = `restart-prompt-${scenario.name.replaceAll(' ', '-')}`; + const transcript = makeTranscript(t, sessionId, 'older'); + transcript.appendResponse('older-response', 'old'); + transcript.appendPrompt('final'); + transcript.appendResponse('final-response', 'finished'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'PostToolUse', session_id: sessionId, + prompt_id: scenario.eventPrompt, transcript_path: transcript.file, + tool_use_id: 'restart-tool', tool_name: 'Read', + tool_input: { file_path: '/tmp/restart' }, tool_response: 'contents', + }); + await daemon.routeEvent({ + hook_event_name: 'SessionEnd', session_id: sessionId, + prompt_id: scenario.endPrompt, transcript_path: transcript.file, reason: 'clear', + }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const tool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === 'restart-tool'); + const chat = spans.find(span => span.attributes[ATTR.RESPONSE_ID] === 'final-response'); + assert.ok(tool && chat); + const turns = turnSpans(spans); + assert.equal(turns.length, scenario.sameRoot ? 1 : 2); + const finalTurn = turns.find(turn => turn.spanContext().spanId === spanParentId(chat)); + assert.ok(String(finalTurn?.attributes[ATTR.INPUT_MESSAGES]).includes('final')); + assert.equal( + spanParentId(tool) === spanParentId(chat), + scenario.sameRoot, + 'prompt identity determines whether recovery reuses the same root', + ); + }); + } +}); + +test('SessionEnd orphans unfinished tools before closing their turn', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'open-tool-session-end'; + const transcript = makeTranscript(t, sessionId, 'keep reading'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, prompt: 'keep reading', + }); + await daemon.routeEvent({ + hook_event_name: 'PreToolUse', session_id: sessionId, + tool_use_id: 'open-read', tool_name: 'Read', tool_input: { file_path: '/tmp/open.txt' }, + }); + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const [tool] = toolSpans(spans); + const [turn] = turnSpans(spans); + assert.ok(tool && turn); + assert.equal(tool.attributes[ATTR.WEAVE_ORPHAN_REASON], 'session_ended'); + assert.equal(tool.status.code, 2); + assert.equal(spanParentId(tool), turn.spanContext().spanId); + assert.ok(spans.indexOf(tool) < spans.indexOf(turn), 'child exports before its parent'); +}); + +test('prompt_id keeps background tools attached to their original turns', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'tool-prompt-ownership'; + const transcript = makeTranscript(t, sessionId, 'first'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + prompt_id: 'prompt-1', prompt: 'first', + }); + const first = { + session_id: sessionId, prompt_id: 'prompt-1', + tool_use_id: 'read-first', tool_name: 'Read', tool_input: { file_path: '/tmp/first' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...first }); + + transcript.appendPrompt('second'); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + prompt_id: 'prompt-2', prompt: 'second', + }); + const second = { + session_id: sessionId, prompt_id: 'prompt-2', + tool_use_id: 'read-second', tool_name: 'Read', tool_input: { file_path: '/tmp/second' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...second }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...second, tool_response: 'second' }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...first, tool_response: 'first' }); + await flushWeave(); + assert.equal(turnSpans(exporter.getFinishedSpans()).length, 1); + + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const turns = turnSpans(spans); + const tools = toolSpans(spans); + assert.equal(turns.length, 2); + assert.equal(tools.length, 2); + const firstTurn = turns.find(span => String(span.attributes[ATTR.INPUT_MESSAGES]).includes('first')); + const secondTurn = turns.find(span => String(span.attributes[ATTR.INPUT_MESSAGES]).includes('second')); + const firstTool = tools.find(span => span.attributes['gen_ai.tool.call.id'] === 'read-first'); + const secondTool = tools.find(span => span.attributes['gen_ai.tool.call.id'] === 'read-second'); + assert.ok(firstTurn && secondTurn && firstTool && secondTool); + assert.equal(spanParentId(firstTool), firstTurn.spanContext().spanId); + assert.equal(spanParentId(secondTool), secondTurn.spanContext().spanId); +}); + +test('Stop(prompt_id) snapshots only its turn and later tools keep their owners', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'stop-selected-prompt'; + const transcript = makeTranscript(t, sessionId, 'older prompt'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + prompt_id: 'prompt-1', prompt: 'older prompt', + }); + await daemon.routeEvent({ + hook_event_name: 'PreToolUse', session_id: sessionId, prompt_id: 'prompt-1', + tool_use_id: 'hold-older', tool_name: 'Read', tool_input: { file_path: '/tmp/hold' }, + }); + + transcript.appendResponse('older-response', 'older answer'); + transcript.appendPrompt('newer prompt'); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + prompt_id: 'prompt-2', prompt: 'newer prompt', + }); + transcript.appendResponse('newer-response', 'newer answer'); + await daemon.routeEvent({ + hook_event_name: 'Stop', session_id: sessionId, prompt_id: 'prompt-1', + }); + await flushWeave(); + + const afterStop = exporter.getFinishedSpans(); + assert.deepEqual(chatsById(afterStop), ['older-response']); + assert.equal(turnSpans(afterStop).length, 0, 'blockable Stop retains both roots'); + + for (const [promptId, toolUseId] of [ + ['prompt-1', 'continued-older'], + ['prompt-2', 'newer-tool'], + ] as const) { + const tool = { + session_id: sessionId, prompt_id: promptId, tool_use_id: toolUseId, + tool_name: 'Read', tool_input: { file_path: `/tmp/${toolUseId}` }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...tool }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...tool, tool_response: 'done' }); + } + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const olderChat = spans.find(span => span.attributes[ATTR.RESPONSE_ID] === 'older-response'); + const newerChat = spans.find(span => span.attributes[ATTR.RESPONSE_ID] === 'newer-response'); + const olderTool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === 'continued-older'); + const newerTool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === 'newer-tool'); + assert.ok(olderChat && newerChat && olderTool && newerTool); + assert.equal(spanParentId(olderTool), spanParentId(olderChat)); + assert.equal(spanParentId(newerTool), spanParentId(newerChat)); + assert.notEqual(spanParentId(olderChat), spanParentId(newerChat)); +}); + +test('a legacy next prompt orphans its open tool and starts a clean turn', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'legacy-tool-prompt-boundary'; + const transcript = makeTranscript(t, sessionId, 'first'); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'SessionStart', session_id: sessionId, + transcript_path: transcript.file, source: 'startup', cwd: '/x', + }); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, prompt: 'first', + }); + const interrupted = { + session_id: sessionId, tool_use_id: 'interrupted-tool', + tool_name: 'Bash', tool_input: { command: 'sleep 999' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...interrupted }); + + transcript.appendPrompt('second'); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, prompt: 'second', + }); + await daemon.routeEvent({ + hook_event_name: 'PostToolUse', ...interrupted, tool_response: 'too late', + }); + const next = { + session_id: sessionId, tool_use_id: 'next-tool', + tool_name: 'Read', tool_input: { file_path: '/tmp/next' }, + }; + await daemon.routeEvent({ hook_event_name: 'PreToolUse', ...next }); + await daemon.routeEvent({ hook_event_name: 'PostToolUse', ...next, tool_response: 'next result' }); + await daemon.routeEvent({ hook_event_name: 'SessionEnd', session_id: sessionId, reason: 'clear' }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const firstTurn = turnSpans(spans).find(span => + String(span.attributes[ATTR.INPUT_MESSAGES]).includes('first')); + const secondTurn = turnSpans(spans).find(span => + String(span.attributes[ATTR.INPUT_MESSAGES]).includes('second')); + const oldTool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === 'interrupted-tool'); + const nextTool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === 'next-tool'); + assert.ok(firstTurn && secondTurn && oldTool && nextTool); + assert.equal(oldTool.attributes[ATTR.WEAVE_ORPHAN_REASON], 'superseded_by_next_prompt'); + assert.equal(oldTool.attributes['gen_ai.tool.call.result'], undefined); + assert.equal(spanParentId(oldTool), firstTurn.spanContext().spanId); + assert.equal(nextTool.attributes['gen_ai.tool.call.result'], 'next result'); + assert.equal(nextTool.attributes[ATTR.WEAVE_ORPHAN_REASON], undefined); + assert.equal(spanParentId(nextTool), secondTurn.spanContext().spanId); +}); + +function chatsById(spans: ReadableSpan[]): unknown[] { + return spans + .filter(span => span.attributes[ATTR.OPERATION_NAME] === 'chat') + .map(span => span.attributes[ATTR.RESPONSE_ID]); +} From 334354cbacd40dfae7fb83d3898f1a450c357b7a Mon Sep 17 00:00:00 2001 From: Rick Gao Date: Thu, 6 Aug 2026 14:52:30 -0700 Subject: [PATCH 2/4] fix(trace): recover tool turns from transcripts --- src/session.ts | 62 ++++++++++++++++++++++++++++--- tests/tool-lifecycle.test.ts | 72 ++++++++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 6 deletions(-) diff --git a/src/session.ts b/src/session.ts index 5eb281e..f2aa5dd 100644 --- a/src/session.ts +++ b/src/session.ts @@ -24,6 +24,7 @@ import type { import { assistantResponses, extractAssistantTextBlocks, + isToolUseBlock, lastAssistantTextEndsWith, parseSessionFd, } from './parser.js'; @@ -43,7 +44,7 @@ export type TurnTrace = { children: Set; /** Number of provider responses already present when this prompt began. */ responseOffset: number; - /** Frozen when a newer prompt starts, preventing cross-turn replay. */ + /** Upper transcript boundary once this turn is known to be complete. */ responseLimit?: number; /** Supports repeated/blockable Stop hooks without duplicate chat spans. */ seenResponses: Set; @@ -67,6 +68,13 @@ type StartTurnOptions = { makeCurrent?: boolean; }; +type TurnCursor = { + responseOffset: number; + responseLimit?: number; + startTime?: Date; + userText?: string; +}; + export class Session { readonly sessionId: string; readonly conversationId: string; @@ -126,7 +134,7 @@ export class Session { } startTool(promptId: string | undefined, tool: ToolDescriptor): boolean { - const turn = this.ensureTurn(promptId); + const turn = this.ensureToolTurn(promptId, tool.toolUseId); const traced = this.tools.start(turn, tool); if (traced) turn.phase = 'active'; return Boolean(traced); @@ -138,7 +146,7 @@ export class Session { outcome: ToolOutcome, ): boolean { const finished = this.tools.finishOrRecover( - () => this.ensureTurn(promptId), + () => this.ensureToolTurn(promptId, tool.toolUseId), tool, outcome, ); @@ -242,7 +250,7 @@ export class Session { private transcriptCursor( options: StartTurnOptions, - ): { responseOffset: number; startTime?: Date; userText?: string } { + ): TurnCursor { const parsed = parseSessionFd(this.transcript.getFd()); if (!parsed) return { responseOffset: 0, userText: options.userMessage }; @@ -261,8 +269,36 @@ export class Session { }; } - private startTurn(options: StartTurnOptions = {}): TurnTrace { - const cursor = this.transcriptCursor(options); + private toolTurnCursor(toolUseId: string): TurnCursor | undefined { + const parsed = parseSessionFd(this.transcript.getFd()); + if (!parsed) return undefined; + + let responseOffset = 0; + let match: TurnCursor | undefined; + for (const [index, turn] of parsed.turns.entries()) { + const containsTool = turn.responses.some(response => + response.content.some(block => isToolUseBlock(block) && block.id === toolUseId)); + if (containsTool) { + // An exact tool_use_id should belong to one transcript turn. Preserve + // the cautious fallback if a malformed transcript makes it ambiguous. + if (match) return undefined; + const responseLimit = responseOffset + turn.responses.length; + match = { + responseOffset, + ...(index < parsed.turns.length - 1 ? { responseLimit } : {}), + startTime: parseTimestamp(turn.startTime), + userText: turn.userText, + }; + } + responseOffset += turn.responses.length; + } + return match; + } + + private startTurn( + options: StartTurnOptions = {}, + cursor: TurnCursor = this.transcriptCursor(options), + ): TurnTrace { const span = this.conversation.startTurn({ agentVersion: VERSION, model: this.initialRequestModel, @@ -283,6 +319,7 @@ export class Session { responseOffset: cursor.responseOffset, seenResponses: new Set(), }; + if (cursor.responseLimit !== undefined) turn.responseLimit = cursor.responseLimit; this.turns.add(turn); if (options.makeCurrent !== false) this.currentTurn = turn; return turn; @@ -298,6 +335,19 @@ export class Session { }); } + private ensureToolTurn(promptId: string | undefined, toolUseId: string): TurnTrace { + const existing = this.turnForPrompt(promptId); + if (existing) return existing; + + const cursor = this.toolTurnCursor(toolUseId); + if (!cursor) return this.ensureTurn(promptId); + return this.startTurn({ + promptId, + makeCurrent: cursor.responseLimit === undefined + && (!this.currentTurn || this.currentTurn.promptId === promptId), + }, cursor); + } + private reconcileFinalTurn( promptId: string | undefined, parsed: ParsedSession | null, diff --git a/tests/tool-lifecycle.test.ts b/tests/tool-lifecycle.test.ts index 4640400..2205d50 100644 --- a/tests/tool-lifecycle.test.ts +++ b/tests/tool-lifecycle.test.ts @@ -42,6 +42,23 @@ function makeTranscript(t: TestContext, sessionId: string, prompt: string) { }, }); }, + appendToolUseResponse( + id: string, + toolUseId: string, + name: string, + input: Record, + ) { + append({ + type: 'assistant', + message: { + role: 'assistant', + id, + model: 'claude-opus-4-8', + usage: { input_tokens: 10, output_tokens: 5 }, + content: [{ type: 'tool_use', id: toolUseId, name, input }], + }, + }); + }, }; } @@ -175,6 +192,61 @@ test('a restart-first terminal hook recovers one exact tool and turn', async (t) ); }); +test('a restart-first tool result recovers its prompt and triggering response', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'terminal-tool-response-restart'; + const transcript = makeTranscript(t, sessionId, 'read after restart'); + const toolUseId = 'recovered-read-response'; + const toolInput = { file_path: '/tmp/recovered-response.txt' }; + transcript.appendToolUseResponse('trigger-response', toolUseId, 'Read', toolInput); + const daemon = makeGenaiDaemon(); + + await daemon.routeEvent({ + hook_event_name: 'PostToolUse', session_id: sessionId, + prompt_id: 'prompt-a', transcript_path: transcript.file, cwd: '/x', + tool_use_id: toolUseId, tool_name: 'Read', tool_input: toolInput, + tool_response: 'contents', + }); + + // A later prompt must not finalize an empty recovered root before SessionEnd + // gets a chance to reconcile it. + transcript.appendPrompt('next prompt'); + await daemon.routeEvent({ + hook_event_name: 'UserPromptSubmit', session_id: sessionId, + prompt_id: 'prompt-b', prompt: 'next prompt', + }); + transcript.appendResponse('next-response', 'next answer'); + await daemon.routeEvent({ + hook_event_name: 'SessionEnd', session_id: sessionId, + prompt_id: 'prompt-b', reason: 'clear', + }); + await flushWeave(); + + const spans = exporter.getFinishedSpans(); + const tool = toolSpans(spans).find(span => + span.attributes['gen_ai.tool.call.id'] === toolUseId); + const triggers = spans.filter(span => span.attributes[ATTR.RESPONSE_ID] === 'trigger-response'); + assert.equal(triggers.length, 1); + const [trigger] = triggers; + assert.ok(tool && trigger); + assert.equal(spanParentId(tool), spanParentId(trigger)); + + const next = spans.find(span => span.attributes[ATTR.RESPONSE_ID] === 'next-response'); + assert.ok(next); + assert.notEqual(spanParentId(trigger), spanParentId(next)); + + const parent = turnSpans(spans).find(turn => + turn.spanContext().spanId === spanParentId(tool)); + assert.equal( + parent?.attributes[ATTR.INPUT_MESSAGES], + JSON.stringify([{ + role: 'user', + parts: [{ type: 'text', content: 'read after restart' }], + }]), + ); +}); + test('restart-first tool results preserve SessionEnd prompt identity', async (t) => { const scenarios: Array<{ name: string; From 5472ced8b34c4fdd8e0ca5ba10d440b14fd59aaa Mon Sep 17 00:00:00 2001 From: Rick Gao Date: Thu, 6 Aug 2026 15:01:06 -0700 Subject: [PATCH 3/4] refactor(trace): drop custom display names --- src/callLifecycle.ts | 6 +----- src/genaiSpans.ts | 28 +--------------------------- tests/tool-lifecycle.test.ts | 1 - 3 files changed, 2 insertions(+), 33 deletions(-) diff --git a/src/callLifecycle.ts b/src/callLifecycle.ts index a68a9da..cef95d0 100644 --- a/src/callLifecycle.ts +++ b/src/callLifecycle.ts @@ -3,7 +3,7 @@ // SPDX-PackageName: weave-claude-code import type { Tool } from 'weave'; -import { ATTR, jsonStr, toolDisplayName } from './genaiSpans.js'; +import { ATTR, jsonStr } from './genaiSpans.js'; import type { TurnTrace } from './session.js'; export type ToolOutcome = @@ -37,10 +37,6 @@ export class ToolLifecycle { args: jsonStr(tool.input), toolCallId: tool.toolUseId, }); - span.setAttributes({ - [ATTR.WEAVE_DISPLAY_NAME]: toolDisplayName(tool.name, tool.input), - }); - const traced = { span, parent, toolUseId: tool.toolUseId }; parent.children.add(traced); this.openById.set(tool.toolUseId, traced); diff --git a/src/genaiSpans.ts b/src/genaiSpans.ts index bef4259..179a102 100644 --- a/src/genaiSpans.ts +++ b/src/genaiSpans.ts @@ -51,7 +51,6 @@ export const ATTR = { WEAVE_SOURCE: 'weave.claude_code.source', WEAVE_PLUGIN_VERSION: 'weave.claude_code.plugin.version', WEAVE_ORPHAN_REASON: 'weave.claude_code.orphan_reason', - WEAVE_DISPLAY_NAME: 'weave.claude_code.display_name', // Integration identity: unlike gen_ai.agent.name, not user-overridable and // never changes per subagent. Set on the conversation; propagated to every span. @@ -225,7 +224,7 @@ export function setCompactionAttrs(turn: Turn, attrs: CompactionAttrs): void { } // ───────────────────────────────────────────────────────────────────────────── -// Display-name helpers +// Formatting helpers // ───────────────────────────────────────────────────────────────────────────── /** Single-line preview of a value: whitespace collapsed, truncated with `…`. */ @@ -233,28 +232,3 @@ export function snippet(value: unknown, maxLen = 60): string { const s = String(value ?? '').replace(/\s+/g, ' ').trim(); return s.length <= maxLen ? s : s.slice(0, maxLen - 1) + '…'; } - -export function toolDisplayName(toolName: string, input: Record): string { - switch (toolName) { - case 'Read': - case 'Edit': - case 'Write': - return `${toolName}: ${snippet(input['file_path'])}`; - case 'Glob': - return `Glob: ${snippet(input['pattern'])}`; - case 'Grep': - return `Grep: ${snippet(input['pattern'])}`; - case 'Bash': - return `Bash: ${snippet(input['command'])}`; - case 'Agent': - return `Agent: ${snippet(input['description'] ?? input['subagent_type'])}`; - case 'WebFetch': - return `WebFetch: ${snippet(input['url'])}`; - case 'WebSearch': - return `WebSearch: ${snippet(input['query'])}`; - default: { - const first = Object.values(input).find((v): v is string => typeof v === 'string'); - return first ? `${toolName}: ${snippet(first)}` : toolName; - } - } -} diff --git a/tests/tool-lifecycle.test.ts b/tests/tool-lifecycle.test.ts index 2205d50..38b4206 100644 --- a/tests/tool-lifecycle.test.ts +++ b/tests/tool-lifecycle.test.ts @@ -123,7 +123,6 @@ test('ordinary tool calls are traced once while Agent-owned calls remain deferre assert.ok(turn); assert.equal(tools[0].attributes['gen_ai.tool.call.id'], 'read-1'); assert.equal(tools[0].attributes['gen_ai.tool.call.result'], 'contents'); - assert.equal(tools[0].attributes[ATTR.WEAVE_DISPLAY_NAME], 'Read: /tmp/input.txt'); assert.equal(spanParentId(tools[0]), turn.spanContext().spanId); }); From 2eb8521980c37e1c02fed9efc219d17c6a28a1c9 Mon Sep 17 00:00:00 2001 From: Rick Gao Date: Thu, 6 Aug 2026 15:20:47 -0700 Subject: [PATCH 4/4] refactor(trace): validate tool lifecycle data --- src/callLifecycle.ts | 118 ++++++++++++++++++++--------------- src/hookHandler.ts | 75 +++++++++++++++++++--- src/session.ts | 32 +++++----- tests/tool-lifecycle.test.ts | 34 ++++++++++ 4 files changed, 180 insertions(+), 79 deletions(-) diff --git a/src/callLifecycle.ts b/src/callLifecycle.ts index cef95d0..9fe5f06 100644 --- a/src/callLifecycle.ts +++ b/src/callLifecycle.ts @@ -6,96 +6,112 @@ import type { Tool } from 'weave'; import { ATTR, jsonStr } from './genaiSpans.js'; import type { TurnTrace } from './session.js'; -export type ToolOutcome = - | { kind: 'success'; value: unknown } - | { kind: 'failure'; error: unknown }; +export type JsonValue = + | null + | boolean + | number + | string + | JsonValue[] + | { [key: string]: JsonValue }; -export type ToolDescriptor = { +export type JsonObject = Record; + +export type ToolCall = { toolUseId: string; name: string; - input: Record; + input: JsonObject; }; -export type TracedTool = { +export type ToolResult = + | { ok: true; output: JsonValue } + | { ok: false; error: string }; + +type OpenTool = { span: Tool; - parent: TurnTrace; + turn: TurnTrace; toolUseId: string; }; /** Owns the exact hook identities for ordinary tools in one session. */ export class ToolLifecycle { - private readonly openById = new Map(); + private readonly openById = new Map(); + private readonly openIdsByTurn = new Map>(); private readonly tombstones = new Set(); - start(parent: TurnTrace, tool: ToolDescriptor): TracedTool | undefined { - if (this.openById.has(tool.toolUseId) || this.tombstones.has(tool.toolUseId)) { + start(turn: TurnTrace, call: ToolCall): boolean { + return Boolean(this.open(turn, call)); + } + + private open(turn: TurnTrace, call: ToolCall): OpenTool | undefined { + if (this.openById.has(call.toolUseId) || this.tombstones.has(call.toolUseId)) { return undefined; } - const span = parent.span.startTool({ - name: tool.name, - args: jsonStr(tool.input), - toolCallId: tool.toolUseId, + const span = turn.span.startTool({ + name: call.name, + args: jsonStr(call.input), + toolCallId: call.toolUseId, }); - const traced = { span, parent, toolUseId: tool.toolUseId }; - parent.children.add(traced); - this.openById.set(tool.toolUseId, traced); - return traced; + const open = { span, turn, toolUseId: call.toolUseId }; + const openIds = this.openIdsByTurn.get(turn) ?? new Set(); + openIds.add(call.toolUseId); + this.openIdsByTurn.set(turn, openIds); + this.openById.set(call.toolUseId, open); + return open; } /** A terminal hook may be the first hook observed after a daemon restart. */ finishOrRecover( - parent: () => TurnTrace, - tool: ToolDescriptor, - outcome: ToolOutcome, + turn: () => TurnTrace, + call: ToolCall, + result: ToolResult, ): boolean { - if (this.tombstones.has(tool.toolUseId)) return false; - const traced = this.openById.get(tool.toolUseId) ?? this.start(parent(), tool); - if (!traced) return false; + if (this.tombstones.has(call.toolUseId)) return false; + const open = this.openById.get(call.toolUseId) ?? this.open(turn(), call); + if (!open) return false; - if (outcome.kind === 'success') { - traced.span.result = jsonStr(outcome.value); - traced.span.end(); + if (result.ok) { + open.span.result = jsonStr(result.output); + open.span.end(); } else { - const error = String(outcome.error); - traced.span.result = error; - traced.span.setAttributes({ [ATTR.ERROR_TYPE]: errorType(outcome.error) }); - traced.span.end({ error: new Error(error) }); + const error = result.error; + open.span.result = error; + open.span.setAttributes({ [ATTR.ERROR_TYPE]: errorType(error) }); + open.span.end({ error: new Error(error) }); } - this.complete(tool.toolUseId, traced); + this.complete(call.toolUseId, open); return true; } /** End every unfinished child before its owning turn. */ - finalizeChildren(parent: TurnTrace, reason: string): string[] { + finalizeChildren(turn: TurnTrace, reason: string): string[] { const closed: string[] = []; - for (const traced of [...parent.children].reverse()) { - traced.span.setAttributes({ [ATTR.WEAVE_ORPHAN_REASON]: reason }); - traced.span.end({ error: new Error(`call did not complete (${reason})`) }); - this.complete(traced.toolUseId, traced); - closed.push(traced.toolUseId); + for (const toolUseId of [...(this.openIdsByTurn.get(turn) ?? [])].reverse()) { + const open = this.openById.get(toolUseId); + if (!open) continue; + open.span.setAttributes({ [ATTR.WEAVE_ORPHAN_REASON]: reason }); + open.span.end({ error: new Error(`call did not complete (${reason})`) }); + this.complete(toolUseId, open); + closed.push(toolUseId); } return closed; } - hasOpenTools(): boolean { - return this.openById.size > 0; + hasOpenTools(turn?: TurnTrace): boolean { + return turn + ? Boolean(this.openIdsByTurn.get(turn)?.size) + : this.openById.size > 0; } - private complete(toolUseId: string, traced: TracedTool): void { - traced.parent.children.delete(traced); + private complete(toolUseId: string, open: OpenTool): void { + const openIds = this.openIdsByTurn.get(open.turn); + openIds?.delete(toolUseId); + if (openIds?.size === 0) this.openIdsByTurn.delete(open.turn); this.openById.delete(toolUseId); this.tombstones.add(toolUseId); } } -function errorType(error: unknown): string { - if (typeof error === 'string') { - return error.trim().match(/^[A-Z][A-Za-z_]*Error/)?.[0] ?? 'tool_error'; - } - if (error && typeof error === 'object' && 'type' in error) { - const type = (error as Record)['type']; - if (typeof type === 'string' && type) return type; - } - return 'tool_error'; +function errorType(error: string): string { + return error.trim().match(/^[A-Z][A-Za-z_]*Error/)?.[0] ?? 'tool_error'; } diff --git a/src/hookHandler.ts b/src/hookHandler.ts index 18fede2..b875b2d 100644 --- a/src/hookHandler.ts +++ b/src/hookHandler.ts @@ -15,7 +15,12 @@ import type { UserPromptSubmitHookInput, } from '@anthropic-ai/claude-agent-sdk'; import * as weave from 'weave'; -import type { ToolDescriptor, ToolOutcome } from './callLifecycle.js'; +import type { + JsonObject, + JsonValue, + ToolCall, + ToolResult, +} from './callLifecycle.js'; import type { CompactionAttrs } from './genaiSpans.js'; import { snippet } from './genaiSpans.js'; import { Session } from './session.js'; @@ -29,6 +34,32 @@ type HookInputFor = Extract< >; type PostToolResultHookInput = HookInputFor<'PostToolUse' | 'PostToolUseFailure'>; +function isJsonValue(value: unknown): value is JsonValue { + if (value === null) return true; + switch (typeof value) { + case 'boolean': + case 'string': + return true; + case 'number': + return Number.isFinite(value); + case 'object': + return Array.isArray(value) + ? value.every(isJsonValue) + : isJsonObject(value); + default: + return false; + } +} + +function isJsonObject(value: unknown): value is JsonObject { + return value !== null + && typeof value === 'object' + && !Array.isArray(value) + && (Object.getPrototypeOf(value) === Object.prototype + || Object.getPrototypeOf(value) === null) + && Object.values(value).every(isJsonValue); +} + function parseHookInput(payload: unknown): HookInput | undefined { if (!payload || typeof payload !== 'object') return undefined; const input = payload as Record; @@ -274,9 +305,11 @@ export class HookHandler { input: PreToolUseHookInput, ): Promise { if (input.tool_name === 'Agent' || input.agent_id) return; + const call = this.toolCall(input); + if (!call) return; const session = await this.getOrReconstructSession(sessionId, input); if (!session) return; - session.startTool(input.prompt_id, this.toolDescriptor(input)); + session.startTool(input.prompt_id, call); } private async handlePostToolResult( @@ -284,25 +317,47 @@ export class HookHandler { input: PostToolResultHookInput, ): Promise { if (input.tool_name === 'Agent' || input.agent_id) return; + const call = this.toolCall(input); + const result = this.toolResult(input); + if (!call || !result) return; const session = await this.getOrReconstructSession(sessionId, input); if (!session) return; - - const outcome: ToolOutcome = input.hook_event_name === 'PostToolUse' - ? { kind: 'success', value: input.tool_response } - : { kind: 'failure', error: input.error }; - session.finishTool(input.prompt_id, this.toolDescriptor(input), outcome); + session.finishTool(input.prompt_id, call, result); } - private toolDescriptor( + private toolCall( input: Pick, - ): ToolDescriptor { + ): ToolCall | undefined { + if ( + typeof input.tool_use_id !== 'string' + || typeof input.tool_name !== 'string' + || !isJsonObject(input.tool_input) + ) { + this.log('ERROR', 'Invalid tool hook payload'); + return undefined; + } return { toolUseId: input.tool_use_id, name: input.tool_name, - input: input.tool_input as Record, + input: input.tool_input, }; } + private toolResult(input: PostToolResultHookInput): ToolResult | undefined { + if (input.hook_event_name === 'PostToolUse') { + if (!isJsonValue(input.tool_response)) { + this.log('ERROR', 'Invalid tool result payload'); + return undefined; + } + return { ok: true, output: input.tool_response }; + } + if (typeof input.error !== 'string') { + this.log('ERROR', 'Invalid tool result payload'); + return undefined; + } + return { ok: false, error: input.error }; + } + private async handleStop(sessionId: string, input: StopHookInput): Promise { const session = await this.getOrReconstructSession(sessionId, input); if (!session) return; diff --git a/src/session.ts b/src/session.ts index f2aa5dd..c6e4fcd 100644 --- a/src/session.ts +++ b/src/session.ts @@ -17,9 +17,8 @@ import { import type { CompactionAttrs } from './genaiSpans.js'; import { ToolLifecycle } from './callLifecycle.js'; import type { - ToolDescriptor, - ToolOutcome, - TracedTool, + ToolCall, + ToolResult, } from './callLifecycle.js'; import { assistantResponses, @@ -40,8 +39,6 @@ export type TurnTrace = { userText?: string; /** A Stop snapshot is quiescent but remains reopenable because hooks block. */ phase: 'active' | 'stopped'; - /** Calls are owned by their parent span; hook ids are only lookup indexes. */ - children: Set; /** Number of provider responses already present when this prompt began. */ responseOffset: number; /** Upper transcript boundary once this turn is known to be complete. */ @@ -133,22 +130,22 @@ export class Session { this.systemInstructions.set(filePath, content); } - startTool(promptId: string | undefined, tool: ToolDescriptor): boolean { - const turn = this.ensureToolTurn(promptId, tool.toolUseId); - const traced = this.tools.start(turn, tool); - if (traced) turn.phase = 'active'; - return Boolean(traced); + startTool(promptId: string | undefined, call: ToolCall): boolean { + const turn = this.ensureToolTurn(promptId, call.toolUseId); + const started = this.tools.start(turn, call); + if (started) turn.phase = 'active'; + return started; } finishTool( promptId: string | undefined, - tool: ToolDescriptor, - outcome: ToolOutcome, + call: ToolCall, + result: ToolResult, ): boolean { const finished = this.tools.finishOrRecover( - () => this.ensureToolTurn(promptId, tool.toolUseId), - tool, - outcome, + () => this.ensureToolTurn(promptId, call.toolUseId), + call, + result, ); if (finished) this.finalizeIdleSupersededTurns(); return finished; @@ -169,7 +166,7 @@ export class Session { parseSessionFd(this.transcript.getFd()) ?? { turns: [] }, ).length; responseOffsetFloor = previous.responseLimit; - if (promptId === undefined || previous.children.size === 0) { + if (promptId === undefined || !this.tools.hasOpenTools(previous)) { this.finalizeTurn(previous, 'superseded_by_next_prompt'); } } @@ -315,7 +312,6 @@ export class Session { promptId: options.promptId, userText: cursor.userText, phase: 'active', - children: new Set(), responseOffset: cursor.responseOffset, seenResponses: new Set(), }; @@ -451,7 +447,7 @@ export class Session { private finalizeIdleSupersededTurns(): void { for (const turn of [...this.turns]) { - if (turn.responseLimit !== undefined && turn.children.size === 0) { + if (turn.responseLimit !== undefined && !this.tools.hasOpenTools(turn)) { this.finalizeTurn(turn, 'superseded_by_next_prompt'); } } diff --git a/tests/tool-lifecycle.test.ts b/tests/tool-lifecycle.test.ts index 38b4206..4178812 100644 --- a/tests/tool-lifecycle.test.ts +++ b/tests/tool-lifecycle.test.ts @@ -158,6 +158,40 @@ test('PostToolUseFailure records the tool result and error type', async (t) => { assert.equal(span.status.code, 2); }); +test('malformed tool payloads are rejected before session recovery', async (t) => { + const exporter = await initWeaveInMemory(); + exporter.reset(); + const sessionId = 'invalid-tool-payloads'; + const transcript = makeTranscript(t, sessionId, 'do not recover'); + const daemon = makeGenaiDaemon(); + const base = { + session_id: sessionId, + transcript_path: transcript.file, + cwd: '/x', + tool_use_id: 'invalid-tool', + tool_name: 'Read', + }; + + await daemon.routeEvent({ + hook_event_name: 'PreToolUse', ...base, tool_input: [], + }); + await daemon.routeEvent({ + hook_event_name: 'PostToolUse', ...base, + tool_input: { file_path: '/tmp/x' }, tool_response: undefined, + }); + await daemon.routeEvent({ + hook_event_name: 'PostToolUseFailure', ...base, + tool_input: { file_path: '/tmp/x' }, error: { message: 'failed' }, + }); + await daemon.routeEvent({ + hook_event_name: 'PreToolUse', ...base, + tool_input: { file_path: '/tmp/x', offset: Number.NaN }, + }); + await flushWeave(); + + assert.deepEqual(exporter.getFinishedSpans(), []); +}); + test('a restart-first terminal hook recovers one exact tool and turn', async (t) => { const exporter = await initWeaveInMemory(); exporter.reset();