Skip to content

fix(trace): harden Agent Team transcript correlation #154

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

Open
rgao-coreweave wants to merge 1 commit into from
Open

fix(trace): harden Agent Team transcript correlation #154

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
3 changes: 2 additions & 1 deletion src/daemon.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,8 +295,9 @@ export class Daemon {
return;
}

socket.end();
void this.routeEvent(payload);
// Capture receipt-time transcript evidence before acknowledging the hook.
socket.end();
});

socket.on('error', (err: Error) => {
Expand Down
158 changes: 138 additions & 20 deletions src/hookHandler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ import type { SpanParent } from './genaiSpans.js';
import { parseSessionFd } from './parser.js';
import { Session } from './session.js';
import { TeamCoordinator } from './teamCoordinator.js';
import type { TeamCompletion } from './teamCoordinator.js';
import type { TeamTranscriptSnapshot } from './teamTranscripts.js';
import {
TranscriptFile,
readSubagentPrompt,
Expand All @@ -67,6 +69,8 @@ type RecoverCallHookInput = HookInputFor<
'PermissionDenied' | 'PostToolUse' | 'PostToolUseFailure'
>;

const MAX_RECENT_TRANSCRIPTS = 512;

function mergeSubagentOutput(transcriptText?: string, lastMessage?: string): string | undefined {
const transcript = transcriptText?.trim();
const latest = lastMessage?.trim();
Expand Down Expand Up @@ -122,6 +126,9 @@ export class HookHandler {
private readonly sessions = new Map<string, Session>();
private eventQueue = Promise.resolve();
private eventSequence = 0;
/** Hooks can wait in the queue before their sessions are reconstructed. Keep
* recently observed roots available for receipt-time Team snapshots. */
private readonly recentTranscripts = new Set<string>();
/** InstructionsLoaded can arrive before SessionStart. */
private readonly pendingInstructions = new Map<string, Map<string, string>>();
private readonly teams = new TeamCoordinator();
Expand All @@ -139,25 +146,60 @@ export class HookHandler {
}

const sequence = ++this.eventSequence;
const next = this.eventQueue.then(() => this.route(input, sequence));
this.rememberTranscript(input);
const transcriptSnapshots = input.hook_event_name === 'TeammateIdle'
? this.teams.snapshotTranscripts(
input.session_id,
input.transcript_path,
this.sessions.values(),
this.recentTranscripts.values(),
)
: undefined;
const next = this.eventQueue.then(() =>
this.route(input, sequence, transcriptSnapshots));
this.eventQueue = next;
await next;
}

private async route(input: HookInput, sequence: number): Promise<void> {
private rememberTranscript(input: HookInput): void {
const transcriptPath = input.transcript_path;
if (typeof transcriptPath !== 'string') return;
try {
const resolvedPath = new TranscriptFile(transcriptPath).resolvedPath;
this.recentTranscripts.delete(resolvedPath);
this.recentTranscripts.add(resolvedPath);
if (this.recentTranscripts.size > MAX_RECENT_TRANSCRIPTS) {
const oldest = this.recentTranscripts.values().next().value;
if (oldest !== undefined) this.recentTranscripts.delete(oldest);
}
} catch {
// Event processing reports invalid transcript paths.
}
}

private async route(
input: HookInput,
sequence: number,
transcriptSnapshots?: TeamTranscriptSnapshot[],
): Promise<void> {
const sessionId = input.session_id;
this.log(
'INFO',
`${input.hook_event_name} session=${sessionId}${input.agent_id ? ` agent=${input.agent_id}` : ''}`,
);
try {
await weave.runIsolated(() => this.dispatchEvent(input, sequence));
await weave.runIsolated(() =>
this.dispatchEvent(input, sequence, transcriptSnapshots));
} catch (err) {
this.log('ERROR', `Error handling ${input.hook_event_name}: ${err}`);
}
}

private async dispatchEvent(input: HookInput, sequence: number): Promise<void> {
private async dispatchEvent(
input: HookInput,
sequence: number,
transcriptSnapshots?: TeamTranscriptSnapshot[],
): Promise<void> {
const sessionId = input.session_id;
switch (input.hook_event_name) {
case 'SessionStart':
Expand All @@ -183,13 +225,18 @@ export class HookHandler {
await this.handlePostToolResult(sessionId, input, sequence);
return;
case 'SubagentStart':
await this.handleSubagentStart(sessionId, input);
await this.handleSubagentStart(sessionId, input, sequence);
return;
case 'SubagentStop':
await this.handleSubagentStop(sessionId, input);
await this.handleSubagentStop(sessionId, input, sequence);
return;
case 'TeammateIdle':
await this.handleTeammateIdle(input, sequence);
await this.handleTeammateIdle(
sessionId,
input,
sequence,
transcriptSnapshots,
);
return;
case 'PreCompact':
this.handlePreCompact(sessionId, input);
Expand Down Expand Up @@ -383,7 +430,13 @@ export class HookHandler {
}
const call = beginCall(session.calls, parent, descriptor);
if (call?.kind === 'agent') {
this.teams.registerDispatch(session, call, sequence);
const team = this.teams.registerDispatch(session, call, sequence);
if (team) {
this.log(
'INFO',
`Team member registered: ${team.teamName ?? 'implicit'}::${team.memberName} (queue depth ${team.depth})`,
);
}
}
if (call && !input.agent_id) call.root.phase = 'active';
}
Expand Down Expand Up @@ -461,12 +514,13 @@ export class HookHandler {
const call = existingCall
?? await this.recoverCall(session, input, descriptor!);
if (!existingCall && call?.kind === 'agent') {
this.teams.registerDispatch(session, call, sequence);
this.teams.registerDispatch(session, call, sequence, true);
}
if (call?.kind === 'agent') {
const update = this.teams.postOutcome(call, result);
const update = await this.teams.postOutcome(call, result);
this.settleTeamCompletions(update.completions);
if (update.handled) {
this.settleTeamCompletions(update.completions, session);
if (!update.completions.length) this.settleSession(session);
return;
}
}
Expand Down Expand Up @@ -548,7 +602,7 @@ export class HookHandler {
const call = existingCall
?? await this.recoverCall(session, input, descriptor!);
const completions = call?.kind === 'agent'
? this.teams.deny(call, input.reason)
? await this.teams.deny(call, input.reason)
: undefined;
if (!completions) denyCall(session.calls, input.tool_use_id, input.reason);
this.settleTeamCompletions(completions ?? [], session);
Expand All @@ -557,6 +611,7 @@ export class HookHandler {
private async handleSubagentStart(
sessionId: string,
input: SubagentStartHookInput,
sequence: number,
): Promise<void> {
const session = await this.getOrReconstructSession(sessionId, input);
if (!session
Expand All @@ -574,10 +629,24 @@ export class HookHandler {
return;
}

const teamLifecycle = this.teams.classifyLifecycle(
session,
input.agent_type,
transcriptPath,
);
if (
match.kind === 'missing'
&& (teamLifecycle === 'dispatch' || teamLifecycle === 'ambiguous')
) {
return;
}

let lifecycle: TracedAgent;
if (match.kind === 'found') {
bindAgent(session.calls, match, input.agent_id, input.agent_type);
lifecycle = match.call;
} else {
this.recoverAgent(
lifecycle = this.recoverAgent(
session,
input.agent_id,
input.agent_type,
Expand All @@ -586,6 +655,15 @@ export class HookHandler {
'SubagentStart',
);
}
if (lifecycle.toolUseId === undefined && teamLifecycle === 'idle') {
this.teams.registerIdle(
session,
lifecycle,
input.agent_type,
transcriptPath,
sequence,
);
}
this.log('INFO', `Subagent started: agentId=${input.agent_id} type=${input.agent_type}`);
}

Expand Down Expand Up @@ -631,6 +709,7 @@ export class HookHandler {
private async handleSubagentStop(
sessionId: string,
input: SubagentStopHookInput,
sequence: number,
): Promise<void> {
const session = await this.getOrReconstructSession(sessionId, input);
if (!session || session.calls.agentTombstones.has(input.agent_id)) return;
Expand All @@ -655,6 +734,18 @@ export class HookHandler {
);
}

const teamLifecycle = this.teams.classifyLifecycle(
session,
input.agent_type,
transcriptPath,
);
if (
match.kind === 'missing'
&& (teamLifecycle === 'dispatch' || teamLifecycle === 'ambiguous')
) {
return;
}

const turn = session.turnForPrompt(input.prompt_id);
const recovered = match.kind === 'missing'
? this.recoverAgent(
Expand All @@ -667,12 +758,29 @@ export class HookHandler {
)
: undefined;
const lifecycle = match.kind === 'found' ? match.call : recovered;
if (lifecycle && this.teams.has(lifecycle)) {
this.log(
'DEBUG',
`Subagent stopped: agentId=${input.agent_id} awaiting TeammateIdle`,
if (
lifecycle
&& lifecycle.toolUseId === undefined
&& teamLifecycle === 'idle'
) {
this.teams.registerIdle(
session,
lifecycle,
input.agent_type,
transcriptPath,
sequence,
);
return;
}
if (lifecycle) {
const update = await this.teams.stop(lifecycle, transcriptPath);
this.settleTeamCompletions(update.completions);
if (update.handled) {
this.log(
'DEBUG',
`Subagent stopped: agentId=${input.agent_id} awaiting TeammateIdle`,
);
return;
}
}
const parent = match.kind === 'found'
? match.call.span
Expand Down Expand Up @@ -723,14 +831,18 @@ export class HookHandler {
}

private async handleTeammateIdle(
sessionId: string,
input: TeammateIdleHookInput,
sequence: number,
transcriptSnapshots?: TeamTranscriptSnapshot[],
): Promise<void> {
const result = await this.teams.recordIdle({
sequence,
sessionId,
teamName: input.team_name,
memberName: input.teammate_name,
transcriptPath: input.transcript_path,
idleTranscriptPath: input.transcript_path,
transcriptSnapshots,
});
if (!result.completions.length) {
this.log(
Expand Down Expand Up @@ -798,9 +910,15 @@ export class HookHandler {
}

private settleTeamCompletions(
completions: Array<{ owner: Session }>,
completions: TeamCompletion[],
fallback?: Session,
): void {
for (const completion of completions) {
this.log(
'DEBUG',
`Team completed: ${completion.teamName ?? 'unknown'}::${completion.memberName} (${completion.mode})`,
);
}
const owners = new Set(completions.map(completion => completion.owner));
if (!owners.size && fallback) owners.add(fallback);
for (const owner of owners) this.settleSession(owner);
Expand Down
15 changes: 11 additions & 4 deletions src/parser.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ export interface ParsedSession {
turns: ParsedTurn[];
}

export function parseSessionFd(fd: number): ParsedSession | null {
return parseSessionReader(() => readUtf8FromFd(fd));
export function parseSessionFd(fd: number, maxBytes?: number): ParsedSession | null {
return parseSessionReader(() => readUtf8FromFd(fd, maxBytes));
}

function parseSessionReader(read: () => string): ParsedSession | null {
Expand All @@ -55,8 +55,12 @@ function parseSessionReader(read: () => string): ParsedSession | null {
}
}

function readUtf8FromFd(fd: number): string {
const size = fs.fstatSync(fd).size;
function readUtf8FromFd(fd: number, maxBytes?: number): string {
const fileSize = fs.fstatSync(fd).size;
if (maxBytes !== undefined && fileSize < maxBytes) {
throw new Error('transcript shortened before bounded read');
}
const size = maxBytes ?? fileSize;
if (size === 0) return '';

const buffer = Buffer.allocUnsafe(size);
Expand All @@ -66,6 +70,9 @@ function readUtf8FromFd(fd: number): string {
if (count === 0) break;
bytesRead += count;
}
if (maxBytes !== undefined && bytesRead !== size) {
throw new Error('transcript shortened during bounded read');
}
return buffer.toString('utf8', 0, bytesRead);
}

Expand Down
Loading
Loading