diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index d83a613d3..19f930cbe 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/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 { // `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; + 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' diff --git a/packages/agent-core-v2/src/wire/wireService.ts b/packages/agent-core-v2/src/wire/wireService.ts index 15e64f237..b11786f21 100644 --- a/packages/agent-core-v2/src/wire/wireService.ts +++ b/packages/agent-core-v2/src/wire/wireService.ts @@ -66,5 +66,5 @@ export interface IWireService { handler: (state: DeepReadonly, prev: DeepReadonly) => void, ): IDisposable; onEmission(handler: (emission: WireEmission) => void): IDisposable; - onRestored(handler: () => void): IDisposable; + onRestored(handler: () => void | Promise): IDisposable; } diff --git a/packages/agent-core-v2/src/wire/wireServiceImpl.ts b/packages/agent-core-v2/src/wire/wireServiceImpl.ts index 5cd716e16..b59ff675c 100644 --- a/packages/agent-core-v2/src/wire/wireServiceImpl.ts +++ b/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, ModelInstance>(); private readonly reducerIndex = new Map(); private readonly emissionEmitter = this._register(new Emitter()); - private readonly restoredEmitter = this._register(new Emitter()); + private readonly restoredHandlers = new Set<() => void | Promise>(); 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): IDisposable { + this.restoredHandlers.add(handler); + return toDisposable(() => this.restoredHandlers.delete(handler)); } attach(model: DerivedModelDef): 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 { @@ -285,6 +286,16 @@ export class WireService extends Disposable implements IWireService { return { type: op.type, payload }; } + private async fireRestored(): Promise { + 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): void { if (this.log === undefined) return; diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index e5e893651..84c8c1b87 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -557,18 +557,22 @@ export function questionServices(service: ISessionQuestionService): TestAgentSer export function externalHookServices( hookRunner: Pick | 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 | undefined, +): IExternalHooksRunnerService { + return hookRunner === undefined + ? noopHookRunner + : isRunnerLike(hookRunner) + ? hookRunner + : { ...noopHookRunner, ...hookRunner }; +} + function isRunnerLike( value: Pick, ): 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 { 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 { - 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((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 { @@ -1989,6 +2015,73 @@ function isSystemReminderMessage(message: ContextMessage): boolean { return text.startsWith(''); } +function pendingTaskNotificationKeys(records: readonly PersistedWireRecord[]): readonly string[] { + const terminal = new Set(); + const delivered = new Set(); + 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; + 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; + 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(); diff --git a/packages/agent-core-v2/test/store/store.test.ts b/packages/agent-core-v2/test/store/store.test.ts index 37b8e4fc2..1106ff054 100644 --- a/packages/agent-core-v2/test/store/store.test.ts +++ b/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); diff --git a/packages/agent-core-v2/test/task/rpc-events.test.ts b/packages/agent-core-v2/test/task/rpc-events.test.ts index 25aa7d9da..c8ee3127a 100644 --- a/packages/agent-core-v2/test/task/rpc-events.test.ts +++ b/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; } diff --git a/packages/agent-core-v2/test/task/service.test.ts b/packages/agent-core-v2/test/task/service.test.ts index 701837031..333b7580a 100644 --- a/packages/agent-core-v2/test/task/service.test.ts +++ b/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, { diff --git a/packages/agent-core-v2/test/wireRecord/resume.test.ts b/packages/agent-core-v2/test/wireRecord/resume.test.ts index 7056761ba..e92b3a238 100644 --- a/packages/agent-core-v2/test/wireRecord/resume.test.ts +++ b/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 {