diff --git a/extensions/codex/src/app-server/run-attempt.ts b/extensions/codex/src/app-server/run-attempt.ts index 07e9f944964f..aa1851d8de79 100644 --- a/extensions/codex/src/app-server/run-attempt.ts +++ b/extensions/codex/src/app-server/run-attempt.ts @@ -2392,6 +2392,7 @@ export async function runCodexAppServerAttempt( runMaintenance: runHarnessContextEngineMaintenance, config: params.config, warn: (message) => embeddedAgentLog.warn(message), + isHeartbeat: params.bootstrapContextRunKind === "heartbeat", }); } runAgentHarnessLlmOutputHook({ diff --git a/src/agents/cli-runner.context-engine.test.ts b/src/agents/cli-runner.context-engine.test.ts index 78293e98930c..6782f31606a8 100644 --- a/src/agents/cli-runner.context-engine.test.ts +++ b/src/agents/cli-runner.context-engine.test.ts @@ -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, }); diff --git a/src/agents/cli-runner.ts b/src/agents/cli-runner.ts index 72aa5e1252a3..af22d2375dfe 100644 --- a/src/agents/cli-runner.ts +++ b/src/agents/cli-runner.ts @@ -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, diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts index aa07837851b4..d9b211bdff89 100644 --- a/src/agents/embedded-agent-runner/run/attempt.ts +++ b/src/agents/embedded-agent-runner/run/attempt.ts @@ -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", }); } 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 74c5effb1628..fb61928eb509 100644 --- a/src/agents/embedded-agent-runner/tool-result-context-guard.ts +++ b/src/agents/embedded-agent-runner/tool-result-context-guard.ts @@ -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, }); } } 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 2deff0233e42..05c747e95723 100644 --- a/src/agents/harness/context-engine-lifecycle.ts +++ b/src/agents/harness/context-engine-lifecycle.ts @@ -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;