fix(agent): delivered checkpoint reminder to active loop

This commit is contained in:
can1357
2026-08-16 01:54:10 +02:00
parent 958a88b25a
commit 65ee481aa2
2 changed files with 20 additions and 25 deletions
@@ -2344,11 +2344,11 @@ export class AgentSession {
/** /**
* Builds the transient checkpoint-active reminder for a successful * Builds the transient checkpoint-active reminder for a successful
* checkpoint tool result, or undefined otherwise. The reminder is appended * checkpoint tool result, or undefined otherwise. The reminder is queued as
* to agent.state synchronously in the message_end handler (before any * steering synchronously in the message_end handler (before any await), so
* await) so the next provider call within the same tool loop sees it, and * the agent loop folds it into the next provider call and persists it through
* is persisted via #persistMessageEnd. Because the entry sits after the * its normal custom-message event. Because the entry sits after the checkpoint
* checkpoint entry, the rewind branch cut drops it from the active path. * entry, the rewind branch cut drops it from the active path.
*/ */
#checkpointActiveReminderFor( #checkpointActiveReminderFor(
message: AgentMessage, message: AgentMessage,
@@ -2527,14 +2527,12 @@ export class AgentSession {
this.agent.appendMessage(interruptedThinkingMessage); this.agent.appendMessage(interruptedThinkingMessage);
} }
// Same pre-await visibility requirement as the interrupted-thinking // message_end listeners are fire-and-forget, and the agent loop runs against
// message: agent-core invokes message_end listeners fire-and-forget, so // a context cloned at prompt start. Appending only to agent.state would not
// the next provider call can start before any awaited session-event // reach the next provider call in this tool loop. Queue the reminder as
// emission or persistence below settles. The checkpoint-active reminder // steering before any await so the loop drains it at the next step boundary;
// must already be in agent.state for that call — append it synchronously // its normal custom-message event persists it after the checkpoint entry,
// here and persist it later (see #persistMessageEnd). It sits after the // allowing the rewind branch cut to drop it from the active path.
// checkpoint entry, so the rewind branch cut
// (branchWithSummary(checkpointEntryId)) drops it from the active path
const checkpointReminder = const checkpointReminder =
event.type === "message_end" && event.message.role === "toolResult" event.type === "message_end" && event.message.role === "toolResult"
? this.#checkpointActiveReminderFor(event.message) ? this.#checkpointActiveReminderFor(event.message)
@@ -2555,7 +2553,7 @@ export class AgentSession {
}; };
this.#pendingRewindReport = undefined; this.#pendingRewindReport = undefined;
this.#lastCompletedRewind = undefined; this.#lastCompletedRewind = undefined;
this.agent.appendMessage(checkpointReminder); this.agent.steer(checkpointReminder);
} }
const messageEndPersistence = const messageEndPersistence =
@@ -2677,15 +2675,6 @@ export class AgentSession {
interruptedThinkingMessage.attribution, interruptedThinkingMessage.attribution,
); );
} }
if (checkpointReminder) {
this.sessionManager.appendCustomMessageEntry(
checkpointReminder.customType,
checkpointReminder.content,
checkpointReminder.display,
checkpointReminder.details,
checkpointReminder.attribution,
);
}
// Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
if (event.message.role === "assistant") { if (event.message.role === "assistant") {
@@ -301,9 +301,10 @@ describe("AgentSession checkpoint rewind branch context", () => {
while ( while (
!session.messages.some( !session.messages.some(
message => message.role === "custom" && message.customType === "checkpoint-active-reminder", message => message.role === "custom" && message.customType === "checkpoint-active-reminder",
) ) ||
mock.calls.length < 2
) { ) {
if (Date.now() > deadline) throw new Error("checkpoint-active-reminder never appeared"); if (Date.now() > deadline) throw new Error("checkpoint reminder/provider call never appeared");
await Bun.sleep(10); await Bun.sleep(10);
} }
const reminder = session.messages.find( const reminder = session.messages.find(
@@ -312,6 +313,11 @@ describe("AgentSession checkpoint rewind branch context", () => {
); );
expect(reminder).toBeDefined(); expect(reminder).toBeDefined();
expect(reminder?.content).toContain("MUST `rewind` before yielding"); expect(reminder?.content).toContain("MUST `rewind` before yielding");
const activeCall = mock.calls[1];
expect(activeCall).toBeDefined();
expect(
activeCall?.context.messages.some(message => messageText(message).includes("Exploration checkpoint active.")),
).toBe(true);
// #checkpointState is set synchronously with the reminder (pre-await), so an // #checkpointState is set synchronously with the reminder (pre-await), so an
// immediate rewind would find an active checkpoint, not "No active checkpoint". // immediate rewind would find an active checkpoint, not "No active checkpoint".
expect(session.getCheckpointState()).toBeDefined(); expect(session.getCheckpointState()).toBeDefined();