Merge branch 'main' into fix/web-design-parity

Author: elkaix

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

@@ -0,0 +1,5 @@
+---
+"@pymodel/pythinker-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,
@@ -18,16 +21,18 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
   declare readonly _serviceBrand: undefined;
 
   private lastPersisted: SessionTurnOutcome | undefined;
+  private lastPersistedTurnId: number | undefined;
   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,24 +57,33 @@ 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') {
-          this.write('completed');
+          this.write('completed', { turnId: event.turnId });
           return;
         }
         if (event.reason === 'failed' || event.reason === 'blocked') {
-          this.write('failed');
+          this.write('failed', { turnId: event.turnId });
           return;
         }
         if (event.reason === 'cancelled' && event.interruptReason === 'user_cancelled') {
-          this.write('cancelled');
+          this.write('cancelled', { turnId: event.turnId });
         }
       }),
     );
@@ -79,31 +93,67 @@ export class SessionOutcomeMirror extends Disposable implements ISessionOutcomeM
         this.write(undefined);
       }),
     );
+    subscription.add(
+      bus.subscribe(ContextUndone, (event) => {
+        if (
+          event.fromTurnId !== undefined &&
+          this.lastPersistedTurnId !== undefined &&
+          this.lastPersistedTurnId < event.fromTurnId
+        ) {
+          return;
+        }
+        this.write(undefined);
+      }),
+    );
     subscription.add(
       bus.subscribe(AgentActivityUpdated, (event) => {
         if (this.turnStartedHere) return;
         if (this.lastPersisted !== undefined) return;
         const reason = event.lastTurn?.reason;
         if (reason === 'completed' || reason === 'cancelled') {
-          this.write(reason, { touchUpdatedAt: false });
+          this.write(reason, { touchUpdatedAt: false, turnId: event.lastTurn?.turnId });
         } else if (reason === 'failed' || reason === 'blocked') {
-          this.write('failed', { touchUpdatedAt: false });
+          this.write('failed', { touchUpdatedAt: false, turnId: event.lastTurn?.turnId });
         }
       }),
     );
   }
 
+  private async reconcileAfterRestore(agentStates: IAgentStateService): Promise<void> {
+    await this.metadataReady;
+    if (this.lastPersisted === undefined) return;
+    if (this.turnStartedHere) return;
+    if (!agentStates.has(turnKey)) return;
+    const lastEnded = agentStates.get(turnKey).lastEnded;
+    if (lastEnded === undefined) {
+      this.write(undefined, { touchUpdatedAt: false });
+      return;
+    }
+    if (this.lastPersistedTurnId === undefined) this.lastPersistedTurnId = lastEnded.turnId;
+  }
+
   private write(
     outcome: SessionTurnOutcome | undefined,
-    opts?: { readonly touchUpdatedAt?: boolean },
+    opts?: { readonly touchUpdatedAt?: boolean; readonly turnId?: number },
   ): void {
-    if (outcome === this.lastPersisted) return;
+    if (outcome === this.lastPersisted) {
+      if (opts?.turnId !== undefined) this.lastPersistedTurnId = opts.turnId;
+      return;
+    }
     this.adopted = true;
     const previous = this.lastPersisted;
+    const previousTurnId = this.lastPersistedTurnId;
     this.lastPersisted = outcome;
-    void this.metadata.update({ lastTurnReason: outcome }, opts).catch(() => {
-      if (this.lastPersisted === outcome) this.lastPersisted = previous;
-    });
+    this.lastPersistedTurnId =
+      outcome === undefined ? undefined : (opts?.turnId ?? this.lastPersistedTurnId);
+    void this.metadata
+      .update({ lastTurnReason: outcome }, { touchUpdatedAt: opts?.touchUpdatedAt })
+      .catch(() => {
+        if (this.lastPersisted === outcome) {
+          this.lastPersisted = previous;
+          this.lastPersistedTurnId = previousTurnId;
+        }
+      });
   }
 }
 

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();
   });