fix(agent-core-v2): clear the turn outcome when an undo rewinds the turn it describes

.changeset/undo-clears-turn-outcome.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix the "manually stopped" state lingering after undoing the interrupted turn.

packages/agent-core-v2/src/agent/activityView/activityViewService.ts

@@ -34,6 +34,7 @@ import { USER_PROMPT_ORIGIN } from '#/agent/contextMemory/types';
 import type { PromptOrigin } from '#/agent/contextMemory/types';
 import type { TurnEndReason } from '#/agent/loop/turnEvents';
 import { IEventDispatcher } from '#/state/eventDispatcher';
+import { ContextUndone } from '#/agent/undo/undoService';
 
 import type {
   ActivityLastTurnState,
@@ -140,6 +141,9 @@ export class AgentActivityView extends Disposable implements IAgentActivityView
     this._register(
       this.eventBus.subscribe(TurnEnded, (e) => this.onTurnEnded(e.turnId, e.reason)),
     );
+    this._register(
+      this.eventBus.subscribe(ContextUndone, (e) => this.onContextUndone(e.fromTurnId)),
+    );
     this._register(
       this.eventBus.subscribe(PermissionApprovalRequested, (e) =>
         this.onApprovalRequested(e.id ?? e.toolCallId, e.toolCallId),
@@ -296,6 +300,14 @@ export class AgentActivityView extends Disposable implements IAgentActivityView
     this.publish();
   }
 
+  private onContextUndone(fromTurnId: number | undefined): void {
+    const last = this.lastTurn;
+    if (last === undefined) return;
+    if (fromTurnId !== undefined && last.turnId < fromTurnId) return;
+    this.lastTurn = undefined;
+    this.publish();
+  }
+
   private onStepStarted(step: number): void {
     this.mutateTurn((t) => {
       t.step = step;

packages/agent-core-v2/src/agent/loop/turnOps.ts

@@ -139,10 +139,19 @@ export const turnKey = defineState(
     return { ...next, anchorTurnIds: [...s.anchorTurnIds, s.nextTurnId] };
   })
   .on(TurnSteer, () => {})
-  .on(ContextUndo, (s, e) => ({
-    ...s,
-    anchorTurnIds: s.anchorTurnIds.slice(0, Math.max(0, s.anchorTurnIds.length - e.count)),
-  }))
+  .on(ContextUndo, (s, e) => {
+    const firstRemoved = s.anchorTurnIds[s.anchorTurnIds.length - e.count];
+    const lastEnded = s.lastEnded;
+    return {
+      ...s,
+      anchorTurnIds: s.anchorTurnIds.slice(0, Math.max(0, s.anchorTurnIds.length - e.count)),
+      lastEnded:
+        lastEnded !== undefined &&
+        (firstRemoved === undefined || lastEnded.turnId >= firstRemoved)
+          ? undefined
+          : lastEnded,
+    };
+  })
   .on(ContextApplyCompaction, (s) => ({ ...s, anchorTurnIds: [] }))
   .on(ContextClear, (s) => ({ ...s, anchorTurnIds: [] }))
   .on(TurnCancel, (s, e) => {

packages/agent-core-v2/src/session/sessionActivity/sessionOutcomeMirrorService.ts

@@ -4,7 +4,10 @@ import { LifecycleScope } from '#/app/scopes';
 import { IEventBus } from '#/app/event/eventBus';
 import { AgentActivityUpdated } from '#/agent/activityView/activityView';
 import { TurnStarted } from '#/agent/loop/turnEvents';
-import { TurnEnded } from '#/agent/loop/turnOps';
+import { TurnEnded, turnKey } from '#/agent/loop/turnOps';
+import { ContextUndone } from '#/agent/undo/undoService';
+import { IAgentStateService } from '#/agent/state/agentState';
+import { IEventDispatcher } from '#/state/eventDispatcher';
 import {
   IAgentLifecycleService,
   MAIN_AGENT_ID,
@@ -21,13 +24,14 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
   private adopted = false;
   private turnStartedHere = false;
   private mainSubscription: DisposableStore | undefined;
+  private readonly metadataReady: Promise<void>;
 
   constructor(
     @IAgentLifecycleService private readonly agents: IAgentLifecycleService,
     @ISessionMetadata private readonly metadata: ISessionMetadata,
   ) {
     super();
-    void this.metadata
+    this.metadataReady = this.metadata
       .read()
       .then((meta) => {
         if (!this.adopted) this.lastPersisted = meta.lastTurnReason;
@@ -52,12 +56,21 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
 
   private attachMain(): void {
     if (this.mainSubscription !== undefined) return;
-    const bus = this.agents.handleOf(MAIN_AGENT_ID)?.accessor.get(IEventBus) as
-      | IEventBus
-      | undefined;
+    const handle = this.agents.handleOf(MAIN_AGENT_ID);
+    const bus = handle?.accessor.get(IEventBus) as IEventBus | undefined;
     if (bus === undefined) return;
     const subscription = new DisposableStore();
     this.mainSubscription = subscription;
+    const dispatcher = handle?.accessor.get(IEventDispatcher) as IEventDispatcher | undefined;
+    const agentStates = handle?.accessor.get(IAgentStateService) as IAgentStateService | undefined;
+    if (dispatcher !== undefined && agentStates !== undefined) {
+      subscription.add(
+        dispatcher.hooks.onDidRestore.register('session-outcome-mirror', async (_ctx, next) => {
+          await next();
+          await this.reconcileAfterRestore(agentStates);
+        }),
+      );
+    }
     subscription.add(
       bus.subscribe(TurnEnded, (event) => {
         if (event.reason === 'completed') {
@@ -79,6 +92,11 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
         this.write(undefined);
       }),
     );
+    subscription.add(
+      bus.subscribe(ContextUndone, () => {
+        this.write(undefined);
+      }),
+    );
     subscription.add(
       bus.subscribe(AgentActivityUpdated, (event) => {
         if (this.turnStartedHere) return;
@@ -93,6 +111,15 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
     );
   }
 
+  private async reconcileAfterRestore(agentStates: IAgentStateService): Promise<void> {
+    await this.metadataReady;
+    if (this.lastPersisted === undefined) return;
+    if (this.turnStartedHere) return;
+    if (!agentStates.has(turnKey)) return;
+    if (agentStates.get(turnKey).lastEnded !== undefined) return;
+    this.write(undefined, { touchUpdatedAt: false });
+  }
+
   private write(
     outcome: SessionTurnOutcome | undefined,
     opts?: { readonly touchUpdatedAt?: boolean },

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

@@ -26,6 +26,7 @@ import {
 } from '#/agent/toolApproval/toolApprovalService';
 import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction';
 import type { FullCompactionTask } from '#/agent/fullCompaction/fullCompaction';
+import { ContextUndone } from '#/agent/undo/undoService';
 import { OrderedHookSlot } from '#/hooks';
 import { IEventDispatcher } from '#/state/eventDispatcher';
 import { stubAgentContext } from '../agentContext/stubs';
@@ -234,6 +235,37 @@ describe('AgentActivityView', () => {
     expect(view.state().lastTurn).toMatchObject({ turnId: 2, reason: 'completed' });
   });
 
+  it('clears the last outcome when an undo rewinds the turn it describes', () => {
+    const { bus, view } = harness();
+
+    bus.publish(new TurnStarted({ agentId: 'main', turnId: 1, origin: { kind: 'user' } }));
+    bus.publish(new TurnEnded({ agentId: 'main', turnId: 1, reason: 'cancelled' }));
+    expect(view.state().lastTurn).toMatchObject({ turnId: 1, reason: 'cancelled' });
+
+    bus.publish(new ContextUndone({ agentId: 'main', turns: 1, fromTurnId: 1 }));
+    expect(view.state().lastTurn).toBeUndefined();
+  });
+
+  it('keeps the last outcome when an undo rewinds only later turns', () => {
+    const { bus, view } = harness();
+
+    bus.publish(new TurnStarted({ agentId: 'main', turnId: 1, origin: { kind: 'user' } }));
+    bus.publish(new TurnEnded({ agentId: 'main', turnId: 1, reason: 'completed' }));
+
+    bus.publish(new ContextUndone({ agentId: 'main', turns: 1, fromTurnId: 2 }));
+    expect(view.state().lastTurn).toMatchObject({ turnId: 1, reason: 'completed' });
+  });
+
+  it('clears the last outcome when the undo range cannot be determined', () => {
+    const { bus, view } = harness();
+
+    bus.publish(new TurnStarted({ agentId: 'main', turnId: 1, origin: { kind: 'user' } }));
+    bus.publish(new TurnEnded({ agentId: 'main', turnId: 1, reason: 'failed' }));
+
+    bus.publish(new ContextUndone({ agentId: 'main', turns: 1 }));
+    expect(view.state().lastTurn).toBeUndefined();
+  });
+
   it('exposes the engine-minted interaction id as the approval id', () => {
     const { bus, view } = harness();
 

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

@@ -67,6 +67,36 @@ describe('turnKey lastEnded', () => {
     expect(s.lastEnded?.reason).toBe('completed');
   });
 
+  it('clears the stored outcome when an undo rewinds the turn it describes', () => {
+    let s = turnKey.initial();
+    s = fold(s, new TurnPrompt({ agentId: 'main', input: [], origin: { kind: 'user' } }));
+    s = fold(s, new TurnEnded({ agentId: 'main', turnId: 0, reason: 'completed', durationMs: 10 }));
+    s = fold(s, new TurnPrompt({ agentId: 'main', input: [], origin: { kind: 'user' } }));
+    s = fold(s, new TurnEnded({ agentId: 'main', turnId: 1, reason: 'cancelled', durationMs: 10 }));
+    expect(s.lastEnded?.reason).toBe('cancelled');
+    s = fold(s, new ContextUndo({ agentId: 'main', count: 1 }));
+    expect(s.anchorTurnIds).toEqual([0]);
+    expect(s.lastEnded).toBeUndefined();
+  });
+
+  it('keeps the stored outcome when an undo rewinds only later turns', () => {
+    let s = turnKey.initial();
+    s = fold(s, new TurnPrompt({ agentId: 'main', input: [], origin: { kind: 'user' } }));
+    s = fold(s, new TurnEnded({ agentId: 'main', turnId: 0, reason: 'completed', durationMs: 10 }));
+    s = fold(s, new TurnPrompt({ agentId: 'main', input: [], origin: { kind: 'user' } }));
+    s = fold(s, new ContextUndo({ agentId: 'main', count: 1 }));
+    expect(s.lastEnded).toMatchObject({ turnId: 0, reason: 'completed' });
+  });
+
+  it('clears the stored outcome when the undo count exceeds the tracked anchors', () => {
+    let s = turnKey.initial();
+    s = fold(s, new TurnPrompt({ agentId: 'main', input: [], origin: { kind: 'user' } }));
+    s = fold(s, new TurnEnded({ agentId: 'main', turnId: 0, reason: 'cancelled', durationMs: 10 }));
+    s = fold(s, new ContextUndo({ agentId: 'main', count: 2 }));
+    expect(s.anchorTurnIds).toEqual([]);
+    expect(s.lastEnded).toBeUndefined();
+  });
+
   it('starts without a stored outcome', () => {
     expect(turnKey.initial().lastEnded).toBeUndefined();
   });

packages/agent-core-v2/test/session/sessionActivity/sessionOutcomeMirror.test.ts

@@ -16,7 +16,10 @@ import type { Event2, Event2Class } from '#/app/event/event2';
 import { AgentActivityUpdated } from '#/agent/activityView/activityView';
 import type { AgentContext } from '#/agent/agentContext/agentContext';
 import { TurnStarted } from '#/agent/loop/turnEvents';
-import { TurnEnded } from '#/agent/loop/turnOps';
+import { TurnEnded, turnKey, type TurnModelState } from '#/agent/loop/turnOps';
+import { ContextUndone } from '#/agent/undo/undoService';
+import { IAgentStateService } from '#/agent/state/agentState';
+import { IEventDispatcher } from '#/state/eventDispatcher';
 import type { SessionMeta } from '#/session/sessionMetadata/sessionMetadata';
 import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata';
 import {
@@ -49,6 +52,8 @@ class FakeBus {
 class FakeAgentLifecycle implements IAgentLifecycleService {
   declare readonly _serviceBrand: undefined;
   readonly bus = new FakeBus();
+  readonly restoreHooks: Array<(ctx: undefined, next: () => Promise<void>) => Promise<void>> = [];
+  lastEnded: TurnModelState['lastEnded'];
   private readonly context: AgentContext = stubAgentContext(MAIN_AGENT_ID, 1);
   private readonly createEmitter = new Emitter<AgentContext>();
   private readonly willCloseEmitter = new Emitter<AgentContext>();
@@ -58,10 +63,36 @@ class FakeAgentLifecycle implements IAgentLifecycleService {
   readonly onWillClose = this.willCloseEmitter.event;
   readonly onDidClose = this.didCloseEmitter.event;
   private mainPresent = false;
+  private readonly dispatcher = {
+    hooks: {
+      onDidRestore: {
+        register: (_id: string, fn: (ctx: undefined, next: () => Promise<void>) => Promise<void>) => {
+          this.restoreHooks.push(fn);
+          return { dispose: () => {} };
+        },
+      },
+    },
+  };
+  private readonly agentStates = {
+    has: (key: unknown) => key === turnKey,
+    get: (key: unknown) =>
+      key === turnKey
+        ? { nextTurnId: 1, cancelledTurnIds: [], anchorTurnIds: [], lastEnded: this.lastEnded }
+        : undefined,
+  };
 
   private readonly mainHandle = {
     id: MAIN_AGENT_ID,
-    accessor: { get: (token: unknown) => (token === IEventBus ? this.bus : undefined) },
+    accessor: {
+      get: (token: unknown) =>
+        token === IEventBus
+          ? this.bus
+          : token === IEventDispatcher
+            ? this.dispatcher
+            : token === IAgentStateService
+              ? this.agentStates
+              : undefined,
+    },
   } as unknown as IAgentScopeHandle;
 
   get(agentId: string): AgentContext | undefined {
@@ -264,6 +295,63 @@ describe('SessionOutcomeMirror (Session scope)', () => {
     expect(writes).toEqual(['completed']);
   });
 
+  it('clears the persisted outcome when an undo rewinds the turn', async () => {
+    lifecycle.addMain();
+    await tick();
+    ended('cancelled', 'user_cancelled');
+    expect(writes).toEqual(['cancelled']);
+    lifecycle.bus.publish(new ContextUndone({ agentId: 'main', turns: 1, fromTurnId: 1 }));
+    expect(writes).toEqual(['cancelled', undefined]);
+  });
+
+  it('an undo with no stored outcome writes nothing', async () => {
+    lifecycle.addMain();
+    await tick();
+    lifecycle.bus.publish(new ContextUndone({ agentId: 'main', turns: 1 }));
+    expect(writes).toEqual([]);
+  });
+
+  it('clears a stale persisted outcome when the replayed wire has no ended turn', async () => {
+    const stale = host.child(LifecycleScope.Session, 'session-stale', [
+      stubPair(ISessionMetadata, {
+        read: async () => ({ lastTurnReason: 'cancelled' }) as SessionMeta,
+        update: async (
+          patch: { lastTurnReason?: SessionMeta['lastTurnReason'] },
+          uopts?: { touchUpdatedAt?: boolean },
+        ) => {
+          writes.push(patch.lastTurnReason);
+          touches.push(uopts?.touchUpdatedAt !== false);
+        },
+      } as unknown as ISessionMetadata),
+    ]);
+    const staleLifecycle = stale.accessor.get(IAgentLifecycleService) as unknown as FakeAgentLifecycle;
+    stale.accessor.get(ISessionOutcomeMirror);
+    staleLifecycle.addMain();
+    await tick();
+    expect(writes).toEqual([]);
+    for (const hook of staleLifecycle.restoreHooks) await hook(undefined, async () => {});
+    expect(writes).toEqual([undefined]);
+    expect(touches).toEqual([false]);
+  });
+
+  it('keeps the adopted outcome when the replayed wire still has its ended turn', async () => {
+    const fresh = host.child(LifecycleScope.Session, 'session-fresh', [
+      stubPair(ISessionMetadata, {
+        read: async () => ({ lastTurnReason: 'cancelled' }) as SessionMeta,
+        update: async (patch: { lastTurnReason?: SessionMeta['lastTurnReason'] }) => {
+          writes.push(patch.lastTurnReason);
+        },
+      } as unknown as ISessionMetadata),
+    ]);
+    const freshLifecycle = fresh.accessor.get(IAgentLifecycleService) as unknown as FakeAgentLifecycle;
+    fresh.accessor.get(ISessionOutcomeMirror);
+    freshLifecycle.addMain();
+    freshLifecycle.lastEnded = { turnId: 3, reason: 'cancelled', durationMs: 5 };
+    await tick();
+    for (const hook of freshLifecycle.restoreHooks) await hook(undefined, async () => {});
+    expect(writes).toEqual([]);
+  });
+
   it('reattaches when the main agent is disposed and recreated', async () => {
     lifecycle.addMain();
     await tick();