// CRITICALimport type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";import type { AssistantMessage, ToolResultMessage } from "@mariozechner/pi-ai"; export type ToolExecutionInfo = { toolName: string; args: Record<string, unknown>; startedAt: string;}; export type AgentEventHandlerHelpers = { runId: string; sessionId: string; publish: (type: string, data: Record<string, unknown>) => void; toolExecutionStarts: Map<string, ToolExecutionInfo>; toolCallToMessageId: Map<string, string>; userMessageId: string; setAssistantId: (id: string | null) => void; setLastAssistantId: (id: string | null) => void; getAssistantId: () => string | null; getLastAssistantId: () => string | null; cleanMessage: (message: AgentMessage) => void; getTurnIndex: () => number; setTurnIndex: (value: number) => void; markError: (message: string, status: "error" | "aborted") => void;}; export type AgentEventHandlerOptions = { createMessageId: () => string; mapToolCallsToMessage: ( assistant: AssistantMessage, messageId: string | null, toolCallToMessageId: Map<string, string>, ) => void; persistAssistantMessage: ( sessionId: string, messageId: string, assistant: AssistantMessage, toolResults: ToolResultMessage[], runId: string, turnIndex?: number, ) => void; addToolExecution: ( runId: string, toolCallId: string, toolName: string, options: { arguments: Record<string, unknown>; resultText: string; isError: boolean; finishedAt: string; startedAt?: string; toolServer?: string; }, ) => void; parseToolServer: (toolName: string) => string | null; extractToolResultText: (result: unknown) => string;}; /** * Translate Pi agent runtime events into controller SSE events and persisted side effects. * @param event - Agent event payload. * @param helpers - Helper callbacks and mutable state for the run. * @param options - Side-effect callbacks for persistence and tool execution tracking. * @returns Nothing. */export function handleAgentEvent( event: AgentEvent, helpers: AgentEventHandlerHelpers, options: AgentEventHandlerOptions,): void { switch (event.type) { case "turn_start": { const nextIndex = helpers.getTurnIndex() + 1; helpers.setTurnIndex(nextIndex); helpers.publish("turn_start", { ...(event as Record<string, unknown>), turn_index: nextIndex }); return; } case "message_start": { const message = event.message as AgentMessage; const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; if (message.role === "assistant") { const id = options.createMessageId(); helpers.setAssistantId(id); helpers.publish("message_start", { message_id: id, message, ...turnPayload }); return; } if (message.role === "user") { helpers.publish("message_start", { message_id: helpers.userMessageId, message, ...turnPayload }); return; } helpers.publish("message_start", { message, ...turnPayload }); return; } case "message_update": { const message = event.message as AgentMessage; helpers.cleanMessage(message); const messageId = message.role === "assistant" ? helpers.getAssistantId() : undefined; const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; helpers.publish("message_update", { ...(messageId ? { message_id: messageId } : {}), message, assistantMessageEvent: event.assistantMessageEvent, ...turnPayload, }); return; } case "message_end": { const message = event.message as AgentMessage; helpers.cleanMessage(message); const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; if (message.role === "assistant") { const messageId = helpers.getAssistantId(); helpers.setLastAssistantId(messageId); options.mapToolCallsToMessage(message as AssistantMessage, messageId, helpers.toolCallToMessageId); helpers.publish("message_end", { ...(messageId ? { message_id: messageId } : {}), message, ...turnPayload }); return; } if (message.role === "user") { helpers.publish("message_end", { message_id: helpers.userMessageId, message, ...turnPayload }); return; } helpers.publish("message_end", { message, ...turnPayload }); return; } case "tool_execution_start": { const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; helpers.toolExecutionStarts.set(event.toolCallId, { toolName: event.toolName, args: event.args ?? {}, startedAt: new Date().toISOString(), }); helpers.publish("tool_execution_start", { toolCallId: event.toolCallId, toolName: event.toolName, args: event.args, message_id: helpers.toolCallToMessageId.get(event.toolCallId), ...turnPayload, }); return; } case "tool_execution_update": { const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; helpers.publish("tool_execution_update", { toolCallId: event.toolCallId, toolName: event.toolName, args: event.args, partialResult: event.partialResult, message_id: helpers.toolCallToMessageId.get(event.toolCallId), ...turnPayload, }); return; } case "tool_execution_end": { const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; const started = helpers.toolExecutionStarts.get(event.toolCallId); const finishedAt = new Date().toISOString(); const toolServer = options.parseToolServer(event.toolName); const toolExecutionOptions = { arguments: started?.args ?? {}, resultText: options.extractToolResultText(event.result?.content ?? event.result), isError: event.isError, finishedAt, ...(toolServer ? { toolServer } : {}), ...(started?.startedAt ? { startedAt: started.startedAt } : {}), }; options.addToolExecution(helpers.runId, event.toolCallId, event.toolName, toolExecutionOptions); helpers.publish("tool_execution_end", { toolCallId: event.toolCallId, toolName: event.toolName, result: event.result, isError: event.isError, message_id: helpers.toolCallToMessageId.get(event.toolCallId), ...turnPayload, }); return; } case "turn_end": { const assistant = event.message as AssistantMessage; helpers.cleanMessage(assistant as unknown as AgentMessage); const messageId = helpers.getLastAssistantId(); const turnIndex = helpers.getTurnIndex(); const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {}; if (assistant.stopReason === "error" || assistant.stopReason === "aborted") { helpers.markError( assistant.errorMessage ?? "Agent error", assistant.stopReason === "aborted" ? "aborted" : "error", ); } if (messageId) { options.persistAssistantMessage( helpers.sessionId, messageId, assistant, event.toolResults ?? [], helpers.runId, turnIndex >= 0 ? turnIndex : undefined, ); } helpers.publish("turn_end", { message: assistant, toolResults: event.toolResults ?? [], message_id: messageId, ...turnPayload, }); return; } case "agent_end": case "agent_start": default: helpers.publish(event.type, { ...(event as Record<string, unknown>) }); }}
// CRITICAL
import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ToolResultMessage } from "@mariozechner/pi-ai";
export type ToolExecutionInfo = {
toolName: string;
args: Record<string, unknown>;
startedAt: string;
};
export type AgentEventHandlerHelpers = {
runId: string;
sessionId: string;
publish: (type: string, data: Record<string, unknown>) => void;
toolExecutionStarts: Map<string, ToolExecutionInfo>;
toolCallToMessageId: Map<string, string>;
userMessageId: string;
setAssistantId: (id: string | null) => void;
setLastAssistantId: (id: string | null) => void;
getAssistantId: () => string | null;
getLastAssistantId: () => string | null;
cleanMessage: (message: AgentMessage) => void;
getTurnIndex: () => number;
setTurnIndex: (value: number) => void;
markError: (message: string, status: "error" | "aborted") => void;
};
export type AgentEventHandlerOptions = {
createMessageId: () => string;
mapToolCallsToMessage: (
assistant: AssistantMessage,
messageId: string | null,
toolCallToMessageId: Map<string, string>,
) => void;
persistAssistantMessage: (
sessionId: string,
messageId: string,
assistant: AssistantMessage,
toolResults: ToolResultMessage[],
runId: string,
turnIndex?: number,
) => void;
addToolExecution: (
runId: string,
toolCallId: string,
toolName: string,
options: {
arguments: Record<string, unknown>;
resultText: string;
isError: boolean;
finishedAt: string;
startedAt?: string;
toolServer?: string;
},
) => void;
parseToolServer: (toolName: string) => string | null;
extractToolResultText: (result: unknown) => string;
};
/**
* Translate Pi agent runtime events into controller SSE events and persisted side effects.
* @param event - Agent event payload.
* @param helpers - Helper callbacks and mutable state for the run.
* @param options - Side-effect callbacks for persistence and tool execution tracking.
* @returns Nothing.
*/
export function handleAgentEvent(
event: AgentEvent,
helpers: AgentEventHandlerHelpers,
options: AgentEventHandlerOptions,
): void {
switch (event.type) {
case "turn_start": {
const nextIndex = helpers.getTurnIndex() + 1;
helpers.setTurnIndex(nextIndex);
helpers.publish("turn_start", { ...(event as Record<string, unknown>), turn_index: nextIndex });
return;
}
case "message_start": {
const message = event.message as AgentMessage;
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
if (message.role === "assistant") {
const id = options.createMessageId();
helpers.setAssistantId(id);
helpers.publish("message_start", { message_id: id, message, ...turnPayload });
return;
}
if (message.role === "user") {
helpers.publish("message_start", { message_id: helpers.userMessageId, message, ...turnPayload });
return;
}
helpers.publish("message_start", { message, ...turnPayload });
return;
}
case "message_update": {
const message = event.message as AgentMessage;
helpers.cleanMessage(message);
const messageId = message.role === "assistant" ? helpers.getAssistantId() : undefined;
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
helpers.publish("message_update", {
...(messageId ? { message_id: messageId } : {}),
message,
assistantMessageEvent: event.assistantMessageEvent,
...turnPayload,
});
return;
}
case "message_end": {
const message = event.message as AgentMessage;
helpers.cleanMessage(message);
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
if (message.role === "assistant") {
const messageId = helpers.getAssistantId();
helpers.setLastAssistantId(messageId);
options.mapToolCallsToMessage(message as AssistantMessage, messageId, helpers.toolCallToMessageId);
helpers.publish("message_end", { ...(messageId ? { message_id: messageId } : {}), message, ...turnPayload });
return;
}
if (message.role === "user") {
helpers.publish("message_end", { message_id: helpers.userMessageId, message, ...turnPayload });
return;
}
helpers.publish("message_end", { message, ...turnPayload });
return;
}
case "tool_execution_start": {
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
helpers.toolExecutionStarts.set(event.toolCallId, {
toolName: event.toolName,
args: event.args ?? {},
startedAt: new Date().toISOString(),
});
helpers.publish("tool_execution_start", {
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
message_id: helpers.toolCallToMessageId.get(event.toolCallId),
...turnPayload,
});
return;
}
case "tool_execution_update": {
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
helpers.publish("tool_execution_update", {
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
partialResult: event.partialResult,
message_id: helpers.toolCallToMessageId.get(event.toolCallId),
...turnPayload,
});
return;
}
case "tool_execution_end": {
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
const started = helpers.toolExecutionStarts.get(event.toolCallId);
const finishedAt = new Date().toISOString();
const toolServer = options.parseToolServer(event.toolName);
const toolExecutionOptions = {
arguments: started?.args ?? {},
resultText: options.extractToolResultText(event.result?.content ?? event.result),
isError: event.isError,
finishedAt,
...(toolServer ? { toolServer } : {}),
...(started?.startedAt ? { startedAt: started.startedAt } : {}),
};
options.addToolExecution(helpers.runId, event.toolCallId, event.toolName, toolExecutionOptions);
helpers.publish("tool_execution_end", {
toolCallId: event.toolCallId,
toolName: event.toolName,
result: event.result,
isError: event.isError,
message_id: helpers.toolCallToMessageId.get(event.toolCallId),
...turnPayload,
});
return;
}
case "turn_end": {
const assistant = event.message as AssistantMessage;
helpers.cleanMessage(assistant as unknown as AgentMessage);
const messageId = helpers.getLastAssistantId();
const turnIndex = helpers.getTurnIndex();
const turnPayload = turnIndex >= 0 ? { turn_index: turnIndex } : {};
if (assistant.stopReason === "error" || assistant.stopReason === "aborted") {
helpers.markError(
assistant.errorMessage ?? "Agent error",
assistant.stopReason === "aborted" ? "aborted" : "error",
);
}
if (messageId) {
options.persistAssistantMessage(
helpers.sessionId,
messageId,
assistant,
event.toolResults ?? [],
helpers.runId,
turnIndex >= 0 ? turnIndex : undefined,
);
}
helpers.publish("turn_end", {
message: assistant,
toolResults: event.toolResults ?? [],
message_id: messageId,
...turnPayload,
});
return;
}
case "agent_end":
case "agent_start":
default:
helpers.publish(event.type, { ...(event as Record<string, unknown>) });
}
}