Skip to content

Commit 1223768

Browse files
zenglingbiaozenglingbiao
authored
Merge 2935fce into 4752e9a 
2 parents 4752e9a + 2935fce commit 1223768

8 files changed

Lines changed: 70 additions & 2 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
@@ -2392,6 +2392,7 @@ export async function runCodexAppServerAttempt(
23922392
runMaintenance: runHarnessContextEngineMaintenance,
23932393
config: params.config,
23942394
warn: (message) => embeddedAgentLog.warn(message),
2395+
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
23952396
});
23962397
}
23972398
runAgentHarnessLlmOutputHook({

‎src/agents/cli-runner.context-engine.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ describe("runPreparedCliAgent context engine lifecycle", () => {
175175
const dispose = vi.fn(async () => {});
176176
const contextEngine = createContextEngine({ bootstrap, afterTurn, maintain, dispose });
177177
const context = buildPreparedContext(contextEngine);
178+
context.params.bootstrapContextRunKind = "heartbeat";
178179
const result = await runPreparedCliAgent(context);
179180

180181
expect(result.meta.agentMeta?.sessionId).toBe("external-cli-session-1");
@@ -198,6 +199,7 @@ describe("runPreparedCliAgent context engine lifecycle", () => {
198199
sessionKey: "agent:main:main",
199200
sessionFile: "session.jsonl",
200201
prePromptMessageCount: 2,
202+
isHeartbeat: true,
201203
tokenBudget: undefined,
202204
runtimeContext: undefined,
203205
});

‎src/agents/cli-runner.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,7 @@ async function finalizeCliContextEngineTurn(params: {
270270
sessionIdUsed: runParams.sessionId,
271271
sessionKey: runParams.sessionKey,
272272
sessionFile: runParams.sessionFile,
273+
isHeartbeat: runParams.bootstrapContextRunKind === "heartbeat",
273274
messagesSnapshot: [...prePromptMessages, ...turnMessages],
274275
prePromptMessageCount: prePromptMessages.length,
275276
config: context.contextEngineConfig,

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2406,6 +2406,7 @@ export async function runEmbeddedAttempt(
24062406
}),
24072407
}),
24082408
}),
2409+
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
24092410
});
24102411
const removeGuard = installToolResultContextGuard({
24112412
agent: activeSession.agent,
@@ -4691,6 +4692,7 @@ export async function runEmbeddedAttempt(
46914692
sessionManager: activeSessionManager,
46924693
config: params.config,
46934694
warn: (message) => log.warn(message),
4695+
isHeartbeat: params.bootstrapContextRunKind === "heartbeat",
46944696
});
46954697
}
46964698

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -509,6 +509,7 @@ describe("installContextEngineLoopHook", () => {
509509
prePromptMessageCount: number;
510510
}) => Record<string, unknown> | undefined,
511511
onAfterTurnCheckpoint?: (messageCount: number) => void,
512+
isHeartbeat?: boolean,
512513
): () => void {
513514
return installContextEngineLoopHook({
514515
agent,
@@ -521,6 +522,7 @@ describe("installContextEngineLoopHook", () => {
521522
...(prePromptCount !== undefined ? { getPrePromptMessageCount: () => prePromptCount } : {}),
522523
...(getRuntimeContext ? { getRuntimeContext } : {}),
523524
...(onAfterTurnCheckpoint ? { onAfterTurnCheckpoint } : {}),
525+
...(isHeartbeat !== undefined ? { isHeartbeat } : {}),
524526
});
525527
}
526528

@@ -984,7 +986,7 @@ describe("installContextEngineLoopHook", () => {
984986
it("ingests new messages in batches when afterTurn is absent", async () => {
985987
const agent = makeGuardableAgent();
986988
const engine = makeMockEngine({ omitAfterTurn: true });
987-
installHook(agent, engine);
989+
installHook(agent, engine, undefined, undefined, undefined, true);
988990

989991
const batch0 = [makeUser("first"), makeToolResult("call_1", "r1")];
990992
await callTransform(agent, batch0);
@@ -1001,7 +1003,9 @@ describe("installContextEngineLoopHook", () => {
10011003
throw new Error("expected ingestBatch mock");
10021004
}
10031005
expect(recordMockArg(ingestBatch).messages).toEqual(batch1.slice(2));
1006+
expect(recordMockArg(ingestBatch).isHeartbeat).toBe(true);
10041007
expect(recordMockArg(ingestBatch, 1).messages).toEqual(batch2.slice(4));
1008+
expect(recordMockArg(ingestBatch, 1).isHeartbeat).toBe(true);
10051009
expect(engine.assemble).toHaveBeenCalledTimes(2);
10061010
});
10071011

@@ -1019,9 +1023,23 @@ describe("installContextEngineLoopHook", () => {
10191023
expect(ingestParams?.sessionId).toBe(sessionId);
10201024
expect(ingestParams?.sessionKey).toBe(sessionKey);
10211025
expect(ingestParams?.message).toBe(toolResult);
1026+
expect(ingestParams?.isHeartbeat).toBeUndefined();
10221027
expect(engine.assemble).toHaveBeenCalledTimes(1);
10231028
});
10241029

1030+
it("passes heartbeat state through per-message ingest fallbacks", async () => {
1031+
const agent = makeGuardableAgent();
1032+
const engine = makeMockEngine({ omitAfterTurn: true, omitIngestBatch: true });
1033+
installHook(agent, engine, 1, undefined, undefined, true);
1034+
1035+
const toolResult = makeToolResult("call_1", "r1");
1036+
const messages = [makeUser("first"), toolResult];
1037+
await callTransform(agent, messages);
1038+
1039+
expect(engine.ingest).toHaveBeenCalledTimes(1);
1040+
expect(recordMockArg(engine.ingest).isHeartbeat).toBe(true);
1041+
});
1042+
10251043
it("falls through to source messages when engine.afterTurn throws", async () => {
10261044
const agent = makeGuardableAgent();
10271045
const engine = makeMockEngine({

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,8 @@ export function installContextEngineLoopHook(params: {
334334
messages: AgentMessage[];
335335
prePromptMessageCount: number;
336336
}) => ContextEngineRuntimeContext | undefined;
337+
/** True when this turn belongs to a heartbeat run. */
338+
isHeartbeat?: boolean;
337339
}): () => void {
338340
const { contextEngine, sessionId, sessionKey, sessionFile, tokenBudget, modelId } = params;
339341
const mutableAgent = params.agent as GuardableAgentRecord;
@@ -398,6 +400,7 @@ export function installContextEngineLoopHook(params: {
398400
messages: transcriptMessages,
399401
prePromptMessageCount,
400402
}),
403+
isHeartbeat: params.isHeartbeat,
401404
});
402405
} else {
403406
const newMessages = transcriptMessages.slice(prePromptMessageCount);
@@ -407,13 +410,15 @@ export function installContextEngineLoopHook(params: {
407410
sessionId,
408411
sessionKey,
409412
messages: newMessages,
413+
isHeartbeat: params.isHeartbeat,
410414
});
411415
} else {
412416
for (const message of newMessages) {
413417
await contextEngine.ingest({
414418
sessionId,
415419
sessionKey,
416420
message,
421+
isHeartbeat: params.isHeartbeat,
417422
});
418423
}
419424
}

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

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,11 +209,45 @@ describe("harness context engine lifecycle", () => {
209209
runtimeContext: {},
210210
runMaintenance: async () => undefined,
211211
warn: () => {},
212+
isHeartbeat: true,
212213
});
213214

214215
const ingestBatchCalls = (ingestBatch as unknown as { mock: { calls: unknown[][] } }).mock
215216
.calls;
216-
const ingestBatchParams = ingestBatchCalls[0]?.[0] as { messages?: AgentMessage[] } | undefined;
217+
const ingestBatchParams = ingestBatchCalls[0]?.[0] as
218+
| { isHeartbeat?: boolean; messages?: AgentMessage[] }
219+
| undefined;
217220
expect(ingestBatchParams?.messages).toEqual([turnUser, turnAssistant]);
221+
expect(ingestBatchParams?.isHeartbeat).toBe(true);
222+
});
223+
224+
it("forwards heartbeat state to per-message ingest fallbacks", async () => {
225+
const turnUser = textMessage("user", "new ask", 4);
226+
const turnAssistant = textMessage("assistant", "new answer", 6);
227+
const ingest = vi.fn(async () => ({ ingested: true }));
228+
229+
await finalizeHarnessContextEngineTurn({
230+
contextEngine: createContextEngine({ ingest }),
231+
promptError: false,
232+
aborted: false,
233+
yieldAborted: false,
234+
sessionIdUsed: sessionParams.sessionIdUsed,
235+
sessionKey: sessionParams.sessionKey,
236+
sessionFile: sessionParams.sessionFile,
237+
messagesSnapshot: [turnUser, turnAssistant],
238+
prePromptMessageCount: 0,
239+
tokenBudget: 2048,
240+
runtimeContext: {},
241+
runMaintenance: async () => undefined,
242+
warn: () => {},
243+
isHeartbeat: true,
244+
});
245+
246+
const ingestCalls = (ingest as unknown as { mock: { calls: unknown[][] } }).mock.calls;
247+
expect(ingestCalls).toHaveLength(2);
248+
for (const call of ingestCalls) {
249+
const ingestParams = call[0] as { isHeartbeat?: boolean };
250+
expect(ingestParams.isHeartbeat).toBe(true);
251+
}
218252
});
219253
});

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,8 @@ export async function finalizeHarnessContextEngineTurn(params: {
147147
sessionManager?: unknown;
148148
config?: SessionWriteLockAcquireTimeoutConfig;
149149
warn: (message: string) => void;
150+
/** True when this turn belongs to a heartbeat run. */
151+
isHeartbeat?: boolean;
150152
}) {
151153
if (!params.contextEngine) {
152154
return { postTurnFinalizationSucceeded: true };
@@ -168,6 +170,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
168170
prePromptMessageCount: conversationSnapshot.prePromptMessageCount,
169171
tokenBudget: params.tokenBudget,
170172
runtimeContext: params.runtimeContext,
173+
isHeartbeat: params.isHeartbeat,
171174
});
172175
} catch (afterTurnErr) {
173176
postTurnFinalizationSucceeded = false;
@@ -184,6 +187,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
184187
sessionId: params.sessionIdUsed,
185188
sessionKey: params.sessionKey,
186189
messages: newMessages,
190+
isHeartbeat: params.isHeartbeat,
187191
});
188192
} catch (ingestErr) {
189193
postTurnFinalizationSucceeded = false;
@@ -196,6 +200,7 @@ export async function finalizeHarnessContextEngineTurn(params: {
196200
sessionId: params.sessionIdUsed,
197201
sessionKey: params.sessionKey,
198202
message: msg,
203+
isHeartbeat: params.isHeartbeat,
199204
});
200205
} catch (ingestErr) {
201206
postTurnFinalizationSucceeded = false;

0 commit comments

Comments
 (0)