Skip to content

Commit 2103a25

Browse files
luyao618luyao618
committed
fix(context-engine): forward isHeartbeat to afterTurn hook
The ContextEngine afterTurn() SDK type declares isHeartbeat?: boolean to let plugins distinguish heartbeat turns from user turns, but every call site (harness finalize, embedded-agent loop hook, attempt.ts post finalize) omitted the field. Plugins such as OpenViking that branch on afterTurnParams.isHeartbeat could not reliably detect heartbeat turns and risked polluting context state. Thread isHeartbeat through finalizeHarnessContextEngineTurn and installContextEngineLoopHook. In attempt.ts pull the value from the already-registered AgentRunContext (registered by auto-reply with isHeartbeat) via getAgentRunContext(params.runId). Only emit the field when defined so the absence-vs-false distinction is preserved. Add it.each-parameterized regression tests covering true, false, and the omitted-by-caller cases. Closes #89302 [AI-assisted]
1 parent 0e16e72 commit 2103a25

6 files changed

Lines changed: 108 additions & 0 deletions

File tree

‎extensions/codex/src/app-server/run-attempt.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2344,6 +2344,7 @@ export async function runCodexAppServerAttempt(
23442344
sessionFile: activeSessionFile,
23452345
messagesSnapshot: finalMessages,
23462346
prePromptMessageCount,
2347+
...(params.trigger === "heartbeat" ? { isHeartbeat: true } : {}),
23472348
tokenBudget: params.contextTokenBudget,
23482349
runtimeContext: buildHarnessContextEngineRuntimeContextFromUsage({
23492350
attempt: buildActiveRunAttemptParams(),

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,7 @@ import {
313313
remapInjectedContextFilesToWorkspace,
314314
} from "./attempt.bootstrap-context.js";
315315
export { buildContextEnginePromptCacheInfo } from "./attempt.context-engine-helpers.js";
316+
import { getAgentRunContext } from "../../../infra/agent-events.js";
316317
import {
317318
rotateTranscriptAfterCompaction,
318319
shouldRotateCompactionTranscript,
@@ -2348,6 +2349,7 @@ export async function runEmbeddedAttempt(
23482349
sessionFile: params.sessionFile,
23492350
tokenBudget: params.contextTokenBudget,
23502351
modelId: params.modelId,
2352+
...(params.trigger === "heartbeat" ? { isHeartbeat: true } : {}),
23512353
getPrePromptMessageCount: () => prePromptMessageCount,
23522354
onAfterTurnCheckpoint: (messageCount) => {
23532355
contextEngineAfterTurnCheckpoint = messageCount;
@@ -4536,6 +4538,7 @@ export async function runEmbeddedAttempt(
45364538
prePromptMessageCount: contextEngineAfterTurnCheckpoint ?? prePromptMessageCount,
45374539
tokenBudget: params.contextTokenBudget,
45384540
runtimeContext: afterTurnRuntimeContext,
4541+
...(params.trigger === "heartbeat" ? { isHeartbeat: true } : {}),
45394542
runMaintenance: async (contextParams) =>
45404543
await runContextEngineMaintenance({
45414544
contextEngine: contextParams.contextEngine as never,

‎src/agents/embedded-agent-runner/tool-result-context-guard.test.ts‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1029,4 +1029,56 @@ describe("installContextEngineLoopHook", () => {
10291029
expect(retryResult).toBe(compactedView);
10301030
expect(engine.assemble).toHaveBeenCalledTimes(1);
10311031
});
1032+
1033+
it.each([
1034+
{ label: "forwards isHeartbeat=true to loop hook afterTurn", isHeartbeat: true },
1035+
{ label: "forwards isHeartbeat=false to loop hook afterTurn", isHeartbeat: false },
1036+
])("$label", async ({ isHeartbeat }) => {
1037+
const agent = makeGuardableAgent();
1038+
const engine = makeMockEngine();
1039+
const remove = installContextEngineLoopHook({
1040+
agent,
1041+
contextEngine: engine,
1042+
sessionId,
1043+
sessionKey,
1044+
sessionFile,
1045+
tokenBudget,
1046+
modelId,
1047+
isHeartbeat,
1048+
getPrePromptMessageCount: () => 1,
1049+
});
1050+
await callAfterInitialToolResult(agent);
1051+
remove();
1052+
const afterTurnCalls = (engine.afterTurn as unknown as { mock: { calls: unknown[][] } }).mock
1053+
.calls;
1054+
expect(afterTurnCalls.length).toBeGreaterThan(0);
1055+
for (const call of afterTurnCalls) {
1056+
const callParams = call[0] as { isHeartbeat?: boolean } | undefined;
1057+
expect(callParams?.isHeartbeat).toBe(isHeartbeat);
1058+
}
1059+
});
1060+
1061+
it("omits isHeartbeat from loop hook afterTurn when caller does not provide it", async () => {
1062+
const agent = makeGuardableAgent();
1063+
const engine = makeMockEngine();
1064+
const remove = installContextEngineLoopHook({
1065+
agent,
1066+
contextEngine: engine,
1067+
sessionId,
1068+
sessionKey,
1069+
sessionFile,
1070+
tokenBudget,
1071+
modelId,
1072+
getPrePromptMessageCount: () => 1,
1073+
});
1074+
await callAfterInitialToolResult(agent);
1075+
remove();
1076+
const afterTurnCalls = (engine.afterTurn as unknown as { mock: { calls: unknown[][] } }).mock
1077+
.calls;
1078+
expect(afterTurnCalls.length).toBeGreaterThan(0);
1079+
for (const call of afterTurnCalls) {
1080+
const callParams = call[0] as Record<string, unknown> | undefined;
1081+
expect(callParams && "isHeartbeat" in callParams).toBe(false);
1082+
}
1083+
});
10321084
});

‎src/agents/embedded-agent-runner/tool-result-context-guard.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,8 @@ export function installContextEngineLoopHook(params: {
330330
messages: AgentMessage[];
331331
prePromptMessageCount: number;
332332
}) => ContextEngineRuntimeContext | undefined;
333+
/** True when this run is a heartbeat run. Forwarded to contextEngine.afterTurn. */
334+
isHeartbeat?: boolean;
333335
}): () => void {
334336
const { contextEngine, sessionId, sessionKey, sessionFile, tokenBudget, modelId } = params;
335337
const mutableAgent = params.agent as GuardableAgentRecord;
@@ -394,6 +396,7 @@ export function installContextEngineLoopHook(params: {
394396
messages: transcriptMessages,
395397
prePromptMessageCount,
396398
}),
399+
...(params.isHeartbeat !== undefined ? { isHeartbeat: params.isHeartbeat } : {}),
397400
});
398401
} else {
399402
const newMessages = transcriptMessages.slice(prePromptMessageCount);

‎src/agents/harness/context-engine-lifecycle.test.ts‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,4 +211,50 @@ describe("harness context engine lifecycle", () => {
211211
const ingestBatchParams = ingestBatchCalls[0]?.[0] as { messages?: AgentMessage[] } | undefined;
212212
expect(ingestBatchParams?.messages).toEqual([turnUser, turnAssistant]);
213213
});
214+
215+
it.each([
216+
{ label: "heartbeat run forwards isHeartbeat=true", isHeartbeat: true },
217+
{ label: "non-heartbeat run forwards isHeartbeat=false", isHeartbeat: false },
218+
])("$label", async ({ isHeartbeat }) => {
219+
const userMsg = textMessage("user", "ask", 1);
220+
const asstMsg = textMessage("assistant", "answer", 2);
221+
const afterTurn = vi.fn(async () => {});
222+
await finalizeHarnessContextEngineTurn({
223+
contextEngine: createContextEngine({ afterTurn }),
224+
promptError: false,
225+
aborted: false,
226+
yieldAborted: false,
227+
sessionIdUsed: sessionParams.sessionIdUsed,
228+
sessionKey: sessionParams.sessionKey,
229+
sessionFile: sessionParams.sessionFile,
230+
messagesSnapshot: [userMsg, asstMsg],
231+
prePromptMessageCount: 1,
232+
isHeartbeat,
233+
warn: () => {},
234+
});
235+
const calls = (afterTurn as unknown as { mock: { calls: unknown[][] } }).mock.calls;
236+
const callParams = calls[0]?.[0] as { isHeartbeat?: boolean } | undefined;
237+
expect(callParams?.isHeartbeat).toBe(isHeartbeat);
238+
});
239+
240+
it("omits isHeartbeat from afterTurn payload when caller does not provide it", async () => {
241+
const userMsg = textMessage("user", "ask", 1);
242+
const asstMsg = textMessage("assistant", "answer", 2);
243+
const afterTurn = vi.fn(async () => {});
244+
await finalizeHarnessContextEngineTurn({
245+
contextEngine: createContextEngine({ afterTurn }),
246+
promptError: false,
247+
aborted: false,
248+
yieldAborted: false,
249+
sessionIdUsed: sessionParams.sessionIdUsed,
250+
sessionKey: sessionParams.sessionKey,
251+
sessionFile: sessionParams.sessionFile,
252+
messagesSnapshot: [userMsg, asstMsg],
253+
prePromptMessageCount: 1,
254+
warn: () => {},
255+
});
256+
const calls = (afterTurn as unknown as { mock: { calls: unknown[][] } }).mock.calls;
257+
const callParams = calls[0]?.[0] as Record<string, unknown> | undefined;
258+
expect(callParams && "isHeartbeat" in callParams).toBe(false);
259+
});
214260
});

‎src/agents/harness/context-engine-lifecycle.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,8 @@ export async function finalizeHarnessContextEngineTurn(params: {
140140
prePromptMessageCount: number;
141141
tokenBudget?: number;
142142
runtimeContext?: ContextEngineRuntimeContext;
143+
/** True when this turn belongs to a heartbeat run. Forwarded to contextEngine.afterTurn. */
144+
isHeartbeat?: boolean;
143145
runMaintenance?: typeof runHarnessContextEngineMaintenance;
144146
sessionManager?: unknown;
145147
config?: SessionWriteLockAcquireTimeoutConfig;
@@ -165,6 +167,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
165167
prePromptMessageCount: conversationSnapshot.prePromptMessageCount,
166168
tokenBudget: params.tokenBudget,
167169
runtimeContext: params.runtimeContext,
170+
...(params.isHeartbeat !== undefined ? { isHeartbeat: params.isHeartbeat } : {}),
168171
});
169172
} catch (afterTurnErr) {
170173
postTurnFinalizationSucceeded = false;

0 commit comments

Comments
 (0)