File size: 37,265 Bytes

f0634fb

import { afterEach, describe, expect, it, vi } from 'vitest';

import { type IDisposable } from '#/_base/di/lifecycle';
import {
  resetUnexpectedErrorHandler,
  setUnexpectedErrorHandler,
} from '#/_base/errors/unexpectedError';
import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants';
import { ContextApplyCompaction } from '#/agent/contextMemory/contextEvents';
import { isPromptOwnedInjection, isUndoAnchor } from '#/agent/contextMemory/conversationTime';
import type { ContextMessage, PromptOrigin, TaskOrigin } from '#/agent/contextMemory/types';
import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction';
import { IAgentLoopService } from '#/agent/loop/loop';
import { turnKey } from '#/agent/loop/turnOps';
import { IAgentPlanService } from '#/features/plan/plan';
import { planKey } from '#/features/plan/planOps';
import { IAgentTaskService, type AgentTask } from '#/agent/task/task';
import { taskNotificationDeliveryKey } from '#/agent/task/taskService';
import { IAgentConversationUndoService } from '#/agent/undo/undo';
import { ContextUndone } from '#/agent/undo/undoService';
import { AgentStatusUpdated } from '#/agent/usage/usageEvents';
import { IEventBus } from '#/app/event/eventBus';
import { ErrorCodes } from '#/errors';
import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata';
import { ToolsUpdateStore } from '#/features/todo/todoOps';
import type { TodoItem } from '#/features/todo/todoItem';
import { IAgentTodoService } from '#/features/todo/todoService';
import type { DurableAgentRuntimeParticipant } from '#/state/eventDispatcher';
import { WIRE_PROTOCOL_VERSION } from '#/wire/migration/migration';
import type { WireRecord } from '#/wire/record';
import { IWireService } from '#/wire/wire';

import { createTestAgent, execEnvServices, telemetryServices, InMemoryWireRecordPersistence, type TestAgentContext } from '../../harness';
import { submitPromptTurn } from '../loop/stubs';
import { createFakeHostFs } from '../../tools/fixtures/fake-exec';
import { recordingTelemetry, type TelemetryRecord } from '../../app/telemetry/stubs';
import { ITelemetryService } from '#/app/telemetry/telemetry';

describe('AgentConversationUndoService', () => {
  let ctx: TestAgentContext;
  let records: TelemetryRecord[];

  afterEach(async () => {
    try {
      await ctx.expectResumeMatches();
    } finally {
      await ctx.dispose();
    }
  });

  async function setup() {
    records = [];
    ctx = createTestAgent(
      telemetryServices(recordingTelemetry(records)),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    ctx.get(IAgentContextMemoryService);
    await ctx.restorePersisted();
    return ctx;
  }

  it('runs transcript reconciliation after restored messages are persisted', async () => {
    await setup();
    ctx.appendTurnExchange('kept', 'answer');
    ctx.appendTurnExchange('removed', 'answer');
    const participants = ctx.get(IAgentConversationUndoParticipantRegistry);
    const observed: string[] = [];
    let restored = false;
    let persisted = false;
    const wire = ctx.get(IWireService);
    const originalFlush = wire.flush.bind(wire);
    const flush = vi.spyOn(wire, 'flush').mockImplementation(async () => {
      await originalFlush();
      if (restored) persisted = true;
    });
    participants.register({
      id: 'test.transcript',
      phase: 'after-flush',
      reconcileAfterUndo: async () => {
        observed.push(persisted ? 'flushed' : 'not flushed');
      },
    });
    participants.register({
      id: 'test.notification',
      reconcileAfterUndo: async () => {
        await Promise.resolve();
        ctx.context.append({
          role: 'user',
          content: [{ type: 'text', text: 'restored notification' }],
          toolCalls: [],
          origin: { kind: 'task', taskId: 'task-1', status: 'completed', notificationId: 'notification-1' },
        });
        restored = true;
      },
    });
    await ctx.get(IAgentConversationUndoService).undo(1);
    expect(observed).toEqual(['flushed']);
    flush.mockRestore();
  });

  it('exposes availability from context history', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    expect(undo.availability()).toEqual({ maxTurns: 0, stoppedAtCompaction: false });

    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    expect(undo.availability()).toEqual({ maxTurns: 2, stoppedAtCompaction: false });
  });

  it('rejects undo with structured reasons', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);

    await expect(undo.undo(1)).rejects.toMatchObject({
      code: ErrorCodes.SESSION_UNDO_UNAVAILABLE,
      details: { reason: 'empty', requestedCount: 1, undoableCount: 0 },
    });

    ctx.appendTurnExchange('u1', 'a1');
    await expect(undo.undo(2)).rejects.toMatchObject({
      code: ErrorCodes.SESSION_UNDO_UNAVAILABLE,
      details: { reason: 'insufficient', requestedCount: 2, undoableCount: 1 },
    });
  });

  it.each([
    0,
    -1,
    0.5,
    Number.MAX_SAFE_INTEGER + 1,
    Number.POSITIVE_INFINITY,
    Number.NaN,
  ])('rejects invalid undo count %s without mutating history', async (count) => {
    await setup();
    ctx.appendTurnExchange('u1', 'a1');
    const history = ctx.context.get();

    await expect(ctx.get(IAgentConversationUndoService).undo(count)).rejects.toMatchObject({
      code: ErrorCodes.REQUEST_INVALID,
      details: { field: 'count' },
    });

    expect(ctx.context.get()).toBe(history);
  });

  it('returns session.busy for an active turn without cancelling it', async () => {
    await setup();
    const loop = ctx.get(IAgentLoopService);
    let started!: () => void;
    let release!: () => void;
    const didStart = new Promise<void>((resolve) => {
      started = resolve;
    });
    const canFinish = new Promise<void>((resolve) => {
      release = resolve;
    });
    const hook = loop.hooks.onWillBeginStep.register('test-invalid-undo', async (_hookCtx, next) => {
      started();
      await canFinish;
      await next();
    });
    ctx.mockNextResponse({ type: 'text', text: 'system result' });
    const turn = submitPromptTurn(loop, {
      message: { role: 'user', content: [{ type: 'text', text: 'system work' }] },
      meta: { origin: { kind: 'system_trigger', name: 'test' } as PromptOrigin },
    }).turn;
    await didStart;
    const history = ctx.context.get();

    await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toMatchObject({
      code: ErrorCodes.SESSION_BUSY,
      details: { reason: 'loop' },
    });
    expect(turn.signal.aborted).toBe(false);
    expect(loop.snapshot().state).toBe('running');
    expect(ctx.context.get()).toBe(history);

    hook.dispose();
    release();
    await expect(turn.result).resolves.toMatchObject({ type: 'completed' });
  });

  it('returns session.busy for active compaction without cancelling it', async () => {
    await setup();
    ctx.appendTurnExchange('u1', 'a1');
    const history = ctx.context.get();
    const compaction = ctx.get(IAgentFullCompactionService);
    const abortController = new AbortController();
    const active = vi.spyOn(compaction, 'compacting', 'get').mockReturnValue({
      abortController,
      promise: new Promise<never>(() => {}),
      trigger: 'manual',
      tokenCount: 2,
    });

    try {
      await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toMatchObject({
        code: ErrorCodes.SESSION_BUSY,
        details: { reason: 'compaction' },
      });
      expect(abortController.signal.aborted).toBe(false);
      expect(ctx.context.get()).toBe(history);
    } finally {
      active.mockRestore();
    }
  });

  it('refuses to cross a compaction boundary', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    ctx.appendTurnExchange('u1', 'a1');
    ctx.get(IAgentContextMemoryService).applyCompaction({
      summary: 'summary of u1',
      compactedCount: 2,
      tokensBefore: 100,
      tokensAfter: 10,
    });
    ctx.appendTurnExchange('u2', 'a2');

    expect(undo.availability()).toEqual({ maxTurns: 1, stoppedAtCompaction: true });
    await expect(undo.undo(2)).rejects.toMatchObject({
      code: ErrorCodes.SESSION_UNDO_UNAVAILABLE,
      details: { reason: 'compaction_boundary', requestedCount: 2, undoableCount: 1 },
    });

    await undo.undo(1);
    const history = ctx.context.get();
    expect(history.map((m) => m.role)).toEqual(['user', 'user', 'user']);
    expect(history[1]?.origin?.kind).toBe('compaction_summary');
    expect(history[2]?.origin).toEqual({ kind: 'injection', variant: 'compaction_continuation' });
  });

  it('rejects undo across a legacy compaction boundary even when the in-memory precheck allows it', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    await ctx.dispatcher.dispatch(
      new ContextApplyCompaction({ agentId: 'main', summary: 'legacy summary', compactedCount: 2 }),
    );
    expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'user', 'assistant']);

    await expect(undo.undo(1)).rejects.toMatchObject({
      code: ErrorCodes.SESSION_UNDO_UNAVAILABLE,
      details: { reason: 'compaction_boundary', requestedCount: 1, undoableCount: 0 },
    });
    expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'user', 'assistant']);
  });

  it('cuts before the anchor that survives a legacy unpaired undo', async () => {
    records = [];
    const prompt = (text: string): WireRecord => ({
      type: 'context.append_message',
      agentId: 'main',
      message: {
        role: 'user',
        content: [{ type: 'text', text }],
        toolCalls: [],
        origin: { kind: 'user' },
      },
      time: 1,
    });
    const reply = (text: string): WireRecord => ({
      type: 'context.append_message',
      agentId: 'main',
      message: {
        role: 'assistant',
        content: [{ type: 'text', text }],
        toolCalls: [],
      },
      time: 1,
    });
    const persistence = new InMemoryWireRecordPersistence([
      { type: 'metadata', protocol_version: WIRE_PROTOCOL_VERSION, created_at: 1 },
      prompt('u1'),
      reply('a1'),
      prompt('u2'),
      reply('a2'),
      { type: 'context.undo', agentId: 'main', count: 1, time: 2 },
      prompt('u3'),
      reply('a3'),
    ] as WireRecord[]);
    ctx = createTestAgent(
      { autoConfigure: false, persistence },
      telemetryServices(recordingTelemetry(records)),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    ctx.get(IAgentContextMemoryService);
    await ctx.restorePersisted();
    expect(
      ctx.context.get().map((m) => m.content.map((p) => (p.type === 'text' ? p.text : '')).join('')),
    ).toEqual(['u1', 'a1', 'u3', 'a3']);

    await ctx.get(IAgentConversationUndoService).undo(2);

    expect(ctx.context.get()).toEqual([]);
    const persisted = await ctx.persistedWireRecords();
    const edgeIndex = persisted.findIndex((record) => record.type === 'agent.switched');
    expect(persisted[edgeIndex]).toMatchObject({
      branch: 'b1',
      base: { branch: 'main', line: 1 },
      turns: 2,
      legacyUndoLine: 10,
    });
    expect(persisted[edgeIndex + 1]).toMatchObject({ type: 'context.undo', count: 2 });
    expect(persisted[edgeIndex + 2]).toMatchObject({ type: 'context.undone', turns: 2 });
  });

  it('restores plan mode and its telemetry mirror to their pre-turn value', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    await ctx.get(IAgentPlanService).enter('plan-x', false);
    const restoredModes: boolean[] = [];
    const subscription = ctx.get(IEventBus).subscribe(AgentStatusUpdated, (event) => {
      if (event.planMode !== undefined) restoredModes.push(event.planMode);
    });

    try {
      await undo.undo(1);

      expect(ctx.agentState.get(planKey).active).toBe(false);
      expect(ctx.get(ITelemetryService).getContext().mode).toBe('agent');
      expect(restoredModes).toEqual([false]);
    } finally {
      subscription.dispose();
    }
  });

  it('keeps machine and wire turn ids aligned across undo, a continued turn, and a restart', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);

    const runTurn = async (
      target: TestAgentContext,
      text: string,
    ): Promise<number | undefined> => {
      target.mockNextResponse({ type: 'text', text: `answer to ${text}` });
      const { turn } = submitPromptTurn(target.get(IAgentLoopService), {
        message: { role: 'user', content: [{ type: 'text', text }] },
        meta: { origin: { kind: 'user' } },
      });
      await expect(turn.result).resolves.toMatchObject({ type: 'completed' });
      return turn.id;
    };

    await runTurn(ctx, 'u1');
    await runTurn(ctx, 'u2');
    expect(ctx.agentState.get(turnKey).nextTurnId).toBe(2);

    await undo.undo(1);

    expect(ctx.agentState.get(turnKey).nextTurnId).toBe(2);

    await expect(runTurn(ctx, 'u3')).resolves.toBe(1);

    const persisted = await ctx.persistedWireRecords();
    expect(
      persisted.filter((record) => record.type === 'turn.prompt').map((record) => record['turnId']),
    ).toEqual([0, 1, 1]);
    expect(
      persisted
        .filter((record) => record.type === 'agent.turn.started')
        .map((record) => record['turnId']),
    ).toEqual([0, 1, 1]);

    const resumed = createTestAgent(
      { autoConfigure: false, persistence: new InMemoryWireRecordPersistence(persisted) },
      telemetryServices(recordingTelemetry(records)),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    try {
      resumed.get(IAgentContextMemoryService);
      await resumed.restorePersisted();
      expect(resumed.agentState.get(turnKey).nextTurnId).toBe(2);
      await expect(runTurn(resumed, 'u4')).resolves.toBe(2);
      const repersisted = await resumed.persistedWireRecords();
      expect(
        repersisted
          .filter((record) => record.type === 'agent.turn.started')
          .map((record) => record['turnId']),
      ).toEqual([0, 1, 1, 2]);
    } finally {
      await resumed.dispose();
    }
  });

  it('reports the removed turn id only when context anchors were opened by engine turns', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    const loop = ctx.get(IAgentLoopService);

    ctx.mockNextResponse({ type: 'text', text: 'a1' });
    const userTurn = submitPromptTurn(loop, {
      message: { role: 'user', content: [{ type: 'text', text: 'u1' }] },
      meta: { origin: { kind: 'user' } },
    }).turn;
    await expect(userTurn.result).resolves.toMatchObject({ type: 'completed' });

    ctx.mockNextResponse({ type: 'text', text: 'cron done' });
    const cronTurn = submitPromptTurn(loop, {
      message: { role: 'user', content: [{ type: 'text', text: 'cron work' }] },
      meta: { origin: {
        kind: 'cron_job',
        jobId: 'j1',
        cron: '0 9 * * *',
        recurring: true,
        coalescedCount: 0,
        stale: false,
      } as PromptOrigin },
    }).turn;
    await expect(cronTurn.result).resolves.toMatchObject({ type: 'completed' });

    let fromTurnId: number | undefined;
    const subscription = ctx.get(IEventBus).subscribe(ContextUndone, (event) => {
      fromTurnId = event.fromTurnId;
    });
    try {
      await undo.undo(1);
      expect(fromTurnId).toBe(userTurn.id);
      expect(ctx.agentState.get(turnKey).anchorTurnIds).toEqual([]);
      expect(ctx.context.get()).toHaveLength(0);
    } finally {
      subscription.dispose();
    }

    ctx.get(IAgentContextMemoryService).append(
      {
        role: 'user',
        content: [{ type: 'text', text: 'u2' }],
        toolCalls: [],
        origin: { kind: 'user' },
      },
      {
        role: 'assistant',
        content: [{ type: 'text', text: 'a2' }],
        toolCalls: [],
      },
    );

    let absentTurnId: number | undefined = Number.NaN;
    const second = ctx.get(IEventBus).subscribe(ContextUndone, (event) => {
      absentTurnId = event.fromTurnId;
    });
    try {
      await undo.undo(1);
      expect(absentTurnId).toBeUndefined();
    } finally {
      second.dispose();
    }
  });

  it('flushes state reconciliation before publishing undo', async () => {
    await setup();
    const wire = ctx.get(IWireService);
    const order: string[] = [];
    const flush = vi.spyOn(wire, 'flush');
    const originalFlush = flush.getMockImplementation();
    flush.mockImplementation(async () => {
      order.push('flush');
      await originalFlush?.();
    });
    const participants = ctx.get(IAgentConversationUndoParticipantRegistry);
    participants.register({
      id: 'test.state',
      reconcileAfterUndo: async () => {
        order.push('state');
      },
    });
    const subscription = ctx.get(IEventBus).subscribe(ContextUndone, () => {
      order.push('context.undone');
    });
    ctx.appendTurnExchange('u1', 'a1');

    try {
      await ctx.get(IAgentConversationUndoService).undo(1);

      expect(order).toEqual(['flush', 'flush', 'state', 'flush', 'context.undone']);
    } finally {
      subscription.dispose();
      flush.mockRestore();
    }
  });

  it.each([
    [1, []],
    [3, ['state']],
  ] as const)(
    'rejects the undo when wire flush %i fails',
    async (failureCall, expectedReconciled) => {
      await setup();
      const wire = ctx.get(IWireService);
      const originalFlush = wire.flush.bind(wire);
      let flushCalls = 0;
      const storageError = new Error('storage unavailable');
      const flush = vi.spyOn(wire, 'flush').mockImplementation(async () => {
        flushCalls += 1;
        if (flushCalls === failureCall) throw storageError;
        await originalFlush();
      });
      const originalAppend = wire.appendRecord.bind(wire);
      const appendRecord = vi.spyOn(wire, 'appendRecord');
      if (failureCall === 1) {
        appendRecord.mockImplementation((record, dehydrate) => {
          if (
            record.type === 'agent.switched' ||
            record.type === 'context.undo' ||
            record.type === 'context.undone'
          ) {
            return;
          }
          originalAppend(record, dehydrate);
        });
      }
      const reconciled: string[] = [];
      const participants = ctx.get(IAgentConversationUndoParticipantRegistry);
      participants.register({
        id: 'test.flush-failure-state',
        reconcileAfterUndo: async () => {
          reconciled.push('state');
        },
      });
      const undone: number[] = [];
      const subscription = ctx.get(IEventBus).subscribe(ContextUndone, ({ turns }) => {
        undone.push(turns);
      });
      ctx.appendTurnExchange('u1', 'a1');

      try {
        await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toBe(storageError);
        if (failureCall === 1) {
          expect(ctx.context.get().map((message) => message.role)).toEqual(['user', 'assistant']);
        } else {
          expect(ctx.context.get()).toEqual([]);
        }
        expect(reconciled).toEqual(expectedReconciled);
        expect(undone).toEqual([]);
        expect(records.filter((record) => record.event === 'conversation_undo')).toEqual([]);
      } finally {
        subscription.dispose();
        appendRecord.mockRestore();
        flush.mockRestore();
      }
    },
  );

  it('serializes concurrent undos through state reconciliation', async () => {
    await setup();
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    let releaseFirst!: () => void;
    const firstBlocked = new Promise<void>((resolve) => {
      releaseFirst = resolve;
    });
    let markFirstStarted!: () => void;
    const firstStarted = new Promise<void>((resolve) => {
      markFirstStarted = resolve;
    });
    let calls = 0;
    let active = 0;
    let maxActive = 0;
    ctx.get(IAgentConversationUndoParticipantRegistry).register({
      id: 'test.serial-state',
      reconcileAfterUndo: async () => {
        calls += 1;
        active += 1;
        maxActive = Math.max(maxActive, active);
        if (calls === 1) {
          markFirstStarted();
          await firstBlocked;
        }
        active -= 1;
      },
    });

    const first = ctx.get(IAgentConversationUndoService).undo(1);
    await firstStarted;
    const second = ctx.get(IAgentConversationUndoService).undo(1);
    await Promise.resolve();

    expect(calls).toBe(1);
    expect(ctx.context.get().map((message) => message.role)).toEqual(['user', 'assistant']);
    releaseFirst();
    await Promise.all([first, second]);

    expect(calls).toBe(2);
    expect(maxActive).toBe(1);
    expect(ctx.context.get()).toEqual([]);
  });

  it('publishes context.undone and tracks conversation_undo', async () => {
    await setup();
    ctx.get(IAgentConversationUndoService);
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');

    await ctx.rpc.undoHistory({ count: 1 });

    expect(records).toContainEqual({
      event: 'conversation_undo',
      properties: {
        agent_id: 'main',
        count: 1,
        mode: 'agent',
        model: 'mock-model',
        protocol: 'openai',
        provider_type: 'kimi',
      },
    });
    expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'assistant']);
  });

  it('reconciles lastPrompt after undo', async () => {
    await setup();
    const metadata = ctx.get(ISessionMetadata);
    await metadata.ready;
    await metadata.update({ lastPrompt: 'u1' });
    ctx.appendTurnExchange('u1', 'a1');

    await ctx.get(IAgentConversationUndoService).undo(1);
    await expect(metadata.read()).resolves.toMatchObject({ lastPrompt: undefined });

  });

  it.each([undefined, 'Save button · Rename it'])('uses the newest pending prompt as lastPrompt after undo (display=%s)', async (displayText) => {
    await setup();
    const metadata = ctx.get(ISessionMetadata);
    await metadata.ready;
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    ctx.appendTurnExchange('u3', 'a3');
    const list = vi.spyOn(ctx.get(IAgentLoopService), 'snapshot').mockReturnValue({
      state: 'idle',
      activeTurnId: undefined,
      activePromptId: undefined,
      queue: [
        {
          message: {
            role: 'user',
            content: [{ type: 'text', text: 'queued prompt' }],
          },
          meta: {
            promptId: 'queued',
            origin: { kind: 'user', clientMetadata: displayText === undefined ? undefined : [{ display_text: displayText }] } as PromptOrigin,
            tracked: true,
            createdAt: new Date(0).toISOString(),
            userMessageId: 'queued',
          },
        },
      ],
      notificationCount: 0,
      paused: false,
      hasPendingRequests: true,
      turn: undefined,
      activeTraceId: undefined,
    });

    try {
      await ctx.get(IAgentConversationUndoService).undo(1);
      await expect(metadata.read()).resolves.toMatchObject({ lastPrompt: displayText ?? 'queued prompt' });
    } finally {
      list.mockRestore();
    }
  });

  it('treats metadata reconciliation failure as non-fatal after committing undo', async () => {
    await setup();
    ctx.appendTurnExchange('u1', 'a1');
    ctx.appendTurnExchange('u2', 'a2');
    const update = vi.spyOn(ctx.get(ISessionMetadata), 'update').mockRejectedValueOnce(
      new Error('metadata write failed'),
    );
    const undone: number[] = [];
    const subscription = ctx.get(IEventBus).subscribe(ContextUndone, ({ turns }) => {
      undone.push(turns);
    });

    try {
      await expect(ctx.get(IAgentConversationUndoService).undo(1)).resolves.toBe(1);

      expect(ctx.context.get().map((message) => message.role)).toEqual(['user', 'assistant']);
      expect(undone).toEqual([1]);
      expect(records).toContainEqual({
        event: 'conversation_undo',
        properties: {
          agent_id: 'main',
          count: 1,
          mode: 'agent',
          model: 'mock-model',
          protocol: 'openai',
          provider_type: 'kimi',
        },
      });
    } finally {
      subscription.dispose();
      update.mockRestore();
    }
  });

  it('re-delivers wait-reported task notifications after conversation undo', async () => {
    await setup();
    const undo = ctx.get(IAgentConversationUndoService);
    const tasks = ctx.get(IAgentTaskService);
    ctx.appendTurnExchange('u1', 'a1');

    const completingTask = (output: string): AgentTask => ({
      idPrefix: 'test',
      kind: 'process',
      description: 'fake process task',
      start: async (sink) => {
        sink.appendOutput(output);
        await sink.settle({ status: 'completed' });
      },
      toInfo: (base) => ({ ...base, kind: 'process', command: 'echo', pid: 0, exitCode: null }),
    });

    const taskA = tasks.registerTask(completingTask('a\n'));
    const taskB = tasks.registerTask(completingTask('b\n'));
    tasks.markTasksDeliveredViaWait([
      { taskId: taskA, status: 'completed' },
      { taskId: taskB, status: 'completed' },
    ]);
    await tasks.wait(taskA, 1000);
    await tasks.wait(taskB, 1000);

    expect(ctx.context.get().some((message) => message.origin?.kind === 'task')).toBe(false);
    expect(ctx.agentState.get(taskNotificationDeliveryKey)).toHaveLength(2);

    await undo.undo(1);

    const redelivered = ctx.context.get().filter((message) => message.origin?.kind === 'task');
    expect(redelivered.map((message) => (message.origin as TaskOrigin).taskId).toSorted()).toEqual(
      [taskA, taskB].toSorted(),
    );
  });

  it('registers a participant that late-attaches inside the undo rerun window', async () => {
    await setup();
    const unexpected: unknown[] = [];
    setUnexpectedErrorHandler((error) => unexpected.push(error));
    try {
      const dispatcher = ctx.dispatcher;
      const box: { todos: readonly TodoItem[] } = { todos: [] };
      const folded: TodoItem[][] = [];
      const participant: DurableAgentRuntimeParticipant<{ todos: readonly TodoItem[] }> = {
        id: 'runtime.test.rerun-late',
        events: [ToolsUpdateStore],
        undoable: true,
        transition: (draft, event) => {
          if (event instanceof ToolsUpdateStore && event.key === 'todo') {
            const value = event.value as TodoItem[];
            draft.todos = value;
            folded.push(value);
          }
        },
        getState: () => box,
        commit: (next) => {
          box.todos = next.todos;
        },
      };
      const update = (title: string) =>
        new ToolsUpdateStore({ agentId: 'main', key: 'todo', value: [{ title, status: 'pending' }] });
      await dispatcher.dispatch(update('kept'));
      ctx.appendTurnExchange('u1', 'a1');
      await dispatcher.dispatch(update('doomed'));

      let lateAttach: Promise<IDisposable> | undefined;
      let liveDispatch: Promise<void> | undefined;
      const originalRestore = dispatcher.restore.bind(dispatcher);
      const restoreSpy = vi.spyOn(dispatcher, 'restore').mockImplementation(async () => {
        const restored = originalRestore();
        lateAttach = dispatcher.attachLate(participant);
        liveDispatch = dispatcher.dispatch(update('live'));
        await restored;
      });
      try {
        await ctx.get(IAgentConversationUndoService).undo(1);

        await expect(lateAttach!).resolves.toBeDefined();
        await liveDispatch;
        expect(folded).toEqual([
          [{ title: 'kept', status: 'pending' }],
          [{ title: 'live', status: 'pending' }],
        ]);
        expect(box.todos).toEqual([{ title: 'live', status: 'pending' }]);

        await dispatcher.dispatch(update('after'));
        expect(box.todos).toEqual([{ title: 'after', status: 'pending' }]);
        expect(ctx.get(IAgentTodoService).get()).toEqual([{ title: 'after', status: 'pending' }]);
        expect(
          unexpected.filter((error) => String((error as Error)?.message).includes('late-attached')),
        ).toEqual([]);
      } finally {
        restoreSpy.mockRestore();
      }
    } finally {
      resetUnexpectedErrorHandler();
    }
  });

  it('recovers the undo when the rerun restore fails transiently', async () => {
    await setup();
    ctx.appendTurnExchange('u1', 'a1');
    const wire = ctx.get(IWireService);
    const originalFlush = wire.flush.bind(wire);
    const failure = new Error('transient storage failure');
    let flushCalls = 0;
    const flush = vi.spyOn(wire, 'flush').mockImplementation(async () => {
      flushCalls += 1;
      if (flushCalls === 2) throw failure;
      await originalFlush();
    });

    try {
      await expect(ctx.get(IAgentConversationUndoService).undo(1)).resolves.toBe(1);

      expect(ctx.context.get()).toEqual([]);
      expect(ctx.dispatcher.restorePhase).toBe('ready');
      const persisted = await ctx.persistedWireRecords();
      expect(persisted.filter((record) => record.type === 'agent.switched')).toHaveLength(1);
      expect(records.filter((record) => record.event === 'conversation_undo')).toHaveLength(1);
    } finally {
      flush.mockRestore();
    }
  });

  it('keeps terminal state equivalent across a legacy record-level downgrade round trip, and documents the orphan-edge crash window', async () => {
    records = [];
    const persistence = new InMemoryWireRecordPersistence();
    ctx = createTestAgent(
      { persistence },
      telemetryServices(recordingTelemetry(records)),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    ctx.get(IAgentContextMemoryService);
    await ctx.restorePersisted();

    ctx.appendTurnExchange('u1', 'a1');
    await ctx.dispatcher.dispatch(
      new ToolsUpdateStore({ agentId: 'main', key: 'todo', value: [{ title: 'kept', status: 'pending' }] }),
    );
    ctx.appendTurnExchange('u2', 'a2');
    await ctx.dispatcher.dispatch(
      new ToolsUpdateStore({ agentId: 'main', key: 'todo', value: [{ title: 'doomed', status: 'pending' }] }),
    );
    ctx.get(IWireService).append({
      type: 'human.agent.turn.ended',
      kind: 'event',
      turnId: 0,
      outcome: 'done',
      time: 100,
    });
    await ctx.get(IAgentConversationUndoService).undo(1);

    expect(ctx.context.get().map(messageText)).toEqual(['user:u1', 'assistant:a1']);
    expect(ctx.get(IAgentTodoService).get().map((item) => item.title)).toEqual(['kept']);
    const afterNew = await ctx.persistedWireRecords();
    const legacyAfterNew = legacyWireFold(afterNew);
    expect(legacyAfterNew.context).toEqual(['user:u1', 'assistant:a1']);
    expect(legacyAfterNew.todo).toEqual(['kept']);
    expect(legacyAfterNew.skippedUnknownTypes).toEqual([
      'human.agent.turn.ended',
      'agent.switched',
    ]);

    persistence.records.push(
      {
        type: 'context.append_message',
        agentId: 'main',
        message: {
          role: 'user',
          content: [{ type: 'text', text: 'u3' }],
          toolCalls: [],
          origin: { kind: 'user' },
        },
        time: 200,
      },
      {
        type: 'context.append_message',
        agentId: 'main',
        message: { role: 'assistant', content: [{ type: 'text', text: 'a3' }], toolCalls: [] },
        time: 201,
      },
      { type: 'context.undo', agentId: 'main', count: 1, time: 202 },
    );
    const finalRecords = [...persistence.records];
    const legacyFinal = legacyWireFold(finalRecords);
    expect(legacyFinal.context).toEqual(['user:u1', 'assistant:a1']);
    expect(legacyFinal.todo).toEqual(['kept']);
    expect(legacyFinal.skippedUnknownTypes).toEqual([
      'human.agent.turn.ended',
      'agent.switched',
    ]);

    const reopened = createTestAgent(
      { autoConfigure: false, persistence: new InMemoryWireRecordPersistence(finalRecords) },
      telemetryServices(recordingTelemetry([])),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    try {
      reopened.get(IAgentContextMemoryService);
      await reopened.restorePersisted();
      expect(reopened.context.get().map(messageText)).toEqual(legacyFinal.context);
      expect(reopened.get(IAgentTodoService).get().map((item) => item.title)).toEqual(
        legacyFinal.todo,
      );
    } finally {
      await reopened.dispose();
    }
    await ctx.dispose();

    const orphanRecords: WireRecord[] = [
      { type: 'metadata', protocol_version: WIRE_PROTOCOL_VERSION, created_at: 1 },
      {
        type: 'context.append_message',
        agentId: 'main',
        message: {
          role: 'user',
          content: [{ type: 'text', text: 'u1' }],
          toolCalls: [],
          origin: { kind: 'user' },
        },
        time: 1,
      },
      {
        type: 'context.append_message',
        agentId: 'main',
        message: { role: 'assistant', content: [{ type: 'text', text: 'a1' }], toolCalls: [] },
        time: 2,
      },
      {
        type: 'context.append_message',
        agentId: 'main',
        message: {
          role: 'user',
          content: [{ type: 'text', text: 'u2' }],
          toolCalls: [],
          origin: { kind: 'user' },
        },
        time: 3,
      },
      {
        type: 'context.append_message',
        agentId: 'main',
        message: { role: 'assistant', content: [{ type: 'text', text: 'a2' }], toolCalls: [] },
        time: 4,
      },
      {
        type: 'agent.switched',
        agentId: 'main',
        branch: 'b1',
        reason: 'undo',
        base: { branch: 'main', line: 3 },
        turns: 1,
        legacyUndoLine: 7,
        time: 5,
      },
    ];
    const legacyOrphan = legacyWireFold(orphanRecords);
    expect(legacyOrphan.context).toEqual(['user:u1', 'assistant:a1', 'user:u2', 'assistant:a2']);
    expect(legacyOrphan.skippedUnknownTypes).toEqual(['agent.switched']);

    ctx = createTestAgent(
      { autoConfigure: false, persistence: new InMemoryWireRecordPersistence(orphanRecords) },
      telemetryServices(recordingTelemetry([])),
      execEnvServices({ hostFs: createFakeHostFs({ mkdir: async () => {} }) }),
    );
    ctx.get(IAgentContextMemoryService);
    await ctx.restorePersisted();
    expect(ctx.context.get().map(messageText)).toEqual(['user:u1', 'assistant:a1']);
  });
});

function messageText(message: ContextMessage): string {
  return `${message.role}:${message.content
    .map((part) => (part.type === 'text' ? part.text : ''))
    .join('')}`;
}

function legacyWireFold(records: readonly WireRecord[]): {
  readonly context: readonly string[];
  readonly todo: readonly string[];
  readonly skippedUnknownTypes: readonly string[];
} {
  const transcript: ContextMessage[] = [];
  const todoCheckpoints: string[][] = [];
  let todo: string[] = [];
  const skippedUnknownTypes: string[] = [];
  let clearFloor = 0;
  const applyUndo = (count: number): void => {
    let removedUserCount = 0;
    for (let i = transcript.length - 1; i >= clearFloor; i--) {
      const message = transcript[i]!;
      if (message.origin?.kind === 'injection') continue;
      if (message.origin?.kind === 'compaction_summary') break;
      transcript.splice(i, 1);
      if (!isUndoAnchor(message)) continue;
      removedUserCount++;
      while (i > clearFloor && isPromptOwnedInjection(transcript[i - 1]!, message)) {
        transcript.splice(i - 1, 1);
        i--;
      }
      if (removedUserCount >= count) break;
    }
    const targetIndex = todoCheckpoints.length - count;
    const target = todoCheckpoints[targetIndex];
    if (target === undefined) return;
    todo = [...target];
    todoCheckpoints.length = targetIndex;
  };
  for (const record of records) {
    switch (record.type) {
      case 'metadata':
        break;
      case 'context.append_message': {
        const message = record['message'] as ContextMessage;
        transcript.push(message);
        if (isUndoAnchor(message)) todoCheckpoints.push([...todo]);
        break;
      }
      case 'context.undo': {
        const count = record['count'];
        if (typeof count === 'number') applyUndo(count);
        break;
      }
      case 'context.clear':
        clearFloor = transcript.length;
        todoCheckpoints.length = 0;
        break;
      case 'tools.update_store': {
        if (record['key'] === 'todo') {
          todo = (record['value'] as { title: string }[]).map((item) => item.title);
        }
        break;
      }
      case 'context.undone':
        break;
      default:
        if (record.type.startsWith('agent.') || record.type.startsWith('human.')) {
          skippedUnknownTypes.push(record.type);
        }
        break;
    }
  }
  return {
    context: transcript.slice(clearFloor).map(messageText),
    todo,
    skippedUnknownTypes,
  };
}