From 599580cd964decfcbefb8668bbb1055a0bf15998 Mon Sep 17 00:00:00 2001 From: Josh Lehman Date: Fri, 5 Jun 2026 09:14:49 -0700 Subject: [PATCH] fix(context-engine): forward heartbeat to ingest fallbacks --- .../tool-result-context-guard.test.ts | 20 ++++++++++- .../tool-result-context-guard.ts | 2 ++ .../harness/context-engine-lifecycle.test.ts | 36 ++++++++++++++++++- .../harness/context-engine-lifecycle.ts | 2 ++ 4 files changed, 58 insertions(+), 2 deletions(-) diff --git a/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts b/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts index edcd47ba7c0a..5aa6479e39f9 100644 --- a/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts +++ b/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts @@ -509,6 +509,7 @@ describe("installContextEngineLoopHook", () => { prePromptMessageCount: number; }) => Record | 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({ diff --git a/src/agents/embedded-agent-runner/tool-result-context-guard.ts b/src/agents/embedded-agent-runner/tool-result-context-guard.ts index 2e6e98b10c27..fb61928eb509 100644 --- a/src/agents/embedded-agent-runner/tool-result-context-guard.ts +++ b/src/agents/embedded-agent-runner/tool-result-context-guard.ts @@ -410,6 +410,7 @@ export function installContextEngineLoopHook(params: { sessionId, sessionKey, messages: newMessages, + isHeartbeat: params.isHeartbeat, }); } else { for (const message of newMessages) { @@ -417,6 +418,7 @@ export function installContextEngineLoopHook(params: { sessionId, sessionKey, message, + isHeartbeat: params.isHeartbeat, }); } } diff --git a/src/agents/harness/context-engine-lifecycle.test.ts b/src/agents/harness/context-engine-lifecycle.test.ts index f6d70f540fcc..723e3b46819b 100644 --- a/src/agents/harness/context-engine-lifecycle.test.ts +++ b/src/agents/harness/context-engine-lifecycle.test.ts @@ -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); + } }); }); diff --git a/src/agents/harness/context-engine-lifecycle.ts b/src/agents/harness/context-engine-lifecycle.ts index e6811590caa2..05c747e95723 100644 --- a/src/agents/harness/context-engine-lifecycle.ts +++ b/src/agents/harness/context-engine-lifecycle.ts @@ -187,6 +187,7 @@ export async function finalizeHarnessContextEngineTurn(params: { sessionId: params.sessionIdUsed, sessionKey: params.sessionKey, messages: newMessages, + isHeartbeat: params.isHeartbeat, }); } catch (ingestErr) { postTurnFinalizationSucceeded = false; @@ -199,6 +200,7 @@ export async function finalizeHarnessContextEngineTurn(params: { sessionId: params.sessionIdUsed, sessionKey: params.sessionKey, message: msg, + isHeartbeat: params.isHeartbeat, }); } catch (ingestErr) { postTurnFinalizationSucceeded = false;