Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
22 changes: 22 additions & 0 deletions packages/junior/src/api/conversations/events.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export const conversationReportSourceEventTypes = [
"message_handled",
"agent_step",
"tool_execution_started",
"tool_execution_completed",
"turn_started",
"turn_routed",
"turn_completed",
Expand Down Expand Up @@ -181,6 +182,9 @@ export function conversationReportToolResultIds(
return [
...new Set(
events.flatMap((event) => {
if (event.data.type === "tool_execution_completed") {
return [event.data.toolCallId];
}
if (event.data.type !== "agent_step") return [];
const result = reportingToolResultMessageSchema.safeParse(
event.data.message,
Expand Down Expand Up @@ -346,6 +350,24 @@ export function projectConversationReportEventPage(args: {
},
],
};
} else if (event.data.type === "tool_execution_completed") {
const start = toolStarts.get(event.data.toolCallId);
data = {
type: "tool_calls",
calls: [
{
toolCallId: event.data.toolCallId,
name: event.data.toolName,
status: event.data.outcome,
...(start && start.seq < event.seq
? {
startedAt: new Date(start.createdAtMs).toISOString(),
startedSeq: start.seq,
}
: {}),
},
],
};
} else if (event.data.type === "subagent_started") {
subagentStarts.set(event.data.subagentInvocationId, {
createdAtMs: event.createdAtMs,
Expand Down
72 changes: 55 additions & 17 deletions packages/junior/src/chat/agent/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ import {
loadConnectedMcpProviders,
loadTurnRoute,
openConversationProjection,
recordToolExecutionCompleted,
recordToolExecutionStarted,
recordMcpProviderConnected,
recordTurnRoute,
Expand Down Expand Up @@ -410,17 +411,37 @@ async function executeAgentRunInPrivacyContext(
surface,
});
const runResume = resume;
const recordParentToolExecutionStart = async (event: {
args: unknown;
toolCallId: string;
toolName: string;
}) => {
let pendingToolActivityWrites = Promise.resolve();
const recordParentToolExecution = async (event:
| {
args: unknown;
toolCallId: string;
toolName: string;
type: "tool_execution_start";
}
| {
isError: boolean;
result: unknown;
toolCallId: string;
toolName: string;
type: "tool_execution_end";
},
) => {
try {
await recordToolExecutionStarted({
conversationId,
toolCallId: event.toolCallId,
toolName: event.toolName,
});
if (event.type === "tool_execution_start") {
await recordToolExecutionStarted({
conversationId,
toolCallId: event.toolCallId,
toolName: event.toolName,
});
} else {
await recordToolExecutionCompleted({
conversationId,
isError: event.isError,
toolCallId: event.toolCallId,
toolName: event.toolName,
});
}
} catch (error) {
// Host-only activity events are best-effort reporting writes; a
// failed append must not abort the in-flight model turn.
Expand All @@ -431,10 +452,17 @@ async function executeAgentRunInPrivacyContext(
{
"gen_ai.tool.name": event.toolName,
},
"Failed to record host-only tool execution start",
"Failed to record host-only tool execution",
);
}
};
const enqueueParentToolExecution = (
event: Parameters<typeof recordParentToolExecution>[0],
): void => {
pendingToolActivityWrites = pendingToolActivityWrites.then(() =>
recordParentToolExecution(event),
);
};
const persistedConfigurationValues = policy.channelConfiguration
? await policy.channelConfiguration.resolveValues()
: {};
Expand Down Expand Up @@ -998,16 +1026,25 @@ async function executeAgentRunInPrivacyContext(
});

const unsubscribe = agent.subscribe((event) => {
if (event.type === "tool_execution_start") {
return recordParentToolExecutionStart(event);
if (
event.type === "tool_execution_start" ||
event.type === "tool_execution_end"
) {
// Pi emits tool_execution_end before appending the corresponding
// toolResult message. Queue reporting writes without blocking that
// lifecycle so timeout recovery can snapshot the continuable result.
enqueueParentToolExecution(event);
return;
}
if (event.type === "turn_end" && event.toolResults.length > 0) {
if (pendingHandoff) {
return;
return pendingToolActivityWrites;
}
return runResume
.persistSafeBoundary([...agent!.state.messages])
.then(() => undefined);
return pendingToolActivityWrites.then(() =>
runResume
.persistSafeBoundary([...agent!.state.messages])
.then(() => undefined),
);
}
if (event.type === "message_end" && isAssistantMessage(event.message)) {
if (
Expand Down Expand Up @@ -1336,6 +1373,7 @@ async function executeAgentRunInPrivacyContext(
return authPauseOutcome;
}
} finally {
await pendingToolActivityWrites;
unsubscribe();
}

Expand Down
11 changes: 11 additions & 0 deletions packages/junior/src/chat/conversations/history.ts
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,15 @@ const toolExecutionStartedEventDataSchema = z
})
.strict();

const toolExecutionCompletedEventDataSchema = z
.object({
type: z.literal("tool_execution_completed"),
toolCallId: z.string().min(1),
toolName: z.string().min(1),
outcome: z.enum(["completed", "error"]),
})
.strict();

const conversationMessageRoleSchema = z.union([
z.literal("user"),
z.literal("assistant"),
Expand Down Expand Up @@ -307,6 +316,7 @@ const appendableConversationEventDataSchema = z.union([
authorizationRequestedEventDataSchema,
authorizationCompletedEventDataSchema,
toolExecutionStartedEventDataSchema,
toolExecutionCompletedEventDataSchema,
messageHandledEventDataSchema,
messagesSummarizedEventDataSchema,
turnStartedEventDataSchema,
Expand Down Expand Up @@ -339,6 +349,7 @@ const knownConversationEventTypeSchema = z.enum([
"authorization_requested",
"authorization_completed",
"tool_execution_started",
"tool_execution_completed",
"message_handled",
"messages_summarized",
"turn_started",
Expand Down
21 changes: 21 additions & 0 deletions packages/junior/src/chat/conversations/projection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -458,3 +458,24 @@ export async function recordToolExecutionStarted(args: {
},
]);
}

/** Record a host-observed parent tool completion without persisting its output. */
export async function recordToolExecutionCompleted(args: {
conversationId: string;
createdAtMs?: number;
isError: boolean;
toolCallId: string;
toolName: string;
}): Promise<void> {
await getConversationEventStore().append(args.conversationId, [
{
data: {
type: "tool_execution_completed",
toolCallId: args.toolCallId,
toolName: args.toolName,
outcome: args.isError ? "error" : "completed",
},
createdAtMs: args.createdAtMs ?? Date.now(),
},
]);
}
Original file line number Diff line number Diff line change
Expand Up @@ -216,13 +216,10 @@ describe("conversation event list API", () => {
},
{
data: {
type: "agent_step",
message: {
role: "toolResult",
toolCallId: "search-before-page",
content: [{ type: "text", text: "two matches" }],
isError: false,
} as PiMessage,
type: "tool_execution_completed",
toolCallId: "search-before-page",
toolName: "search",
outcome: "completed",
},
createdAtMs: 3,
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,15 @@ vi.mock("@/chat/conversations/projection", async (importOriginal) => {
}
return actual.recordToolExecutionStarted(...args);
},
recordToolExecutionCompleted: async (
...args: Parameters<typeof actual.recordToolExecutionCompleted>
) => {
sessionLogState.toolExecutionAppendCalls += 1;
if (sessionLogState.failToolExecutionAppend) {
throw new Error("store blip during host-only append");
}
return actual.recordToolExecutionCompleted(...args);
},
};
});

Expand Down Expand Up @@ -223,6 +232,22 @@ vi.mock("@earendil-works/pi-agent-core", () => {
this.recordRunFailure(error);
return {};
}
try {
await Promise.all(
this.subscribers.map((subscriber) =>
subscriber({
type: "tool_execution_end",
toolCallId: "call_1",
toolName: "bash",
result: { content: [{ type: "text", text: "ok" }] },
isError: false,
}),
),
);
} catch (error) {
this.recordRunFailure(error);
return {};
}
this.state.messages.push({
role: "toolResult",
toolName: "bash",
Expand Down Expand Up @@ -1094,7 +1119,7 @@ describe("executeAgentRun provider retry", () => {
}),
);

expect(sessionLogState.toolExecutionAppendCalls).toBe(1);
expect(sessionLogState.toolExecutionAppendCalls).toBe(2);
expect(reply.diagnostics.outcome).toBe("success");
expect(reply.text).toBe("Tool done.");
});
Expand Down
Loading