fix(agent-core-v2): remove interruption reminder on undo of its turn (#4076)

Author: 7Sageer

.changeset/undo-interruption-reminder.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Remove interruption reminders from the conversation when their interrupted turns are undone.

packages/agent-core-v2/src/agent/interruptionReminder/interruptionReminderService.ts

@@ -2,6 +2,7 @@ import { Disposable } from '#/_base/di/lifecycle';
 import { LifecycleScope } from '#/app/scopes';
 import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
 import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
+import { isUndoAnchor } from '#/agent/contextMemory/conversationTime';
 import type { ContextMessage } from '#/agent/contextMemory/types';
 import { isVacuousContentPart } from '#/agent/contextMemory/vacuousContent';
 import { TurnEnded } from '#/agent/loop/turnOps';
@@ -35,10 +36,12 @@ export class AgentInterruptionReminderService
     this._register(
       eventBus.subscribe(TurnEnded, (event) => {
         if (event.reason !== 'cancelled' || event.interruptReason !== 'user_cancelled') return;
-        const origin = lastComparableMessage(this.context.get())?.origin;
+        const history = this.context.get();
+        const origin = lastComparableMessage(history)?.origin;
         if (origin?.kind === 'injection' && origin.variant === INTERRUPTION_REMINDER_VARIANT) return;
         this.reminder.notify(INTERRUPTION_REMINDER, {
           variant: INTERRUPTION_REMINDER_VARIANT,
+          ownerPromptId: history.findLast(isUndoAnchor)?.id,
         });
       }),
     );

packages/agent-core-v2/test/agent/loop/loop.test.ts

@@ -1981,22 +1981,56 @@ describe('interruption reminder', () => {
 
     await ctx.undoHistory(1);
 
-    expect(
-      ctx.contextData().history.map((message) => ({
-        role: message.role,
-        origin: message.origin,
-      })),
-    ).toEqual([
-      {
-        role: 'user',
-        origin: { kind: 'injection', variant: 'interruption', ownerPromptId: undefined },
-      },
-    ]);
+    expect(ctx.contextData().history).toEqual([]);
+    await ctx.expectResumeMatches();
 
     ctx.mockNextResponse({ type: 'text', text: 'second answer' });
     await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Next' }] });
     await ctx.untilTurnEnd();
+    expect(interruptionReminders()).toHaveLength(0);
+  });
+
+  it('keeps an interruption reminder when undo only removes a later turn', async () => {
+    ctx.mockNextResponse({ type: 'text', text: 'partial answer' });
+    const subscription = cancelOnFirstDelta();
+    await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Hello' }] });
+    await ctx.untilTurnEnd();
+    subscription.dispose();
+    const interruptedHistory = ctx.contextData().history;
+
+    ctx.mockNextResponse({ type: 'text', text: 'second answer' });
+    await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Next' }] });
+    await ctx.untilTurnEnd();
+
+    await ctx.undoHistory(1);
+
+    expect(ctx.contextData().history).toEqual(interruptedHistory);
+    expect(interruptionReminders()).toHaveLength(1);
+    await ctx.expectResumeMatches();
+
+    await ctx.undoHistory(1);
+
+    expect(ctx.contextData().history).toEqual([]);
+  });
+
+  it('undo removes a cancelled retry reminder with its preceding user prompt', async () => {
+    ctx.mockNextResponse({ type: 'text', text: 'first answer' });
+    await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Hello' }] });
+    await ctx.untilTurnEnd();
+
+    ctx.mockNextResponse({ type: 'text', text: 'partial retry' });
+    const subscription = cancelOnFirstDelta();
+    const retry = submitPromptTurn(loop, {
+      message: { role: 'user', content: [] },
+      meta: { origin: { kind: 'retry' } },
+    }).turn;
+    await expect(retry.result).resolves.toMatchObject({ type: 'cancelled' });
+    subscription.dispose();
     expect(interruptionReminders()).toHaveLength(1);
+
+    await ctx.undoHistory(1);
+
+    expect(ctx.contextData().history).toEqual([]);
   });
 
   it('drops unsigned thinking but keeps signed thinking on user cancel', async () => {