Skip to content

feat(trace): trace Agent lifecycles #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rgao-coreweave merged 1 commit into from
Aug 10, 2026
Merged

feat(trace): trace Agent lifecycles #131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
392 changes: 323 additions & 69 deletions src/callLifecycle.ts
View file
Open in desktop

Large diffs are not rendered by default.

300 changes: 289 additions & 11 deletions src/hookHandler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,26 +5,50 @@
import * as fs from 'fs';
import * as path from 'path';
import type {
BaseHookInput,
HookInput,
InstructionsLoadedHookInput,
PreToolUseHookInput,
PreCompactHookInput,
SessionEndHookInput,
SessionStartHookInput,
StopHookInput,
SubagentStartHookInput,
SubagentStopHookInput,
UserPromptSubmitHookInput,
} from '@anthropic-ai/claude-agent-sdk';
import * as weave from 'weave';
import { emitChatSpans } from './chatSpans.js';
import {
backfillAgentPrompt,
beginCall,
bindAgent,
matchAgent,
recordAgentStop,
recordCallOutcome,
recoverAgentCall,
responseKeysForAgent,
} from './callLifecycle.js';
import type {
AgentMatch,
CallParent,
JsonObject,
JsonValue,
ToolCall,
ToolResult,
TracedAgent,
TracedCall,
} from './callLifecycle.js';
import type { CompactionAttrs } from './genaiSpans.js';
import { snippet } from './genaiSpans.js';
import { ATTR, assistantOutputMessages, snippet } from './genaiSpans.js';
import type { SpanParent } from './genaiSpans.js';
import { parseSessionFd } from './parser.js';
import { Session } from './session.js';
import { TranscriptFile } from './transcriptFile.js';
import {
TranscriptFile,
readSubagentPrompt,
subagentTranscriptPath,
} from './transcriptFile.js';

type TraceLog = (level: 'DEBUG' | 'INFO' | 'ERROR', message: string) => void;

Expand All @@ -33,6 +57,23 @@ type HookInputFor<Event extends HookInput['hook_event_name']> = Extract<
{ hook_event_name: Event }
>;
type PostToolResultHookInput = HookInputFor<'PostToolUse' | 'PostToolUseFailure'>;
type RecoverCallHookInput = PostToolResultHookInput;

function mergeSubagentOutput(transcriptText?: string, lastMessage?: string): string | undefined {
const transcript = transcriptText?.trim();
const latest = lastMessage?.trim();
if (!transcript) return latest || undefined;
if (!latest) return transcript;

const contains = (text: string, message: string) =>
text === message
|| text.startsWith(`${message}\n`)
|| text.endsWith(`\n${message}`)
|| text.includes(`\n${message}\n`);
if (contains(transcript, latest)) return transcript;
if (contains(latest, transcript)) return latest;
return `${transcript}\n${latest}`;
}

function isJsonValue(value: unknown): value is JsonValue {
if (value === null) return true;
Expand Down Expand Up @@ -132,6 +173,12 @@ export class HookHandler {
case 'PostToolUseFailure':
await this.handlePostToolResult(sessionId, input);
return;
case 'SubagentStart':
await this.handleSubagentStart(sessionId, input);
return;
case 'SubagentStop':
await this.handleSubagentStop(sessionId, input);
return;
case 'PreCompact':
this.handlePreCompact(sessionId, input);
return;
Expand Down Expand Up @@ -304,25 +351,95 @@ export class HookHandler {
sessionId: string,
input: PreToolUseHookInput,
): Promise<void> {
if (input.tool_name === 'Agent' || input.agent_id) return;
const call = this.toolCall(input);
if (!call) return;
const descriptor = this.toolCall(input);
if (!descriptor) return;
const session = await this.getOrReconstructSession(sessionId, input);
if (!session) return;
session.startTool(input.prompt_id, call);

const parent = await this.resolveCallParent(session, input);
if (!parent) {
this.log(
'ERROR',
`PreToolUse: unknown parent session=${sessionId} tool=${input.tool_name} agent=${input.agent_id ?? 'root'}`,
);
return;
}
const call = beginCall(session.calls, parent, descriptor);
if (call && !input.agent_id) call.root.phase = 'active';
}

/** Resolve a call's owning span. After restart, nested hooks can arrive
* before SubagentStart; recover only from the stable id, type, and prompt. */
private async resolveCallParent(
session: Session,
input: Pick<BaseHookInput, 'agent_id' | 'agent_type' | 'prompt_id'>,
): Promise<CallParent | undefined> {
if (!input.agent_id) return session.ensureTurn(input.prompt_id);
const active = session.calls.byAgentId.get(input.agent_id);
if (active) return active;
if (!input.agent_type || session.calls.agentTombstones.has(input.agent_id)) {
return undefined;
}

const transcriptPath = subagentTranscriptPath(
session.transcriptPath,
input.agent_id,
);
const prompt = await readSubagentPrompt(transcriptPath);
if (!prompt) {
this.log(
'ERROR',
`Nested hook: cannot recover owner agentId=${input.agent_id} type=${input.agent_type} without its dispatch prompt`,
);
return undefined;
}
return this.recoverAgent(
session,
input.agent_id,
input.agent_type,
input.prompt_id,
prompt,
'SubagentStart',
);
}

/** Recreate a call from its exact tool_use_id after a restart. */
private async recoverCall(
session: Session,
input: RecoverCallHookInput,
descriptor: ToolCall,
): Promise<TracedCall | undefined> {
const existing = session.calls.byToolUseId.get(descriptor.toolUseId);
if (existing || session.calls.toolUseTombstones.has(descriptor.toolUseId)) {
return existing;
}
const parent = input.agent_id
? await this.resolveCallParent(session, input)
: session.ensureToolTurn(input.prompt_id, descriptor.toolUseId);
if (!parent) return undefined;
return beginCall(session.calls, parent, descriptor);
}

private async handlePostToolResult(
sessionId: string,
input: PostToolResultHookInput,
): Promise<void> {
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 (typeof input.tool_use_id !== 'string' || !result) return;

const existingSession = this.sessions.get(sessionId);
if (existingSession?.calls.toolUseTombstones.has(input.tool_use_id)) return;
const existingCall = existingSession?.calls.byToolUseId.get(input.tool_use_id);
const descriptor = existingCall ? undefined : this.toolCall(input);
if (!existingCall && !descriptor) return;

const session = existingSession
?? await this.getOrReconstructSession(sessionId, input);
if (!session) return;
session.finishTool(input.prompt_id, call, result);

if (!existingCall) await this.recoverCall(session, input, descriptor!);
recordCallOutcome(session.calls, input.tool_use_id, result);
session.finishSupersededTurns();
}

private toolCall(
Expand Down Expand Up @@ -358,6 +475,167 @@ export class HookHandler {
return { ok: false, error: input.error };
}

private async handleSubagentStart(
sessionId: string,
input: SubagentStartHookInput,
): Promise<void> {
const session = await this.getOrReconstructSession(sessionId, input);
if (!session
|| session.calls.agentTombstones.has(input.agent_id)
|| session.calls.byAgentId.has(input.agent_id)) return;

const transcriptPath = subagentTranscriptPath(session.transcriptPath, input.agent_id);
const prompt = await readSubagentPrompt(transcriptPath);
const match = matchAgent(session.calls, input.agent_type, prompt, input.prompt_id);
if (match.kind === 'ambiguous') {
this.log(
'ERROR',
`SubagentStart: ambiguous dispatch agentId=${input.agent_id} type=${input.agent_type}`,
);
return;
}

if (match.kind === 'found') {
bindAgent(session.calls, match, input.agent_id, input.agent_type);
} else {
this.recoverAgent(
session,
input.agent_id,
input.agent_type,
input.prompt_id,
prompt ?? '',
'SubagentStart',
);
}
this.log('INFO', `Subagent started: agentId=${input.agent_id} type=${input.agent_type}`);
}

private recoverAgent(
session: Session,
agentId: string,
agentType: string,
promptId: string | undefined,
prompt: string,
event: 'SubagentStart' | 'SubagentStop',
): TracedAgent {
const recovered = recoverAgentCall(session.calls, session.ensureTurn(promptId), {
agentId,
agentType,
prompt,
event,
});
this.log('INFO', `${event}: recovered agentId=${agentId} type=${agentType}`);
return recovered;
}

private emitSubagentTranscript(
parent: SpanParent,
transcriptPath: string,
agentType: string,
seen?: Set<string>,
): { model?: string; text?: string } {
let transcript: TranscriptFile | undefined;
try {
transcript = new TranscriptFile(transcriptPath);
const turn = parseSessionFd(transcript.getFd())?.turns.at(-1);
if (!turn) return {};
emitChatSpans(parent, turn.responses, { agentName: agentType, seen });
return { model: turn.model, text: turn.text.join('\n') || undefined };
} catch (error) {
this.log('DEBUG', `SubagentStop: could not parse transcript: ${error}`);
return {};
} finally {
transcript?.close();
}
}

private async handleSubagentStop(
sessionId: string,
input: SubagentStopHookInput,
): Promise<void> {
const session = await this.getOrReconstructSession(sessionId, input);
if (!session || session.calls.agentTombstones.has(input.agent_id)) return;

const transcriptPath = input.agent_transcript_path
?? subagentTranscriptPath(session.transcriptPath, input.agent_id);
const active = session.calls.byAgentId.get(input.agent_id);
let prompt = active?.prompt;
if (!prompt) {
prompt = await readSubagentPrompt(transcriptPath);
if (active && prompt) backfillAgentPrompt(active, prompt);
}
const match: AgentMatch = active
? { kind: 'found', call: active }
: matchAgent(session.calls, input.agent_type, prompt, input.prompt_id);

if (match.kind === 'found' && !match.call.agentId) {
bindAgent(session.calls, match, input.agent_id, input.agent_type);
this.log(
'INFO',
`SubagentStop: late-matched agentId=${input.agent_id} type=${input.agent_type}`,
);
}

const turn = session.turnForPrompt(input.prompt_id);
const recovered = match.kind === 'missing'
? this.recoverAgent(
session,
input.agent_id,
input.agent_type,
input.prompt_id,
prompt ?? '',
'SubagentStop',
)
: undefined;
const lifecycle = match.kind === 'found' ? match.call : recovered;
const parent = match.kind === 'found'
? match.call.span
: recovered?.span ?? turn?.span;
if (!parent) {
this.log(
'ERROR',
`SubagentStop: no parent agentId=${input.agent_id} type=${input.agent_type}`,
);
return;
}

const seen = responseKeysForAgent(session.calls, input.agent_id, lifecycle);
const transcript = this.emitSubagentTranscript(
parent,
transcriptPath,
input.agent_type,
seen,
);
const text = mergeSubagentOutput(transcript.text, input.last_assistant_message);

if (match.kind === 'found') {
if (transcript.model) {
match.call.span.setAttributes({ [ATTR.RESPONSE_MODEL]: transcript.model });
}
if (!match.call.toolUseId && text) {
match.call.span.setAttributes({
[ATTR.OUTPUT_MESSAGES]: assistantOutputMessages([text]),
});
}
recordAgentStop(session.calls, match);
} else if (recovered) {
if (transcript.model) {
recovered.span.setAttributes({ [ATTR.RESPONSE_MODEL]: transcript.model });
}
if (text) {
recovered.span.setAttributes({
[ATTR.OUTPUT_MESSAGES]: assistantOutputMessages([text]),
});
}
}

this.log(
'DEBUG',
`Subagent stopped: agentId=${input.agent_id} type=${input.agent_type} model=${transcript.model ?? 'unknown'} match=${match.kind}`,
);
session.finishSupersededTurns();
}

private async handleStop(sessionId: string, input: StopHookInput): Promise<void> {
const session = await this.getOrReconstructSession(sessionId, input);
if (!session) return;
Expand Down
Loading
Loading