Skip to content

Commit d942194

Browse files
vincentkocvincentkoc
committed
Hooks: rebase lifecycle boundaries onto current runner
1 parent d967009 commit d942194

7 files changed

Lines changed: 361 additions & 13 deletions

File tree

‎src/agents/pi-embedded-runner/run.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1471,7 +1471,9 @@ export async function runEmbeddedPiAgent(
14711471
pendingToolCalls: attempt.clientToolCall
14721472
? [
14731473
{
1474-
id: randomBytes(5).toString("hex").slice(0, 9),
1474+
id:
1475+
attempt.clientToolCall.toolCallId.trim() ||
1476+
randomBytes(5).toString("hex").slice(0, 9),
14751477
name: attempt.clientToolCall.name,
14761478
arguments: JSON.stringify(attempt.clientToolCall.params),
14771479
},

‎src/agents/pi-embedded-runner/run/attempt.ts‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,12 @@ import {
130130
} from "./compaction-timeout.js";
131131
import { pruneProcessedHistoryImages } from "./history-image-prune.js";
132132
import { detectAndLoadPromptImages } from "./images.js";
133+
import {
134+
createInternalAgentHookEmitter,
135+
createResponseLifecycleTracker,
136+
emitThinkingEnd,
137+
emitThinkingStart,
138+
} from "./lifecycle-hooks.js";
133139
import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js";
134140

135141
type PromptBuildHookRunner = {
@@ -750,6 +756,9 @@ export async function runEmbeddedAttempt(
750756
const prevCwd = process.cwd();
751757
const runAbortController = new AbortController();
752758
ensureGlobalUndiciStreamTimeouts();
759+
const emitInternalAgentHook = createInternalAgentHookEmitter(params, (message) =>
760+
log.warn(message),
761+
);
753762

754763
log.debug(
755764
`embedded run start: runId=${params.runId} sessionId=${params.sessionId} provider=${params.provider} model=${params.modelId} thinking=${params.thinkLevel} messageChannel=${params.messageChannel ?? params.messageProvider ?? "unknown"}`,
@@ -1153,16 +1162,20 @@ export async function runEmbeddedAttempt(
11531162
});
11541163

11551164
// Add client tools (OpenResponses hosted tools) to customTools
1156-
let clientToolCallDetected: { name: string; params: Record<string, unknown> } | null = null;
1165+
let clientToolCallDetected: {
1166+
name: string;
1167+
params: Record<string, unknown>;
1168+
toolCallId: string;
1169+
} | null = null;
11571170
const clientToolLoopDetection = resolveToolLoopDetectionConfig({
11581171
cfg: params.config,
11591172
agentId: sessionAgentId,
11601173
});
11611174
const clientToolDefs = clientTools
11621175
? toClientToolDefinitions(
11631176
clientTools,
1164-
(toolName, toolParams) => {
1165-
clientToolCallDetected = { name: toolName, params: toolParams };
1177+
(toolName, toolParams, toolCallId) => {
1178+
clientToolCallDetected = { name: toolName, params: toolParams, toolCallId };
11661179
},
11671180
{
11681181
agentId: sessionAgentId,
@@ -1507,6 +1520,13 @@ export async function runEmbeddedAttempt(
15071520
});
15081521
};
15091522

1523+
const responseLifecycle = createResponseLifecycleTracker({
1524+
params,
1525+
emitInternalAgentHook,
1526+
onAssistantMessageStart: params.onAssistantMessageStart,
1527+
formatError: describeUnknownError,
1528+
});
1529+
15101530
const subscription = subscribeEmbeddedPiSession({
15111531
session: activeSession,
15121532
runId: params.runId,
@@ -1524,7 +1544,7 @@ export async function runEmbeddedAttempt(
15241544
blockReplyBreak: params.blockReplyBreak,
15251545
blockReplyChunking: params.blockReplyChunking,
15261546
onPartialReply: params.onPartialReply,
1527-
onAssistantMessageStart: params.onAssistantMessageStart,
1547+
onAssistantMessageStart: responseLifecycle.handleAssistantMessageStart,
15281548
onAgentEvent: params.onAgentEvent,
15291549
enforceFinalTag: params.enforceFinalTag,
15301550
config: params.config,
@@ -1627,6 +1647,7 @@ export async function runEmbeddedAttempt(
16271647
const prePromptMessageCount = activeSession.messages.length;
16281648
try {
16291649
const promptStartedAt = Date.now();
1650+
void emitThinkingStart(params, emitInternalAgentHook);
16301651

16311652
// Run before_prompt_build hooks to allow plugins to inject prompt context.
16321653
// Legacy compatibility: before_agent_start is also checked for context fields.
@@ -1788,6 +1809,13 @@ export async function runEmbeddedAttempt(
17881809
log.debug(
17891810
`embedded run prompt end: runId=${params.runId} sessionId=${params.sessionId} durationMs=${Date.now() - promptStartedAt}`,
17901811
);
1812+
void emitThinkingEnd({
1813+
lifecycleParams: params,
1814+
emitInternalAgentHook,
1815+
promptStartedAt,
1816+
promptError,
1817+
formatError: describeUnknownError,
1818+
});
17911819
}
17921820

17931821
// Capture snapshot before compaction wait so we have complete messages if timeout occurs
@@ -1941,6 +1969,12 @@ export async function runEmbeddedAttempt(
19411969
: undefined,
19421970
});
19431971
anthropicPayloadLogger?.recordUsage(messagesSnapshot, promptError);
1972+
const hasResponseOutput =
1973+
assistantTexts.length > 0 ||
1974+
didSendViaMessagingTool() ||
1975+
getMessagingToolSentTexts().length > 0;
1976+
await responseLifecycle.emitResponseStartIfNeeded(hasResponseOutput);
1977+
await responseLifecycle.emitResponseEnd(promptError);
19441978

19451979
// Run agent_end hooks to allow plugins to analyze the conversation
19461980
// This is fire-and-forget, so we don't await

‎src/agents/pi-embedded-runner/run/lifecycle-hooks.test.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+
3+
const { triggerInternalHook } = vi.hoisted(() => ({
4+
triggerInternalHook: vi.fn(async () => {}),
5+
}));
6+
7+
vi.mock("../../../hooks/internal-hooks.js", async () => {
8+
const actual = await vi.importActual<typeof import("../../../hooks/internal-hooks.js")>(
9+
"../../../hooks/internal-hooks.js",
10+
);
11+
return {
12+
...actual,
13+
triggerInternalHook,
14+
};
15+
});
16+
17+
import {
18+
createInternalAgentHookEmitter,
19+
createResponseLifecycleTracker,
20+
emitThinkingEnd,
21+
emitThinkingStart,
22+
} from "./lifecycle-hooks.js";
23+
24+
type HookEventRecord = {
25+
type?: string;
26+
action?: string;
27+
sessionKey?: string;
28+
context?: Record<string, unknown>;
29+
};
30+
31+
describe("lifecycle hooks", () => {
32+
beforeEach(() => {
33+
triggerInternalHook.mockClear();
34+
vi.useFakeTimers();
35+
vi.setSystemTime(new Date("2026-03-07T12:00:00.000Z"));
36+
});
37+
38+
afterEach(() => {
39+
vi.useRealTimers();
40+
});
41+
42+
it("emits thinking lifecycle hooks with session fallback and error context", async () => {
43+
const warn = vi.fn();
44+
const emitInternalAgentHook = createInternalAgentHookEmitter(
45+
{
46+
sessionKey: "",
47+
runId: "run-1",
48+
},
49+
warn,
50+
);
51+
const lifecycleParams = {
52+
sessionId: "session-1",
53+
runId: "run-1",
54+
provider: "openai",
55+
modelId: "gpt-5",
56+
messageProvider: "telegram",
57+
messageChannel: "telegram",
58+
} as const;
59+
60+
await emitThinkingStart(lifecycleParams, emitInternalAgentHook);
61+
vi.advanceTimersByTime(275);
62+
await emitThinkingEnd({
63+
lifecycleParams,
64+
emitInternalAgentHook,
65+
promptStartedAt: Date.now() - 275,
66+
promptError: new Error("prompt failed"),
67+
formatError: (err) => String(err),
68+
});
69+
70+
expect(warn).not.toHaveBeenCalled();
71+
expect(triggerInternalHook).toHaveBeenCalledTimes(2);
72+
73+
const recordedCalls = triggerInternalHook.mock.calls as unknown as Array<[HookEventRecord]>;
74+
const thinkingStart = recordedCalls[0]?.[0];
75+
const thinkingEnd = recordedCalls[1]?.[0];
76+
expect(thinkingStart?.type).toBe("agent");
77+
expect(thinkingStart?.action).toBe("thinking:start");
78+
expect(thinkingStart?.sessionKey).toBe("run:run-1");
79+
expect(thinkingStart?.context).toMatchObject({
80+
sessionId: "session-1",
81+
runId: "run-1",
82+
provider: "openai",
83+
model: "gpt-5",
84+
messageProvider: "telegram",
85+
messageChannel: "telegram",
86+
});
87+
expect(thinkingEnd?.type).toBe("agent");
88+
expect(thinkingEnd?.action).toBe("thinking:end");
89+
expect(thinkingEnd?.context).toMatchObject({
90+
sessionId: "session-1",
91+
runId: "run-1",
92+
provider: "openai",
93+
model: "gpt-5",
94+
durationMs: 275,
95+
error: "Error: prompt failed",
96+
});
97+
});
98+
99+
it("emits response hooks only once when assistant output begins", async () => {
100+
const warn = vi.fn();
101+
const onAssistantMessageStart = vi.fn();
102+
const emitInternalAgentHook = createInternalAgentHookEmitter(
103+
{
104+
sessionKey: "agent:main:session-1",
105+
runId: "run-2",
106+
},
107+
warn,
108+
);
109+
const tracker = createResponseLifecycleTracker({
110+
params: {
111+
sessionId: "session-1",
112+
runId: "run-2",
113+
provider: "openai",
114+
modelId: "gpt-5",
115+
messageProvider: "discord",
116+
messageChannel: "discord",
117+
},
118+
emitInternalAgentHook,
119+
onAssistantMessageStart,
120+
formatError: (err) => String(err),
121+
});
122+
123+
await tracker.emitResponseStartIfNeeded(false);
124+
expect(triggerInternalHook).not.toHaveBeenCalled();
125+
126+
tracker.handleAssistantMessageStart();
127+
vi.advanceTimersByTime(120);
128+
await tracker.emitResponseStartIfNeeded(true);
129+
await tracker.emitResponseEnd(undefined);
130+
131+
expect(warn).not.toHaveBeenCalled();
132+
expect(onAssistantMessageStart).toHaveBeenCalledTimes(1);
133+
expect(triggerInternalHook).toHaveBeenCalledTimes(2);
134+
const recordedCalls = triggerInternalHook.mock.calls as unknown as Array<[HookEventRecord]>;
135+
const responseStart = recordedCalls[0]?.[0];
136+
const responseEnd = recordedCalls[1]?.[0];
137+
138+
expect(responseStart).toMatchObject({
139+
type: "agent",
140+
action: "response:start",
141+
sessionKey: "agent:main:session-1",
142+
context: {
143+
sessionId: "session-1",
144+
runId: "run-2",
145+
provider: "openai",
146+
model: "gpt-5",
147+
messageProvider: "discord",
148+
messageChannel: "discord",
149+
},
150+
});
151+
expect(responseEnd).toMatchObject({
152+
type: "agent",
153+
action: "response:end",
154+
context: {
155+
sessionId: "session-1",
156+
runId: "run-2",
157+
provider: "openai",
158+
model: "gpt-5",
159+
durationMs: 120,
160+
error: undefined,
161+
},
162+
});
163+
});
164+
});

‎src/agents/pi-embedded-runner/run/lifecycle-hooks.ts‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import { createInternalHookEvent, triggerInternalHook } from "../../../hooks/internal-hooks.js";
2+
import type { EmbeddedRunAttemptParams } from "./types.js";
3+
4+
type LifecycleParams = Pick<
5+
EmbeddedRunAttemptParams,
6+
"sessionId" | "runId" | "provider" | "modelId" | "messageProvider" | "messageChannel"
7+
>;
8+
9+
type EmitInternalAgentHook = (action: string, context: Record<string, unknown>) => Promise<void>;
10+
11+
type ResponseLifecycleTrackerOptions = {
12+
params: LifecycleParams;
13+
emitInternalAgentHook: EmitInternalAgentHook;
14+
onAssistantMessageStart?: () => void;
15+
formatError: (err: unknown) => string;
16+
};
17+
18+
function buildLifecycleContext(params: LifecycleParams): Record<string, unknown> {
19+
return {
20+
sessionId: params.sessionId,
21+
runId: params.runId,
22+
provider: params.provider,
23+
model: params.modelId,
24+
messageProvider: params.messageProvider,
25+
messageChannel: params.messageChannel,
26+
};
27+
}
28+
29+
export function createInternalAgentHookEmitter(
30+
params: Pick<EmbeddedRunAttemptParams, "sessionKey" | "runId">,
31+
warn: (message: string) => void,
32+
): EmitInternalAgentHook {
33+
const hookSessionKey = params.sessionKey?.trim() || `run:${params.runId}`;
34+
35+
return async (action: string, context: Record<string, unknown>) => {
36+
try {
37+
const hookEvent = createInternalHookEvent("agent", action, hookSessionKey, context);
38+
await triggerInternalHook(hookEvent);
39+
} catch (err) {
40+
warn(`agent:${action} hook failed: ${String(err)}`);
41+
}
42+
};
43+
}
44+
45+
export async function emitThinkingStart(
46+
params: LifecycleParams,
47+
emitInternalAgentHook: EmitInternalAgentHook,
48+
): Promise<void> {
49+
await emitInternalAgentHook("thinking:start", buildLifecycleContext(params));
50+
}
51+
52+
export async function emitThinkingEnd(params: {
53+
lifecycleParams: LifecycleParams;
54+
emitInternalAgentHook: EmitInternalAgentHook;
55+
promptStartedAt: number;
56+
promptError: unknown;
57+
formatError: (err: unknown) => string;
58+
}): Promise<void> {
59+
await params.emitInternalAgentHook("thinking:end", {
60+
...buildLifecycleContext(params.lifecycleParams),
61+
durationMs: Date.now() - params.promptStartedAt,
62+
error: params.promptError ? params.formatError(params.promptError) : undefined,
63+
});
64+
}
65+
66+
export function createResponseLifecycleTracker(options: ResponseLifecycleTrackerOptions): {
67+
handleAssistantMessageStart: () => void;
68+
emitResponseStartIfNeeded: (hasResponseOutput: boolean) => Promise<void>;
69+
emitResponseEnd: (promptError: unknown) => Promise<void>;
70+
} {
71+
let responseStartedAt: number | null = null;
72+
73+
const emitResponseStart = async (startedAt: number) => {
74+
if (responseStartedAt !== null) {
75+
return;
76+
}
77+
responseStartedAt = startedAt;
78+
await options.emitInternalAgentHook("response:start", buildLifecycleContext(options.params));
79+
};
80+
81+
return {
82+
handleAssistantMessageStart: () => {
83+
void emitResponseStart(Date.now());
84+
void options.onAssistantMessageStart?.();
85+
},
86+
emitResponseStartIfNeeded: async (hasResponseOutput: boolean) => {
87+
if (!hasResponseOutput) {
88+
return;
89+
}
90+
await emitResponseStart(Date.now());
91+
},
92+
emitResponseEnd: async (promptError: unknown) => {
93+
if (responseStartedAt === null) {
94+
return;
95+
}
96+
await options.emitInternalAgentHook("response:end", {
97+
...buildLifecycleContext(options.params),
98+
durationMs: Date.now() - responseStartedAt,
99+
error: promptError ? options.formatError(promptError) : undefined,
100+
});
101+
},
102+
};
103+
}

‎src/agents/pi-embedded-runner/run/types.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,5 +62,5 @@ export type EmbeddedRunAttemptResult = {
6262
attemptUsage?: NormalizedUsage;
6363
compactionCount?: number;
6464
/** Client tool call detected (OpenResponses hosted tools). */
65-
clientToolCall?: { name: string; params: Record<string, unknown> };
65+
clientToolCall?: { name: string; params: Record<string, unknown>; toolCallId: string };
6666
};

0 commit comments

Comments
 (0)