fix(context-engine): forward isHeartbeat to afterTurn (fixes #89302) (#90632)

Merged via squash.

Prepared head SHA: 2f6da84c4b
Co-authored-by: zenglingbiao <290951975+zenglingbiao@users.noreply.github.com>
Co-authored-by: jalehman <550978+jalehman@users.noreply.github.com>
Reviewed-by: @jalehman
This commit is contained in:
zenglingbiao
2026-06-06 02:27:37 +08:00
committed by GitHub
parent b3eba2ff38
commit d896a4c7a3
8 changed files with 70 additions and 2 deletions

View File

@@ -2392,6 +2392,7 @@ export async function runCodexAppServerAttempt(
runMaintenance: runHarnessContextEngineMaintenance,
config: params.config,
warn: (message) => embeddedAgentLog.warn(message),
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
});
}
runAgentHarnessLlmOutputHook({

View File

@@ -175,6 +175,7 @@ describe("runPreparedCliAgent context engine lifecycle", () => {
const dispose = vi.fn(async () => {});
const contextEngine = createContextEngine({ bootstrap, afterTurn, maintain, dispose });
const context = buildPreparedContext(contextEngine);
context.params.bootstrapContextRunKind = "heartbeat";
const result = await runPreparedCliAgent(context);
expect(result.meta.agentMeta?.sessionId).toBe("external-cli-session-1");
@@ -198,6 +199,7 @@ describe("runPreparedCliAgent context engine lifecycle", () => {
sessionKey: "agent:main:main",
sessionFile: "session.jsonl",
prePromptMessageCount: 2,
isHeartbeat: true,
tokenBudget: undefined,
runtimeContext: undefined,
});

View File

@@ -270,6 +270,7 @@ async function finalizeCliContextEngineTurn(params: {
sessionIdUsed: runParams.sessionId,
sessionKey: runParams.sessionKey,
sessionFile: runParams.sessionFile,
isHeartbeat: runParams.bootstrapContextRunKind === "heartbeat",
messagesSnapshot: [...prePromptMessages, ...turnMessages],
prePromptMessageCount: prePromptMessages.length,
config: context.contextEngineConfig,

View File

@@ -2406,6 +2406,7 @@ export async function runEmbeddedAttempt(
}),
}),
}),
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
});
const removeGuard = installToolResultContextGuard({
agent: activeSession.agent,
@@ -4691,6 +4692,7 @@ export async function runEmbeddedAttempt(
sessionManager: activeSessionManager,
config: params.config,
warn: (message) => log.warn(message),
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
});
}

View File

@@ -509,6 +509,7 @@ describe("installContextEngineLoopHook", () => {
prePromptMessageCount: number;
}) => Record<string, unknown> | undefined,
onAfterTurnCheckpoint?: (messageCount: number) => void,
isHeartbeat?: boolean,
): () => void {
return installContextEngineLoopHook({
agent,
@@ -521,6 +522,7 @@ describe("installContextEngineLoopHook", () => {
...(prePromptCount !== undefined ? { getPrePromptMessageCount: () => prePromptCount } : {}),
...(getRuntimeContext ? { getRuntimeContext } : {}),
...(onAfterTurnCheckpoint ? { onAfterTurnCheckpoint } : {}),
...(isHeartbeat !== undefined ? { isHeartbeat } : {}),
});
}
@@ -984,7 +986,7 @@ describe("installContextEngineLoopHook", () => {
it("ingests new messages in batches when afterTurn is absent", async () => {
const agent = makeGuardableAgent();
const engine = makeMockEngine({ omitAfterTurn: true });
installHook(agent, engine);
installHook(agent, engine, undefined, undefined, undefined, true);
const batch0 = [makeUser("first"), makeToolResult("call_1", "r1")];
await callTransform(agent, batch0);
@@ -1001,7 +1003,9 @@ describe("installContextEngineLoopHook", () => {
throw new Error("expected ingestBatch mock");
}
expect(recordMockArg(ingestBatch).messages).toEqual(batch1.slice(2));
expect(recordMockArg(ingestBatch).isHeartbeat).toBe(true);
expect(recordMockArg(ingestBatch, 1).messages).toEqual(batch2.slice(4));
expect(recordMockArg(ingestBatch, 1).isHeartbeat).toBe(true);
expect(engine.assemble).toHaveBeenCalledTimes(2);
});
@@ -1019,9 +1023,23 @@ describe("installContextEngineLoopHook", () => {
expect(ingestParams?.sessionId).toBe(sessionId);
expect(ingestParams?.sessionKey).toBe(sessionKey);
expect(ingestParams?.message).toBe(toolResult);
expect(ingestParams?.isHeartbeat).toBeUndefined();
expect(engine.assemble).toHaveBeenCalledTimes(1);
});
it("passes heartbeat state through per-message ingest fallbacks", async () => {
const agent = makeGuardableAgent();
const engine = makeMockEngine({ omitAfterTurn: true, omitIngestBatch: true });
installHook(agent, engine, 1, undefined, undefined, true);
const toolResult = makeToolResult("call_1", "r1");
const messages = [makeUser("first"), toolResult];
await callTransform(agent, messages);
expect(engine.ingest).toHaveBeenCalledTimes(1);
expect(recordMockArg(engine.ingest).isHeartbeat).toBe(true);
});
it("falls through to source messages when engine.afterTurn throws", async () => {
const agent = makeGuardableAgent();
const engine = makeMockEngine({

View File

@@ -334,6 +334,8 @@ export function installContextEngineLoopHook(params: {
messages: AgentMessage[];
prePromptMessageCount: number;
}) => ContextEngineRuntimeContext | undefined;
/** True when this turn belongs to a heartbeat run. */
isHeartbeat?: boolean;
}): () => void {
const { contextEngine, sessionId, sessionKey, sessionFile, tokenBudget, modelId } = params;
const mutableAgent = params.agent as GuardableAgentRecord;
@@ -398,6 +400,7 @@ export function installContextEngineLoopHook(params: {
messages: transcriptMessages,
prePromptMessageCount,
}),
isHeartbeat: params.isHeartbeat,
});
} else {
const newMessages = transcriptMessages.slice(prePromptMessageCount);
@@ -407,6 +410,7 @@ export function installContextEngineLoopHook(params: {
sessionId,
sessionKey,
messages: newMessages,
isHeartbeat: params.isHeartbeat,
});
} else {
for (const message of newMessages) {
@@ -414,6 +418,7 @@ export function installContextEngineLoopHook(params: {
sessionId,
sessionKey,
message,
isHeartbeat: params.isHeartbeat,
});
}
}

View File

@@ -209,11 +209,45 @@ describe("harness context engine lifecycle", () => {
runtimeContext: {},
runMaintenance: async () => undefined,
warn: () => {},
isHeartbeat: true,
});
const ingestBatchCalls = (ingestBatch as unknown as { mock: { calls: unknown[][] } }).mock
.calls;
const ingestBatchParams = ingestBatchCalls[0]?.[0] as { messages?: AgentMessage[] } | undefined;
const ingestBatchParams = ingestBatchCalls[0]?.[0] as
| { isHeartbeat?: boolean; messages?: AgentMessage[] }
| undefined;
expect(ingestBatchParams?.messages).toEqual([turnUser, turnAssistant]);
expect(ingestBatchParams?.isHeartbeat).toBe(true);
});
it("forwards heartbeat state to per-message ingest fallbacks", async () => {
const turnUser = textMessage("user", "new ask", 4);
const turnAssistant = textMessage("assistant", "new answer", 6);
const ingest = vi.fn(async () => ({ ingested: true }));
await finalizeHarnessContextEngineTurn({
contextEngine: createContextEngine({ ingest }),
promptError: false,
aborted: false,
yieldAborted: false,
sessionIdUsed: sessionParams.sessionIdUsed,
sessionKey: sessionParams.sessionKey,
sessionFile: sessionParams.sessionFile,
messagesSnapshot: [turnUser, turnAssistant],
prePromptMessageCount: 0,
tokenBudget: 2048,
runtimeContext: {},
runMaintenance: async () => undefined,
warn: () => {},
isHeartbeat: true,
});
const ingestCalls = (ingest as unknown as { mock: { calls: unknown[][] } }).mock.calls;
expect(ingestCalls).toHaveLength(2);
for (const call of ingestCalls) {
const ingestParams = call[0] as { isHeartbeat?: boolean };
expect(ingestParams.isHeartbeat).toBe(true);
}
});
});

View File

@@ -147,6 +147,8 @@ export async function finalizeHarnessContextEngineTurn(params: {
sessionManager?: unknown;
config?: SessionWriteLockAcquireTimeoutConfig;
warn: (message: string) => void;
/** True when this turn belongs to a heartbeat run. */
isHeartbeat?: boolean;
}) {
if (!params.contextEngine) {
return { postTurnFinalizationSucceeded: true };
@@ -168,6 +170,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
prePromptMessageCount: conversationSnapshot.prePromptMessageCount,
tokenBudget: params.tokenBudget,
runtimeContext: params.runtimeContext,
isHeartbeat: params.isHeartbeat,
});
} catch (afterTurnErr) {
postTurnFinalizationSucceeded = false;
@@ -184,6 +187,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
sessionId: params.sessionIdUsed,
sessionKey: params.sessionKey,
messages: newMessages,
isHeartbeat: params.isHeartbeat,
});
} catch (ingestErr) {
postTurnFinalizationSucceeded = false;
@@ -196,6 +200,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
sessionId: params.sessionIdUsed,
sessionKey: params.sessionKey,
message: msg,
isHeartbeat: params.isHeartbeat,
});
} catch (ingestErr) {
postTurnFinalizationSucceeded = false;