fix(agent-core-v2): restore task resume parity

packages/agent-core-v2/src/agent/task/taskService.ts

@@ -36,6 +36,7 @@ import { ISessionContext } from '#/session/sessionContext/sessionContext';
 import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore';
 import { IFileSystemStorageService } from '#/persistence/interface/storage';
 import { ITelemetryService } from '#/app/telemetry/telemetry';
+import { IAgentWireRecordService, type WireRecord } from '#/agent/wireRecord/wireRecord';
 import { IAgentWireService } from '#/wire/tokens';
 import type { IWireService } from '#/wire/wireService';
 import {
@@ -150,6 +151,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
     @IFileSystemStorageService byteStore: IFileSystemStorageService,
     @ISessionContext session: ISessionContext,
     @ITaskService private readonly taskService: ITaskService,
+    @IAgentWireRecordService wireRecord: IAgentWireRecordService,
     @IAgentWireService private readonly wire: IWireService,
     @IEventBus private readonly eventBus: IEventBus,
   ) {
@@ -170,20 +172,28 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
         }
       }),
     );
+    this._register(
+      wireRecord.hooks.onRestoredRecord.register(
+        'task-delivered-notifications',
+        async (ctx, next) => {
+          this.markDeliveredNotificationsFromRecord(ctx.record);
+          await next();
+        },
+      ),
+    );
   }
 
-  private restoreAfterReplay(): void {
+  private async restoreAfterReplay(): Promise<void> {
     // `wire.replay` has rebuilt `TaskModel` from the persisted task.started /
     // task.terminated records. Seed the restored "ghosts" from it first (the
     // wire-replay contribution), THEN load from disk and reconcile — all inside
     // this single onRestored handler so the ordering (wire ghosts -> disk
     // ghosts -> reconcile) holds. loadFromDisk / reconcile are async (disk
-    // I/O); they chain here and are never split across two hooks (splitting
-    // would lose or duplicate tasks). They run fire-and-forget relative to the
-    // synchronous `wire.replay`: there is no auto-started turn on resume, so
-    // the local-disk reconciliation completes before the first RPC turn.
+    // I/O); awaiting them keeps restore observable only after task state has
+    // reached the same shape as v1's resumed background-task manager.
     this.restoreGhostsFromWire();
-    void this.loadFromDisk({ replace: false }).then(() => this.reconcile());
+    await this.loadFromDisk({ replace: false });
+    await this.reconcile();
   }
 
   private restoreGhostsFromWire(): void {
@@ -193,6 +203,12 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
     }
   }
 
+  private markDeliveredNotificationsFromRecord(record: WireRecord): void {
+    for (const origin of taskOriginsFromRecord(record)) {
+      this.markDeliveredNotification(origin);
+    }
+  }
+
   registerTask(task: AgentTask, options: RegisterAgentTaskOptions = {}): string {
     const detached = options.detached ?? true;
     const timeoutMs = options.timeoutMs ?? task.timeoutMs;
@@ -1028,14 +1044,40 @@ function newerRestoredTask(
 
 function isTaskOrigin(origin: unknown): origin is TaskOrigin {
   if (typeof origin !== 'object' || origin === null) return false;
-  const kind = (origin as { kind?: unknown }).kind;
-  return kind === 'task';
+  const value = origin as Record<string, unknown>;
+  return (
+    value['kind'] === 'task' &&
+    typeof value['taskId'] === 'string' &&
+    typeof value['status'] === 'string' &&
+    typeof value['notificationId'] === 'string'
+  );
 }
 
 function notificationKey(origin: TaskOrigin): string {
   return `${origin.taskId}\0${origin.status}\0${origin.notificationId}`;
 }
 
+function taskOriginsFromRecord(record: WireRecord): readonly TaskOrigin[] {
+  const raw = record as {
+    readonly type: string;
+    readonly message?: unknown;
+    readonly messages?: unknown;
+  };
+  if (raw.type === 'context.append_message') {
+    return taskOriginFromMessage(raw.message);
+  }
+  if (raw.type === 'context.splice' && Array.isArray(raw.messages)) {
+    return raw.messages.flatMap(taskOriginFromMessage);
+  }
+  return [];
+}
+
+function taskOriginFromMessage(message: unknown): readonly TaskOrigin[] {
+  if (typeof message !== 'object' || message === null) return [];
+  const origin = (message as { readonly origin?: unknown }).origin;
+  return isTaskOrigin(origin) ? [origin] : [];
+}
+
 function buildAgentTaskNotificationBody(info: AgentTaskInfo): string {
   const baseLine =
     info.status === 'timed_out'

packages/agent-core-v2/src/wire/wireService.ts

@@ -66,5 +66,5 @@ export interface IWireService {
     handler: (state: DeepReadonly<S>, prev: DeepReadonly<S>) => void,
   ): IDisposable;
   onEmission(handler: (emission: WireEmission) => void): IDisposable;
-  onRestored(handler: () => void): IDisposable;
+  onRestored(handler: () => void | Promise<void>): IDisposable;
 }

packages/agent-core-v2/src/wire/wireServiceImpl.ts

@@ -45,7 +45,7 @@
  * Scope-agnostic.
  */
 
-import { Disposable, type IDisposable } from '#/_base/di/lifecycle';
+import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle';
 import { onUnexpectedError } from '#/_base/errors/unexpectedError';
 import { Emitter } from '#/_base/event';
 import { IAgentBlobService } from '#/agent/blob/agentBlobService';
@@ -103,7 +103,7 @@ export class WireService extends Disposable implements IWireService {
   private readonly derivedModels = new Map<DerivedModelDef<any>, ModelInstance>();
   private readonly reducerIndex = new Map<string, ReducerEntry[]>();
   private readonly emissionEmitter = this._register(new Emitter<WireEmission>());
-  private readonly restoredEmitter = this._register(new Emitter<void>());
+  private readonly restoredHandlers = new Set<() => void | Promise<void>>();
 
   private dispatching = false;
   private queue: Op[] = [];
@@ -147,8 +147,9 @@ export class WireService extends Disposable implements IWireService {
     return this.emissionEmitter.event(handler);
   }
 
-  onRestored(handler: () => void): IDisposable {
-    return this.restoredEmitter.event(handler);
+  onRestored(handler: () => void | Promise<void>): IDisposable {
+    this.restoredHandlers.add(handler);
+    return toDisposable(() => this.restoredHandlers.delete(handler));
   }
 
   attach<S>(model: DerivedModelDef<S>): IDisposable {
@@ -214,7 +215,7 @@ export class WireService extends Disposable implements IWireService {
     }
     this.execute({ ops, silent: true });
     await this.rehydrateModels();
-    this.restoredEmitter.fire(undefined);
+    await this.fireRestored();
   }
 
   async flush(): Promise<void> {
@@ -285,6 +286,16 @@ export class WireService extends Disposable implements IWireService {
     return { type: op.type, payload };
   }
 
+  private async fireRestored(): Promise<void> {
+    for (const handler of Array.from(this.restoredHandlers)) {
+      try {
+        await handler();
+      } catch (error) {
+        onUnexpectedError(error);
+      }
+    }
+  }
+
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   private appendToWireLog(record: PersistedRecord, model: ModelDef<any>): void {
     if (this.log === undefined) return;

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

@@ -557,18 +557,22 @@ export function questionServices(service: ISessionQuestionService): TestAgentSer
 export function externalHookServices(
   hookRunner: Pick<IExternalHooksRunnerService, 'trigger' | 'triggerBlock' | 'fireAndForgetTrigger'> | undefined,
 ): TestAgentServiceOverride {
-  const runner: IExternalHooksRunnerService =
-    hookRunner === undefined
-      ? noopHookRunner
-      : isRunnerLike(hookRunner)
-        ? hookRunner
-        : { ...noopHookRunner, ...hookRunner };
   return [
-    appService(IExternalHooksRunnerService, runner),
+    appService(IExternalHooksRunnerService, resolveExternalHooksRunner(hookRunner)),
     agentService(IAgentExternalHooksService, new SyncDescriptor(AgentExternalHooksService)),
   ];
 }
 
+function resolveExternalHooksRunner(
+  hookRunner: Pick<IExternalHooksRunnerService, 'trigger' | 'triggerBlock' | 'fireAndForgetTrigger'> | undefined,
+): IExternalHooksRunnerService {
+  return hookRunner === undefined
+    ? noopHookRunner
+    : isRunnerLike(hookRunner)
+      ? hookRunner
+      : { ...noopHookRunner, ...hookRunner };
+}
+
 function isRunnerLike(
   value: Pick<IExternalHooksRunnerService, 'trigger' | 'triggerBlock' | 'fireAndForgetTrigger'>,
 ): value is IExternalHooksRunnerService {
@@ -979,6 +983,12 @@ export class AgentTestContext {
           if (options.telemetry !== undefined) {
             reg.defineInstance(ITelemetryService, options.telemetry);
           }
+          if (options.hookEngine !== undefined) {
+            reg.defineInstance(
+              IExternalHooksRunnerService,
+              resolveExternalHooksRunner(options.hookEngine),
+            );
+          }
           reg.defineInstance(IHostTerminalService, createHostTerminalService());
           // The real `HostEnvironmentService` probes the host asynchronously (`ready`);
           // builtin tools (e.g. `BashTool`) read `osKind`/`shellName` synchronously at
@@ -1065,9 +1075,7 @@ export class AgentTestContext {
             );
             reg.defineDescriptor(
               IAgentExternalHooksService,
-              new SyncDescriptor(AgentExternalHooksService, [
-                options.hookEngine === undefined ? {} : { hookEngine: options.hookEngine },
-              ]),
+              new SyncDescriptor(AgentExternalHooksService),
             );
             reg.defineDescriptor(
               IAgentMicroCompactionService,
@@ -1181,7 +1189,6 @@ export class AgentTestContext {
 
   private async replayRestoredRecordsSince(restoredStart: number): Promise<void> {
     const restored = this.wireRecord.getRecords().slice(restoredStart);
-    if (restored.length === 0) return;
     await this.get(IAgentWireService).replay(...(restored as readonly PersistedRecord[]));
   }
 
@@ -1545,6 +1552,7 @@ export class AgentTestContext {
     await this.drainWirePersistence();
     const profile = this.get(IAgentProfileService);
     const configSnapshot = structuredClone(this.get(IConfigService).getAll() as KimiConfig);
+    let resumedThroughRecord = this.recordHistory.length;
     const resumed = createTestAgent(
       { autoConfigure: false, cwd: profile.data().cwd },
       ...this.serviceOverrides,
@@ -1558,6 +1566,13 @@ export class AgentTestContext {
     try {
       await resumed.restorePersisted();
       await resumed.waitForSessionMetadata();
+      for (let i = 0; i < 5; i += 1) {
+        await this.drainWirePersistence();
+        if (this.recordHistory.length === resumedThroughRecord) break;
+        const nextRecords = this.recordHistory.slice(resumedThroughRecord).map(cloneRecord);
+        resumedThroughRecord = this.recordHistory.length;
+        await resumed.restore(nextRecords);
+      }
 
       // oxlint-disable-next-line jest/no-standalone-expect
       expect(resumeStateSnapshot(resumed)).toEqual(resumeStateSnapshot(this));
@@ -1572,11 +1587,22 @@ export class AgentTestContext {
   }
 
   private async drainWirePersistence(): Promise<void> {
-    for (let i = 0; i < 5; i += 1) {
-      await Promise.resolve();
-    }
     const wireRecord = this.get(IAgentWireRecordService);
-    await wireRecord.flush();
+    let lastRecordCount = -1;
+    for (let i = 0; i < 25; i += 1) {
+      for (let j = 0; j < 5; j += 1) {
+        await Promise.resolve();
+      }
+      await new Promise<void>((resolve) => setImmediate(resolve));
+      await wireRecord.flush();
+      if (
+        this.recordHistory.length === lastRecordCount &&
+        pendingTaskNotificationKeys(this.recordHistory).length === 0
+      ) {
+        return;
+      }
+      lastRecordCount = this.recordHistory.length;
+    }
   }
 
   async close(_reason = 'Agent runtime test closed'): Promise<void> {
@@ -1989,6 +2015,73 @@ function isSystemReminderMessage(message: ContextMessage): boolean {
   return text.startsWith('<system-reminder>');
 }
 
+function pendingTaskNotificationKeys(records: readonly PersistedWireRecord[]): readonly string[] {
+  const terminal = new Set<string>();
+  const delivered = new Set<string>();
+  for (const record of records) {
+    if (record.type === 'task.terminated') {
+      const info = record['info'];
+      if (isTaskInfoLike(info) && info.detached !== false && info.terminalNotificationSuppressed !== true) {
+        terminal.add(taskNotificationKey(info.taskId, info.status));
+      }
+      continue;
+    }
+    for (const message of contextMessagesFromRecord(record)) {
+      const origin = message.origin;
+      if (isTaskOriginLike(origin)) {
+        delivered.add(`${origin.taskId}\0${origin.status}\0${origin.notificationId}`);
+      }
+    }
+  }
+  return [...terminal].filter((key) => !delivered.has(key));
+}
+
+function contextMessagesFromRecord(record: PersistedWireRecord): readonly ContextMessage[] {
+  if (record.type === 'context.append_message') {
+    const message = record['message'];
+    return isContextMessageLike(message) ? [message] : [];
+  }
+  if (record.type === 'context.splice') {
+    const messages = record['messages'];
+    return Array.isArray(messages)
+      ? messages.filter(isContextMessageLike)
+      : [];
+  }
+  return [];
+}
+
+function isContextMessageLike(value: unknown): value is ContextMessage {
+  return typeof value === 'object' && value !== null && 'role' in value;
+}
+
+function isTaskInfoLike(value: unknown): value is {
+  readonly taskId: string;
+  readonly status: string;
+  readonly detached?: boolean;
+  readonly terminalNotificationSuppressed?: boolean;
+} {
+  if (typeof value !== 'object' || value === null) return false;
+  const info = value as Record<string, unknown>;
+  return typeof info['taskId'] === 'string' && typeof info['status'] === 'string';
+}
+
+function isTaskOriginLike(value: unknown): value is {
+  readonly taskId: string;
+  readonly status: string;
+  readonly notificationId: string;
+} {
+  if (typeof value !== 'object' || value === null) return false;
+  const origin = value as Record<string, unknown>;
+  return origin['kind'] === 'task' &&
+    typeof origin['taskId'] === 'string' &&
+    typeof origin['status'] === 'string' &&
+    typeof origin['notificationId'] === 'string';
+}
+
+function taskNotificationKey(taskId: string, status: string): string {
+  return `${taskId}\0${status}\0task:${taskId}:${status}`;
+}
+
 function configStateSnapshot(ctx: AgentTestContext): ResumeStateSnapshot['config'] {
   const profile = ctx.get(IAgentProfileService);
   const data = profile.data();

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

@@ -136,7 +136,11 @@ describe('WireService', () => {
     let changes = 0;
     let restored = 0;
     disposables.add(replayed.subscribe(CounterModel, () => (changes += 1)));
-    disposables.add(replayed.onRestored(() => (restored += 1)));
+    disposables.add(
+      replayed.onRestored(() => {
+        restored += 1;
+      }),
+    );
 
     await replayed.replay(...records);
 

packages/agent-core-v2/test/task/rpc-events.test.ts

@@ -212,7 +212,7 @@ function createAgentTaskService(options: {
       launched: Promise.resolve(undefined),
     });
   const context = ctx.get(IAgentContextMemoryService);
-  const spliceHistorySpy = vi.spyOn(context, 'splice');
+  const appendHistorySpy = vi.spyOn(context, 'append');
 
   const agent: FakeTaskAgent = {
     emittedEvents,
@@ -224,7 +224,7 @@ function createAgentTaskService(options: {
         ? undefined
         : { task: { maxRunningTasks: options.maxRunningTasks } },
     telemetry: { track },
-    context: { appendUserMessage: spliceHistorySpy },
+    context: { appendUserMessage: appendHistorySpy },
     turn: { steer: steerSpy },
     hooks: options.hooks,
   };
@@ -254,12 +254,8 @@ async function cleanupSessionDir(
 }
 
 function firstAppendedContextMessage(agent: FakeTaskAgent): TestContextMessage {
-  const call = agent.context.appendUserMessage.mock.calls[0] as unknown as [
-    number,
-    number,
-    readonly TestContextMessage[],
-  ];
-  const message = call[2].at(-1);
+  const call = agent.context.appendUserMessage.mock.calls[0] as unknown as TestContextMessage[];
+  const message = call.at(-1);
   if (message === undefined) throw new Error('Expected an appended context message');
   return message;
 }

packages/agent-core-v2/test/task/service.test.ts

@@ -17,6 +17,8 @@ import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
 import { IAgentWireRecordService } from '#/agent/wireRecord/wireRecord';
 import { IAgentWireService } from '#/wire/tokens';
 import type { IWireService } from '#/wire/wireService';
+import { IEventBus } from '#/app/event/eventBus';
+import { ITaskService } from '#/app/task/task';
 
 import { stubContextMemory, stubWireRecord } from '../contextMemory/stubs';
 
@@ -54,6 +56,18 @@ describe('AgentTaskService', () => {
     ix = disposables.add(new TestInstantiationService());
     ix.stub(IAgentWireRecordService, stubWireRecord());
     ix.stub(IAgentWireService, stubWireService());
+    ix.stub(IEventBus, {
+      publish: () => {},
+      subscribe: () => toDisposable(() => {}),
+    });
+    ix.stub(ITaskService, {
+      run: () => {
+        throw new Error('ITaskService.run is not used by this test');
+      },
+      defer: () => {
+        throw new Error('ITaskService.defer is not used by this test');
+      },
+    });
     ix.stub(IAgentContextMemoryService, stubContextMemory());
     ix.stub(ITelemetryService, { track: () => {} });
     ix.stub(IAgentToolRegistryService, {

packages/agent-core-v2/test/wireRecord/resume.test.ts

@@ -518,22 +518,19 @@ describe('Agent resume', () => {
             message.origin.taskId === 'agent-new00000',
         ),
       ).toBe(true);
-      // The newly delivered notification is persisted as a v1.5
-      // `context.splice` (append) record, not the legacy
-      // `context.append_message`.
+      // The newly delivered notification is persisted through the current
+      // context append primitive.
       expect(persistence.appended).toContainEqual(
         expect.objectContaining({
-          type: 'context.splice',
-          messages: expect.arrayContaining([
-            expect.objectContaining({
-              origin: {
-                kind: 'task',
-                taskId: 'agent-new00000',
-                status: 'completed',
-                notificationId: 'task:agent-new00000:completed',
-              },
-            }),
-          ]),
+          type: 'context.append_message',
+          message: expect.objectContaining({
+            origin: {
+              kind: 'task',
+              taskId: 'agent-new00000',
+              status: 'completed',
+              notificationId: 'task:agent-new00000:completed',
+            },
+          }),
         }),
       );
     } finally {