From e3e64b8d241001296174a2258674b59387d12fe6 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Wed, 9 Sep 2026 23:49:56 +0800 Subject: [PATCH 1/3] feat(agent-core-v2): add event-sourced store core and rewire human agent/session human/eventStore core (schema-first defineEvent factories, serial fold-first dispatch with snapshotting, late-join slices, journal reset), agent events/slices/historySchema, machine rewrite around a store actor with mirror consumption, session stores with turnIndex-boundary undo, todo slice, v2 migration to the new event types, and the loop engine bridge. --- .../src/agent/loop/machine/engine.ts | 36 +- .../agent-core-v2/src/human/agent/events.ts | 75 +++ .../src/human/agent/historySchema.ts | 101 ++++ .../agent-core-v2/src/human/agent/machine.ts | 273 ++++++---- .../agent-core-v2/src/human/agent/replay.ts | 40 -- .../agent-core-v2/src/human/agent/slices.ts | 131 +++++ .../src/human/eventStore/actor.ts | 41 ++ .../src/human/eventStore/eventStore.ts | 342 ++++++++++++ .../src/human/eventStore/events.ts | 61 +++ .../src/human/eventStore/index.ts | 5 + .../src/human/eventStore/journal.ts | 104 ++++ .../src/human/eventStore/slice.ts | 27 + packages/agent-core-v2/src/human/index.ts | 13 +- .../agent-core-v2/src/human/persist/agent.ts | 134 ----- .../agent-core-v2/src/human/persist/open.ts | 8 +- .../src/human/persist/session.ts | 207 -------- .../src/human/persist/v2/migrate.ts | 31 +- .../agent-core-v2/src/human/session/events.ts | 29 + .../src/human/session/machine.ts | 68 +-- .../agent-core-v2/src/human/session/slices.ts | 39 ++ .../agent-core-v2/src/human/session/stores.ts | 155 ++++++ .../agent-core-v2/src/human/session/undo.ts | 118 ----- .../src/human/test/agent/machine.test.ts | 413 +++++++++------ .../src/human/test/agent/persist.test.ts | 264 --------- .../human/test/eventStore/eventStore.test.ts | 294 ++++++++++ .../src/human/test/media/tool.test.ts | 18 +- .../src/human/test/session/machine.test.ts | 194 ++++--- .../src/human/test/session/migrate-v2.test.ts | 111 ++-- .../src/human/test/session/persist.test.ts | 500 ------------------ .../src/human/test/session/stores.test.ts | 223 ++++++++ .../src/human/test/todo/plugin.test.ts | 68 ++- .../src/human/test/tool-select/plugin.test.ts | 22 +- .../src/human/test/usage/machine.test.ts | 18 +- .../agent-core-v2/src/human/todo/plugin.ts | 19 +- .../agent-core-v2/src/human/todo/slice.ts | 34 ++ .../agent-core-v2/src/human/todo/state.ts | 33 -- packages/agent-core-v2/src/human/todo/tool.ts | 16 +- 37 files changed, 2449 insertions(+), 1816 deletions(-) create mode 100644 packages/agent-core-v2/src/human/agent/events.ts create mode 100644 packages/agent-core-v2/src/human/agent/historySchema.ts delete mode 100644 packages/agent-core-v2/src/human/agent/replay.ts create mode 100644 packages/agent-core-v2/src/human/agent/slices.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/actor.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/eventStore.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/events.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/index.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/journal.ts create mode 100644 packages/agent-core-v2/src/human/eventStore/slice.ts delete mode 100644 packages/agent-core-v2/src/human/persist/agent.ts delete mode 100644 packages/agent-core-v2/src/human/persist/session.ts create mode 100644 packages/agent-core-v2/src/human/session/events.ts create mode 100644 packages/agent-core-v2/src/human/session/slices.ts create mode 100644 packages/agent-core-v2/src/human/session/stores.ts delete mode 100644 packages/agent-core-v2/src/human/session/undo.ts delete mode 100644 packages/agent-core-v2/src/human/test/agent/persist.test.ts create mode 100644 packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts delete mode 100644 packages/agent-core-v2/src/human/test/session/persist.test.ts create mode 100644 packages/agent-core-v2/src/human/test/session/stores.test.ts create mode 100644 packages/agent-core-v2/src/human/todo/slice.ts delete mode 100644 packages/agent-core-v2/src/human/todo/state.ts diff --git a/packages/agent-core-v2/src/agent/loop/machine/engine.ts b/packages/agent-core-v2/src/agent/loop/machine/engine.ts index 6ff915033b0..f110df1b399 100644 --- a/packages/agent-core-v2/src/agent/loop/machine/engine.ts +++ b/packages/agent-core-v2/src/agent/loop/machine/engine.ts @@ -6,6 +6,10 @@ import type { ToolInfo, ToolResult as AgentToolResult, ToolUpdate as AgentToolUp import type { ToolInputDisplay } from '#/tool/toolInputDisplay'; import { createAgentMachine } from '#human/agent/machine'; import { createTurnMachine, type AssistantEntry, type HistoryMessage } from '#human/agent/turn'; +import { messageAppended, turnEnded } from '#human/agent/events'; +import { agentSlices, type AgentEventStore } from '#human/agent/slices'; +import { createEventStoreSync } from '#human/eventStore/eventStore'; +import { memoryJournal } from '#human/eventStore/journal'; import type { LlmErrorMessage } from '#human/llm/errors'; import type { FinishInfo } from '#human/llm/finish-reason'; import type { StreamedMessagePart, UserMessage } from '#human/llm/message'; @@ -264,6 +268,11 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin publish({ type: 'toolBatchFailed', error }); }, }); + const store: AgentEventStore = createEventStoreSync({ journal: memoryJournal(), slices: agentSlices }); + const initialTurnId = options.initialTurnId ?? 0; + if (initialTurnId > 0) { + void store.dispatch(turnEnded({ turnId: initialTurnId - 1, outcome: 'done' })); + } const actor = createActor( createAgentMachine({ tools: tools.tools, @@ -278,7 +287,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin ), abortTimeoutMs: options.abortTimeoutMs, }), - { input: { request: { model: options.model, systemPrompt: options.systemPrompt }, turnId: options.initialTurnId } }, + { input: { request: { model: options.model, systemPrompt: options.systemPrompt }, store } }, ); const subscriptions: Subscription[] = [ actor.on('turn.started', (event) => { @@ -433,7 +442,19 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin actor.send({ type: 'input.abort' }); }, resetHistory: (history) => { - actor.send({ type: 'context.reset', history }); + const journal = memoryJournal(); + for (const message of history) { + void journal.append({ type: messageAppended.type, kind: 'event', data: messageAppended({ message }) }); + } + const nextTurnId = store.slice('turnIndex').nextTurnId; + if (nextTurnId > 0) { + void journal.append({ + type: turnEnded.type, + kind: 'event', + data: turnEnded({ turnId: nextTurnId - 1, outcome: 'done' }), + }); + } + void store.reset(journal); }, stop: () => { for (const subscription of subscriptions) subscription.unsubscribe(); @@ -442,6 +463,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin snapshot: () => { const snapshot = actor.getSnapshot() as unknown as MachineSnapshotLike; const value = snapshot.value; + const storeState = store.getState(); const turnRef = snapshot.children['turn']; let turn: MachineEngineTurnSnapshot | undefined; if (turnRef !== undefined) { @@ -455,7 +477,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin : ('running' as const); const context = turnSnapshot.context; turn = { - turnId: snapshot.context.turnId, + turnId: storeState.turnIndex.nextTurnId, phase, step: context.steps, retry: @@ -479,10 +501,10 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin waitingForBackground: typeof value === 'object' && value !== null && 'idle' in value && (value as { idle?: unknown }).idle === 'waiting', - queueLength: snapshot.context.queue.length, - queueIds: snapshot.context.queue.map((entry) => entry.id), - notificationCount: snapshot.context.notifications.length, - reminderCount: snapshot.context.reminders.length, + queueLength: storeState.queue.length, + queueIds: storeState.queue.map((entry) => entry.id), + notificationCount: storeState.notifications.length, + reminderCount: storeState.reminders.length, backgroundCount: Object.keys(snapshot.context.background).length, turn, }; diff --git a/packages/agent-core-v2/src/human/agent/events.ts b/packages/agent-core-v2/src/human/agent/events.ts new file mode 100644 index 00000000000..9c45f46a046 --- /dev/null +++ b/packages/agent-core-v2/src/human/agent/events.ts @@ -0,0 +1,75 @@ +import { z } from 'zod'; + +import { defineEvent } from '#/eventStore/events'; + +import { historyMessageSchema, systemMessageSchema, userMessageSchema } from './historySchema'; + +export const messageAppended = defineEvent({ + type: 'message.appended', + schema: z.object({ message: historyMessageSchema }), +}); +export type MessageAppended = ReturnType; + +export const turnStarted = defineEvent({ + type: 'turn.started', + schema: z.object({ turnId: z.number().int(), queueItemId: z.string().optional() }), +}); +export type TurnStarted = ReturnType; + +export const turnEnded = defineEvent({ + type: 'turn.ended', + schema: z.object({ + turnId: z.number().int(), + outcome: z.enum(['done', 'failed', 'aborted']), + errorMessage: z.string().optional(), + }), +}); +export type TurnEnded = ReturnType; + +export const inputSubmitted = defineEvent({ + type: 'input.submitted', + schema: z.object({ id: z.string().optional(), message: userMessageSchema }), +}); +export type InputSubmitted = ReturnType; + +export const inputNotified = defineEvent({ + type: 'input.notified', + schema: z.object({ message: userMessageSchema, source: z.string().optional() }), +}); +export type InputNotified = ReturnType; + +export const inputReminded = defineEvent({ + type: 'input.reminded', + schema: z.object({ key: z.string(), message: z.union([userMessageSchema, systemMessageSchema]) }), +}); +export type InputReminded = ReturnType; + +export const inputSteered = defineEvent({ + type: 'input.steered', + schema: z.object({ id: z.string(), message: userMessageSchema }), +}); +export type InputSteered = ReturnType; + +export const inputCancelled = defineEvent({ + type: 'input.cancelled', + schema: z.object({ id: z.string() }), +}); +export type InputCancelled = ReturnType; + +export const queueDrained = defineEvent({ + type: 'queue.drained', + schema: z.object({ id: z.string().optional() }), +}); +export type QueueDrained = ReturnType; + +export const inputDrained = defineEvent({ type: 'input.drained', schema: z.object({}) }); +export type InputDrained = ReturnType; + +export const notificationsDrained = defineEvent({ type: 'notifications.drained', schema: z.object({}) }); +export type NotificationsDrained = ReturnType; + +export const stateUpdated = defineEvent({ + type: 'state.updated', + schema: z.object({ name: z.string(), value: z.unknown() }), +}); +export type StateUpdated = ReturnType; diff --git a/packages/agent-core-v2/src/human/agent/historySchema.ts b/packages/agent-core-v2/src/human/agent/historySchema.ts new file mode 100644 index 00000000000..07b6324b1b6 --- /dev/null +++ b/packages/agent-core-v2/src/human/agent/historySchema.ts @@ -0,0 +1,101 @@ +import { z } from 'zod'; + +import type { HistoryMessage } from './turn'; + +const textPartSchema = z.object({ type: z.literal('text'), text: z.string() }); +const thinkPartSchema = z.object({ + type: z.literal('think'), + think: z.string(), + encrypted: z.string().optional(), + detailsIndex: z.number().optional(), +}); +const imageUrlPartSchema = z.object({ + type: z.literal('image_url'), + imageUrl: z.object({ url: z.string(), id: z.string().optional(), name: z.string().optional() }), +}); +const audioUrlPartSchema = z.object({ + type: z.literal('audio_url'), + audioUrl: z.object({ url: z.string(), id: z.string().optional() }), +}); +const videoUrlPartSchema = z.object({ + type: z.literal('video_url'), + videoUrl: z.object({ url: z.string(), id: z.string().optional(), name: z.string().optional() }), +}); + +export const contentPartSchema = z.discriminatedUnion('type', [ + textPartSchema, + thinkPartSchema, + imageUrlPartSchema, + audioUrlPartSchema, + videoUrlPartSchema, +]); + +const toolDescriptionSchema = z.object({ + name: z.string(), + description: z.string(), + parameters: z.record(z.string(), z.unknown()), + deferred: z.literal(true).optional(), +}); + +const toolCallSchema = z.object({ + type: z.literal('function'), + id: z.string(), + name: z.string(), + arguments: z.string().nullable(), + extras: z.record(z.string(), z.unknown()).optional(), + rawId: z.string().optional(), + _streamIndex: z.union([z.number(), z.string()]).optional(), +}); + +export const systemMessageSchema = z.object({ + role: z.literal('system'), + content: z.array(contentPartSchema), + tools: z.array(toolDescriptionSchema).optional(), +}); + +export const userMessageSchema = z.object({ + role: z.literal('user'), + content: z.array(contentPartSchema), +}); + +export const assistantMessageSchema = z.object({ + role: z.literal('assistant'), + content: z.array(contentPartSchema), + toolCalls: z.array(toolCallSchema), +}); + +export const toolMessageSchema = z.object({ + role: z.literal('tool'), + content: z.array(contentPartSchema), + toolCallId: z.string(), +}); + +const tokenUsageSchema = z.object({ + inputOther: z.number(), + output: z.number(), + inputCacheRead: z.number(), + inputCacheCreation: z.number(), + raw: z.record(z.string(), z.unknown()).optional(), +}); + +const finishInfoSchema = z.object({ + finishReason: z.enum(['completed', 'tool_calls', 'truncated', 'filtered', 'paused', 'other']).nullable(), + rawFinishReason: z.string().nullable(), +}); + +const entryMetaSchema = z.object({ source: z.string().optional(), key: z.string().optional() }); + +const assistantMetaSchema = entryMetaSchema.extend({ + model: z.object({ provider: z.string(), model: z.string() }).optional(), + usage: tokenUsageSchema, + headers: z.record(z.string(), z.string()).optional(), + finish: finishInfoSchema.optional(), + messageId: z.string().optional(), +}); + +export const historyMessageSchema = z.union([ + z.object({ message: systemMessageSchema, meta: entryMetaSchema }), + z.object({ message: userMessageSchema, meta: entryMetaSchema }), + z.object({ message: assistantMessageSchema, meta: assistantMetaSchema }), + z.object({ message: toolMessageSchema, meta: entryMetaSchema }), +]) as z.ZodType; diff --git a/packages/agent-core-v2/src/human/agent/machine.ts b/packages/agent-core-v2/src/human/agent/machine.ts index 0a22ae7870b..35eb923acad 100644 --- a/packages/agent-core-v2/src/human/agent/machine.ts +++ b/packages/agent-core-v2/src/human/agent/machine.ts @@ -1,11 +1,6 @@ import { assign, emit, enqueueActions, fromCallback, sendTo, setup } from '#/xstate2'; -import { - createUserMessage, - type SystemMessage, - type ToolCall, - type UserMessage, -} from '#/llm/message'; +import { createUserMessage, type SystemMessage, type ToolCall, type UserMessage } from '#/llm/message'; import type { LlmRequestConfig } from '#/llm/requester/requester'; import type { ToolExecutor, ToolResult } from '#/tool/executor'; import { createToolMachine, type ToolEvent, type ToolOutput } from '#/tool/machine'; @@ -13,21 +8,28 @@ import type { ToolDefinition } from '#/tool/tool'; import { createWaitForTasks, type ToolActorRef } from './wait-for'; import { interruptReasonOf, type TurnInterruptReason } from './errors'; -import { createSystemEntry, createUserEntry } from './turn'; +import { + inputCancelled, + inputDrained, + inputNotified, + inputReminded, + inputSteered, + inputSubmitted, + messageAppended, + notificationsDrained, + queueDrained, + turnEnded, + turnStarted, +} from './events'; +import { createUserEntry } from './turn'; import { createAbortScope, withAbort, type AbortScope } from '#/utils/abort'; -import type { - createTurnMachine, - HistoryMessage, - TurnLlmEvent, - TurnOutput, - UserEntry, -} from './turn'; +import type { createTurnMachine, HistoryMessage, TurnLlmEvent, TurnOutput, UserEntry } from './turn'; +import type { AgentStoreActorLogic, AgentStoreState, QueuedPrompt } from './slices'; + +export type { QueuedPrompt } from './slices'; export interface AgentInput { request: LlmRequestConfig; - history?: readonly HistoryMessage[]; - turnId?: number; - branchId?: string; } export type AgentEvent = @@ -43,7 +45,10 @@ export type AgentEvent = | { type: 'turn.drain' } | { type: 'turn.reminders_consumed'; reminders: HistoryMessage[] } | { type: 'step.started'; step: number } - | { type: 'context.reset'; history: readonly HistoryMessage[]; branchId?: string }; + | { type: 'store.ready'; state: AgentStoreState; branch: string } + | { type: 'store.changed'; state: AgentStoreState } + | { type: 'store.reset'; state: AgentStoreState; branch: string } + | { type: 'store.error'; error: unknown }; export type AgentEmitted = | TurnLlmEvent @@ -69,11 +74,6 @@ interface ToolEntry { ref: ToolActorRef; } -export interface QueuedPrompt { - id?: string; - message: UserMessage; -} - export interface AgentMachineContext { input: AgentInput; messages: HistoryMessage[]; @@ -84,6 +84,7 @@ export interface AgentMachineContext { reminders: HistoryMessage[]; queue: QueuedPrompt[]; turnId: number; + activeTurnId?: number; branchId: string; drainedId?: string; } @@ -186,24 +187,22 @@ function drainPendingPatch( }; } -function steerPatch( - context: AgentMachineContext, - id: string, -): Partial> { - const index = context.queue.findIndex((entry) => entry.id === id); - if (index === -1) { - return {}; - } - const entry = context.queue[index] as QueuedPrompt; +function mirrorPatch(state: AgentStoreState): Pick< + AgentMachineContext, + 'messages' | 'queue' | 'notifications' | 'reminders' +> { return { - queue: context.queue.filter((_, i) => i !== index), - notifications: [...context.notifications, createUserEntry(entry.message, { source: 'input' })], + messages: [...state.history], + queue: [...state.queue], + notifications: [...state.notifications], + reminders: [...state.reminders], }; } export interface CreateAgentMachineOptions { tools?: readonly ToolDefinition[]; turnActor: ReturnType; + store: AgentStoreActorLogic; abortTimeoutMs?: number; maxStepsPerTurn?: number; } @@ -233,6 +232,7 @@ function dispatchTools(tools: readonly ToolDefinition[]): ToolExecutor { export function createAgentMachine({ tools, turnActor, + store, abortTimeoutMs, maxStepsPerTurn, }: CreateAgentMachineOptions) { @@ -247,6 +247,7 @@ export function createAgentMachine({ actors: { turnActor, toolActor: createToolMachine(executor), + storeActor: store, controllerGuard: fromCallback( ({ input }) => () => @@ -257,6 +258,24 @@ export function createAgentMachine({ forwardToParent: ({ self, event }) => { self._parent?.send(event); }, + resetMirror: assign(({ event }) => { + if (event.type !== 'store.reset') return {}; + return { + ...mirrorPatch(event.state), + turnTools: {}, + background: {}, + scope: createAbortScope(), + turnId: event.state.turnIndex.nextTurnId, + activeTurnId: undefined, + branchId: event.branch, + }; + }), + emitReset: emit(({ context }) => ({ type: 'context.reset' as const, branchId: context.branchId })), + abortBackgroundScopes: enqueueActions(({ context }) => { + for (const entry of Object.values(context.background)) { + entry.scope.abort(); + } + }), spawnTurnTools: assign(({ context, spawn, self, event }) => { if (event.type !== 'turn.spawn_tools') { return {}; @@ -307,58 +326,73 @@ export function createAgentMachine({ }, }).createMachine({ id: 'agent', - initial: 'idle', + initial: 'restoring', context: ({ input }) => ({ input, - messages: [...(input.history ?? [])], + messages: [], turnTools: {}, background: {}, scope: createAbortScope(), notifications: [], reminders: [], queue: [], - turnId: input.turnId ?? 0, - branchId: input.branchId ?? 'main', + turnId: 0, + branchId: 'main', }), - invoke: { - src: 'controllerGuard', - input: ({ context }) => ({ scope: context.scope }), - }, + invoke: [ + { + src: 'controllerGuard', + input: ({ context }) => ({ scope: context.scope }), + }, + { + id: 'store', + src: 'storeActor', + }, + ], on: { 'input.submit': { - actions: assign({ - queue: ({ context, event }) => [ - ...context.queue, - { id: event.id, message: event.message }, - ], - }), + actions: sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputSubmitted({ id: event.id, message: event.message }), + })), }, 'input.notify': { - actions: assign({ - notifications: ({ context, event }) => [ - ...context.notifications, - createUserEntry(event.message, { source: 'notify' }), - ], - }), + actions: sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputNotified({ message: event.message }), + })), }, 'input.remind': { - actions: assign({ - reminders: ({ context, event }) => [ - ...context.reminders.filter((entry) => entry.meta.key !== event.key), - event.message.role === 'system' - ? createSystemEntry(event.message, { source: 'reminder', key: event.key }) - : createUserEntry(event.message, { source: 'reminder', key: event.key }), - ], - }), + actions: sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputReminded({ key: event.key, message: event.message }), + })), }, 'input.steer': { - actions: assign(({ context, event }) => steerPatch(context, event.id)), + actions: enqueueActions(({ context, event, enqueue }) => { + const entry = context.queue.find((item) => item.id === event.id); + if (entry === undefined) return; + enqueue.sendTo('store', { + type: 'store.append' as const, + event: inputSteered({ id: event.id, message: entry.message }), + }); + }), }, 'input.cancel': { - actions: assign({ - queue: ({ context, event }) => - context.queue.filter((entry) => entry.id !== event.id), - }), + actions: sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputCancelled({ id: event.id }), + })), + }, + 'store.changed': { + actions: assign(({ event }) => mirrorPatch(event.state)), + }, + 'store.reset': { + target: '.idle', + actions: ['abortBackgroundScopes', 'resetMirror', 'emitReset', 'forwardToParent'], + }, + 'store.error': { + actions: 'forwardToParent', }, 'tool.update': { actions: [emit(({ event }) => event), 'forwardToParent'], @@ -381,24 +415,40 @@ export function createAgentMachine({ }, }, states: { + restoring: { + on: { + 'store.ready': { + target: 'idle', + actions: assign(({ event }) => ({ + ...mirrorPatch(event.state), + turnId: event.state.turnIndex.nextTurnId, + branchId: event.branch, + })), + }, + }, + }, idle: { initial: 'ready', always: { guard: ({ context }) => hasPendingWork(context), target: 'running', - actions: assign(({ context }) => drainPendingPatch(context)), - }, - on: { - 'context.reset': { - actions: [ - assign(({ context, event }) => ({ - messages: [...event.history], - branchId: event.branchId ?? context.branchId, - })), - emit(({ context }) => ({ type: 'context.reset' as const, branchId: context.branchId })), - 'forwardToParent', - ], - }, + actions: [ + sendTo('store', ({ context }) => { + const head = context.queue[0]; + return { + type: 'store.append' as const, + event: [ + ...context.notifications.map((entry) => messageAppended({ message: entry })), + ...(head === undefined + ? [] + : [messageAppended({ message: createUserEntry(head.message, { source: 'input' }) })]), + queueDrained({ id: head?.id }), + notificationsDrained({}), + ], + }; + }), + assign(({ context }) => drainPendingPatch(context)), + ], }, states: { ready: { @@ -425,26 +475,53 @@ export function createAgentMachine({ actions: [ assign(({ context, event }) => turnOutputPatch(context, event.output)), emit(({ context, event }) => turnOutcomeEvent(context, event.output)), + sendTo('store', ({ context, event }) => ({ + type: 'store.append' as const, + event: [ + ...event.output.produced.map((message) => messageAppended({ message })), + turnEnded({ + turnId: context.activeTurnId ?? context.turnId, + outcome: event.output.type, + errorMessage: + event.output.type === 'failed' ? String(event.output.error) : undefined, + }), + ], + })), ], }, onError: { target: '#agent.idle', - actions: emit(({ context, event }) => ({ - type: 'turn.failed' as const, - error: event.error, - messages: context.messages, - interruptReason: interruptReasonOf(event.error), - branchId: context.branchId, - })), + actions: [ + emit(({ context, event }) => ({ + type: 'turn.failed' as const, + error: event.error, + messages: context.messages, + interruptReason: interruptReasonOf(event.error), + branchId: context.branchId, + })), + sendTo('store', ({ context, event }) => ({ + type: 'store.append' as const, + event: turnEnded({ + turnId: context.activeTurnId ?? context.turnId, + outcome: 'failed', + errorMessage: String(event.error), + }), + })), + ], }, }, entry: [ + assign({ activeTurnId: ({ context }) => context.turnId }), emit(({ context }) => ({ type: 'turn.started' as const, turnId: context.turnId, branchId: context.branchId, queueItemId: context.drainedId, })), + sendTo('store', ({ context }) => ({ + type: 'store.append' as const, + event: turnStarted({ turnId: context.turnId, queueItemId: context.drainedId }), + })), ], exit: [ assign({ turnTools: {} }), @@ -452,6 +529,18 @@ export function createAgentMachine({ ], initial: 'active', on: { + 'store.reset': { + target: '#agent.idle', + actions: [ + 'abortTurn', + 'abortTurnTools', + 'stopTurnTools', + 'abortBackgroundScopes', + 'resetMirror', + 'emitReset', + 'forwardToParent', + ], + }, 'turn.drain': { actions: [ sendTo('turn', ({ context }) => ({ @@ -459,6 +548,10 @@ export function createAgentMachine({ messages: [...context.notifications, ...context.reminders], })), assign({ notifications: [], reminders: [] }), + sendTo('store', { + type: 'store.append' as const, + event: inputDrained({}), + }), ], }, 'tool.detached': { @@ -544,11 +637,7 @@ export function createAgentMachine({ }, 'input.abort': { target: 'aborting', - actions: [ - 'abortTurn', - 'abortTurnTools', - emit({ type: 'turn.aborting' as const }), - ], + actions: ['abortTurn', 'abortTurnTools', emit({ type: 'turn.aborting' as const })], }, }, }, diff --git a/packages/agent-core-v2/src/human/agent/replay.ts b/packages/agent-core-v2/src/human/agent/replay.ts deleted file mode 100644 index d410d1d6d77..00000000000 --- a/packages/agent-core-v2/src/human/agent/replay.ts +++ /dev/null @@ -1,40 +0,0 @@ -import type { HistoryMessage } from '#/agent/turn'; -import type { Tree } from '#/store/tree'; - -export type TurnOutcome = 'done' | 'failed' | 'aborted'; - -export type TurnEntryData = - | { phase: 'start'; turnId: number } - | { phase: 'end'; turnId?: number; outcome: TurnOutcome }; - -export interface StateEntryData { - name: string; - value: unknown; -} - -export interface LoadedAgentState { - messages: HistoryMessage[]; - turnId: number; - states: Record; -} - -export async function loadAgentState(tree: Tree, branch: string): Promise { - const loadedBranch = tree.openBranch(branch); - const entries = [...loadedBranch.walk()].toReversed(); - const messages: HistoryMessage[] = []; - let turnId = 0; - const states: Record = {}; - for (const entry of entries) { - const data = await tree.resolve(entry); - if (entry.type === 'message') { - messages.push(data as HistoryMessage); - } else if (entry.type === 'turn') { - const turn = data as TurnEntryData; - if (turn.phase === 'start') turnId = Math.max(turnId, turn.turnId + 1); - } else if (entry.type === 'state') { - const state = data as StateEntryData; - states[state.name] = state.value; - } - } - return { messages, turnId, states }; -} diff --git a/packages/agent-core-v2/src/human/agent/slices.ts b/packages/agent-core-v2/src/human/agent/slices.ts new file mode 100644 index 00000000000..82b211f2e3b --- /dev/null +++ b/packages/agent-core-v2/src/human/agent/slices.ts @@ -0,0 +1,131 @@ +import type { StoreActorLogic } from '#/eventStore/actor'; +import type { CombinedState, EventStore } from '#/eventStore/eventStore'; +import { createSlice } from '#/eventStore/slice'; +import type { BranchRef } from '#/store/types'; +import type { UserMessage } from '#/llm/message'; + +import { + inputCancelled, + inputDrained, + inputNotified, + inputReminded, + inputSteered, + inputSubmitted, + messageAppended, + notificationsDrained, + queueDrained, + turnEnded, + turnStarted, + type InputCancelled, + type InputNotified, + type InputReminded, + type InputSteered, + type InputSubmitted, + type MessageAppended, + type TurnEnded, + type TurnStarted, +} from './events'; +import { createSystemEntry, createUserEntry, type HistoryMessage, type UserEntry } from './turn'; + +export interface QueuedPrompt { + id?: string; + message: UserMessage; +} + +export const historySlice = createSlice({ + name: 'history', + initialState: () => [] as HistoryMessage[], + reducers: { + [messageAppended.type]: (draft, event: MessageAppended) => { + draft.push(event.message); + }, + }, +}); + +export const queueSlice = createSlice({ + name: 'queue', + initialState: () => [] as QueuedPrompt[], + reducers: { + [inputSubmitted.type]: (draft, event: InputSubmitted) => { + draft.push({ id: event.id, message: event.message }); + }, + [inputCancelled.type]: (draft, event: InputCancelled) => + draft.filter((entry) => entry.id !== event.id), + [inputSteered.type]: (draft, event: InputSteered) => + draft.filter((entry) => entry.id !== event.id), + [queueDrained.type]: (draft) => { + draft.shift(); + }, + }, +}); + +export const notificationsSlice = createSlice({ + name: 'notifications', + initialState: () => [] as UserEntry[], + reducers: { + [inputNotified.type]: (draft, event: InputNotified) => { + draft.push(createUserEntry(event.message, { source: event.source ?? 'notify' })); + }, + [inputSteered.type]: (draft, event: InputSteered) => { + draft.push(createUserEntry(event.message, { source: 'input' })); + }, + [inputDrained.type]: () => [], + [notificationsDrained.type]: () => [], + }, +}); + +export const remindersSlice = createSlice({ + name: 'reminders', + initialState: () => [] as HistoryMessage[], + reducers: { + [inputReminded.type]: (draft, event: InputReminded) => { + const kept = draft.filter((entry) => entry.meta.key !== event.key); + kept.push( + event.message.role === 'system' + ? createSystemEntry(event.message, { source: 'reminder', key: event.key }) + : createUserEntry(event.message, { source: 'reminder', key: event.key }), + ); + return kept; + }, + [inputDrained.type]: () => [], + }, +}); + +export interface TurnIndexEntry { + turnId: number; + start: BranchRef; + end?: BranchRef; +} + +export interface TurnIndexState { + turns: TurnIndexEntry[]; + nextTurnId: number; +} + +export const turnIndexSlice = createSlice({ + name: 'turnIndex', + initialState: (): TurnIndexState => ({ turns: [], nextTurnId: 0 }), + reducers: { + [turnStarted.type]: (draft, event: TurnStarted, ctx) => { + draft.turns.push({ turnId: event.turnId, start: ctx.ref }); + }, + [turnEnded.type]: (draft, event: TurnEnded, ctx) => { + const entry = draft.turns.findLast((turn) => turn.turnId === event.turnId); + if (entry !== undefined) entry.end = ctx.ref; + draft.nextTurnId = event.turnId + 1; + }, + }, +}); + +export const agentSlices = { + history: historySlice, + queue: queueSlice, + notifications: notificationsSlice, + reminders: remindersSlice, + turnIndex: turnIndexSlice, +}; + +export type AgentSlices = typeof agentSlices; +export type AgentEventStore = EventStore; +export type AgentStoreState = CombinedState; +export type AgentStoreActorLogic = StoreActorLogic; diff --git a/packages/agent-core-v2/src/human/eventStore/actor.ts b/packages/agent-core-v2/src/human/eventStore/actor.ts new file mode 100644 index 00000000000..79b28d8cfe1 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/actor.ts @@ -0,0 +1,41 @@ +import { fromCallback } from '#/xstate2'; + +import type { CombinedState, EventStore, SliceMap } from './eventStore'; + +export type StoreActorEvent = { + type: 'store.append'; + event: { type: string } & Record | readonly ({ type: string } & Record)[]; +}; + +export type StoreActorEmitted = + | { type: 'store.ready'; state: CombinedState; branch: string } + | { type: 'store.changed'; state: CombinedState } + | { type: 'store.reset'; state: CombinedState; branch: string } + | { type: 'store.error'; error: unknown }; + +export function fromEventStore(store: EventStore) { + return fromCallback>(({ emit, sendBack, receive }) => { + const publish = (event: StoreActorEmitted): void => { + sendBack(event); + emit(event); + }; + publish({ type: 'store.ready', state: store.getState(), branch: store.ref.branch }); + const unsubscribe = store.subscribe((state, cause) => { + if (cause.kind === 'reset') { + publish({ type: 'store.reset', state, branch: store.ref.branch }); + } else { + publish({ type: 'store.changed', state }); + } + }); + receive((event) => { + if (event.type === 'store.append') { + void store.dispatch(event.event).catch((error: unknown) => { + publish({ type: 'store.error', error }); + }); + } + }); + return unsubscribe; + }); +} + +export type StoreActorLogic = ReturnType>; diff --git a/packages/agent-core-v2/src/human/eventStore/eventStore.ts b/packages/agent-core-v2/src/human/eventStore/eventStore.ts new file mode 100644 index 00000000000..af87db77dc3 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/eventStore.ts @@ -0,0 +1,342 @@ +import { produce } from 'immer'; + +import type { BranchRef, EntryLine } from '#/store/types'; +import { StoreError } from '#/store/types'; + +import type { ExternalEvent, InternalEvent } from './events'; +import { eventSchemaFor, parseEvent, validateEvent } from './events'; +import type { JournalRecord, StoreJournal } from './journal'; +import type { FoldContext, Slice } from './slice'; + +export type SliceMap = Record>; + +export type CombinedState = { + readonly [K in keyof SM]: SM[K] extends Slice ? S : never; +}; + +export type Cause = + | { kind: 'event'; event: ExternalEvent; entry: EntryLine } + | { kind: 'internal'; event: InternalEvent } + | { kind: 'reset'; state: CombinedState } + | { kind: 'slice-joined'; name: string }; + +export interface EventStoreOptions { + journal: StoreJournal; + slices: SM; + snapshot?: { everyEvents?: number } | false; + drainLimit?: number; + onError?: (error: unknown) => void; +} + +export interface EventStore { + readonly ref: { tree: string; branch: string }; + readonly phase: 'open' | 'closed'; + + getState(): CombinedState; + slice(name: K): CombinedState[K]; + select(selector: (state: CombinedState) => T): T; + subscribe(listener: (state: CombinedState, cause: Cause) => void): () => void; + + dispatch(event: E | readonly E[]): Promise; + registerSlice(slice: Slice): Promise<() => void>; + reset(journal: StoreJournal): Promise; + flush(): Promise; + close(): Promise; +} + +const SNAPSHOT_ENTRY_TYPE = 'snapshot'; +const SNAPSHOT_ENTRY_KIND = 'snapshot'; +const EVENT_ENTRY_KIND = 'event'; +const DEFAULT_SNAPSHOT_EVERY_EVENTS = 500; +const DEFAULT_DRAIN_LIMIT = 100; + +type Listener = (state: CombinedState, cause: Cause) => void; + +export async function createEventStore( + opts: EventStoreOptions, +): Promise> { + const store = new EventStoreImpl(opts); + await store.refold(opts.journal); + return store; +} + +class EventStoreImpl implements EventStore { + private journal: StoreJournal; + private slices: SliceMap; + private state: Record; + private phaseValue: 'open' | 'closed' = 'open'; + private tail: Promise = Promise.resolve(); + private eventsSinceSnapshot = 0; + private readonly snapshotEvery: number | false; + private readonly drainLimit: number; + private readonly report: (error: unknown) => void; + private readonly listeners = new Set>(); + + constructor(opts: EventStoreOptions) { + this.journal = opts.journal; + this.slices = { ...opts.slices }; + this.state = {}; + this.snapshotEvery = + opts.snapshot === false ? false : (opts.snapshot?.everyEvents ?? DEFAULT_SNAPSHOT_EVERY_EVENTS); + this.drainLimit = opts.drainLimit ?? DEFAULT_DRAIN_LIMIT; + this.report = opts.onError ?? ((error) => console.error(error)); + } + + get ref(): { tree: string; branch: string } { + return this.journal.ref; + } + + get phase(): 'open' | 'closed' { + return this.phaseValue; + } + + getState(): CombinedState { + return this.state as CombinedState; + } + + slice(name: K): CombinedState[K] { + return this.state[name as string] as CombinedState[K]; + } + + select(selector: (state: CombinedState) => T): T { + return selector(this.getState()); + } + + subscribe(listener: Listener): () => void { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + } + + dispatch(event: E | readonly E[]): Promise { + if (this.phaseValue !== 'open') { + return Promise.reject(new StoreError('closed', 'store is closed')); + } + const events = (Array.isArray(event) ? event : [event]) as readonly E[]; + for (const item of events) { + const invalid = validateEvent(item); + if (invalid !== undefined) { + return Promise.reject(invalid); + } + } + const result = this.tail.then(() => this.foldAndAppend(events)); + this.tail = result.then(noop, noop); + return result; + } + + async registerSlice(slice: Slice): Promise<() => void> { + if (this.phaseValue !== 'open') { + throw new StoreError('closed', 'store is closed'); + } + if (this.slices[slice.name] !== undefined) { + throw new StoreError('duplicate-slice', `slice '${slice.name}' is already registered`); + } + const op = this.tail.then(async () => { + this.slices = { ...this.slices, [slice.name]: slice }; + await this.refold(this.journal); + }); + this.tail = op.then(noop, noop); + await op; + this.notify([{ kind: 'slice-joined', name: slice.name }]); + return () => { + const slices = { ...this.slices }; + delete slices[slice.name]; + this.slices = slices; + const state = { ...this.state }; + delete state[slice.name]; + this.state = state; + }; + } + + reset(journal: StoreJournal): Promise { + if (this.phaseValue !== 'open') { + return Promise.reject(new StoreError('closed', 'store is closed')); + } + const op = this.tail.then(async () => { + await this.journal.settled(); + this.journal = journal; + await this.refold(journal); + }); + this.tail = op.then(noop, noop); + return op.then(() => { + this.notify([{ kind: 'reset', state: this.getState() }]); + }); + } + + flush(): Promise { + return this.tail.then(() => this.journal.settled()); + } + + async close(): Promise { + await this.flush(); + this.phaseValue = 'closed'; + this.listeners.clear(); + } + + async refold(journal: StoreJournal): Promise { + const records: JournalRecord[] = []; + for await (const record of journal.read()) { + records.push(record); + } + let snapshotIndex = -1; + for (let i = records.length - 1; i >= 0; i--) { + if (records[i]?.kind === SNAPSHOT_ENTRY_KIND) { + snapshotIndex = i; + break; + } + } + const saved = + snapshotIndex >= 0 + ? (((records[snapshotIndex] as JournalRecord).data as { slices?: Record }) + .slices ?? {}) + : {}; + const seeded: Record = {}; + for (const [name, slice] of Object.entries(this.slices)) { + seeded[name] = + name in saved + ? slice.deserialize !== undefined + ? slice.deserialize(saved[name]) + : saved[name] + : slice.initialState(); + } + this.state = seeded; + let sinceSnapshot = 0; + for (let i = snapshotIndex + 1; i < records.length; i++) { + const record = records[i] as JournalRecord; + if (record.kind !== EVENT_ENTRY_KIND) continue; + sinceSnapshot += 1; + this.replayRecord(record); + } + this.eventsSinceSnapshot = sinceSnapshot; + } + + private replayRecord(record: JournalRecord): void { + if (eventSchemaFor(record.type) === undefined) return; + const event = parseEvent(record.type, record.data); + if (event === undefined) { + this.report( + new StoreError('schema', `event '${record.type}' at seq ${record.seq} failed schema validation`), + ); + return; + } + const ref: BranchRef = { branch: record.branch, seq: record.seq }; + const { raised } = this.applyEvent(event, ref, record.ts, true); + this.drain(raised, ref, record.ts, true); + } + + private async foldAndAppend(events: readonly ExternalEvent[]): Promise { + const causes: Cause[] = []; + const entries: EntryLine[] = []; + for (const event of events) { + const ts = typeof event.time === 'number' ? event.time : Date.now(); + const ref: BranchRef = { branch: this.journal.ref.branch, seq: this.journal.nextSeq() }; + const { raised, effects } = this.applyEvent(event, ref, ts, false); + const internalCauses = this.drain(raised, ref, ts, false); + for (const effect of effects) { + try { + effect(); + } catch (error) { + this.report(error); + } + } + const entry = await this.journal.append({ type: event.type, kind: EVENT_ENTRY_KIND, data: event }); + entries.push(entry); + causes.push({ kind: 'event', event, entry }, ...internalCauses); + this.maybeSnapshot(); + } + this.notify(causes); + return entries[entries.length - 1] as EntryLine; + } + + private applyEvent( + event: { type: string }, + ref: BranchRef, + ts: number, + replaying: boolean, + ): { raised: InternalEvent[]; effects: (() => void)[] } { + const raised: InternalEvent[] = []; + const effects: (() => void)[] = []; + const ctx: FoldContext = { + ref, + ts, + replaying, + enqueue: { + raise: (internal) => { + raised.push(internal); + }, + effect: (fn) => { + effects.push(fn); + }, + }, + }; + let changed = false; + const next: Record = { ...this.state }; + for (const [name, slice] of Object.entries(this.slices)) { + const reducer = slice.reducers[event.type]; + if (reducer === undefined) continue; + changed = true; + next[name] = produce(next[name], (draft) => reducer(draft, event, ctx)); + } + if (changed) { + this.state = next; + } + return { raised, effects }; + } + + private drain( + initial: InternalEvent[], + ref: BranchRef, + ts: number, + replaying: boolean, + ): Cause[] { + const causes: Cause[] = []; + const queue = [...initial]; + let count = 0; + while (queue.length > 0) { + count += 1; + if (count > this.drainLimit) { + throw new StoreError('drain-limit', `internal event drain exceeded limit ${this.drainLimit}`); + } + const internal = queue.shift() as InternalEvent; + const { raised, effects } = this.applyEvent(internal, ref, ts, replaying); + queue.push(...raised); + if (!replaying) { + for (const effect of effects) { + try { + effect(); + } catch (error) { + this.report(error); + } + } + causes.push({ kind: 'internal', event: internal }); + } + } + return causes; + } + + private maybeSnapshot(): void { + if (this.snapshotEvery === false) return; + this.eventsSinceSnapshot += 1; + if (this.eventsSinceSnapshot < this.snapshotEvery) return; + this.eventsSinceSnapshot = 0; + const slices: Record = {}; + for (const [name, slice] of Object.entries(this.slices)) { + slices[name] = slice.serialize !== undefined ? slice.serialize(this.state[name]) : this.state[name]; + } + void this.journal + .append({ type: SNAPSHOT_ENTRY_TYPE, kind: SNAPSHOT_ENTRY_KIND, data: { slices } }) + .catch((error) => this.report(error)); + } + + private notify(causes: Cause[]): void { + const state = this.getState(); + for (const cause of causes) { + for (const listener of this.listeners) { + listener(state, cause); + } + } + } +} + +function noop(): void {} diff --git a/packages/agent-core-v2/src/human/eventStore/events.ts b/packages/agent-core-v2/src/human/eventStore/events.ts new file mode 100644 index 00000000000..0f82bda31f8 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/events.ts @@ -0,0 +1,61 @@ +import { z } from 'zod'; + +import { StoreError } from '#/store/types'; + +export type ExternalEvent

> = P & { type: string; time: number }; + +export interface EventFactory

{ + (payload: P): ExternalEvent

; + readonly type: string; + readonly schema: z.ZodTypeAny; +} + +export type EventOf = F extends EventFactory ? ExternalEvent

: never; + +export type InternalEvent = { type: string } & Record; + +const registry = new Map(); + +export function defineEvent

(def: { type: string; schema: z.ZodType

}): EventFactory

{ + if (registry.has(def.type)) { + throw new StoreError('duplicate-event', `duplicate event type '${def.type}'`); + } + registry.set(def.type, def.schema); + const factory = (payload: P): ExternalEvent

=> ({ + ...payload, + type: def.type, + time: Date.now(), + }); + factory.type = def.type; + factory.schema = def.schema; + return factory; +} + +export function eventSchemaFor(type: string): z.ZodTypeAny | undefined { + return registry.get(type); +} + +export function parseEvent( + type: string, + record: unknown, +): (ExternalEvent & Record) | undefined { + const schema = registry.get(type); + if (schema === undefined) return undefined; + if (typeof record !== 'object' || record === null) return undefined; + const { type: _type, time: _time, ...payload } = record as Record; + if (!schema.safeParse(payload).success) return undefined; + return record as ExternalEvent & Record; +} + +export function validateEvent(event: ExternalEvent): StoreError | undefined { + const schema = registry.get(event.type); + if (schema === undefined) { + return new StoreError('unregistered-event', `event '${event.type}' is not a registered external event`); + } + const { type: _type, time: _time, ...payload } = event; + const parsed = schema.safeParse(payload); + if (!parsed.success) { + return new StoreError('schema', `event '${event.type}' failed schema validation: ${parsed.error.message}`); + } + return undefined; +} diff --git a/packages/agent-core-v2/src/human/eventStore/index.ts b/packages/agent-core-v2/src/human/eventStore/index.ts new file mode 100644 index 00000000000..66ff89537b5 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/index.ts @@ -0,0 +1,5 @@ +export * from './events'; +export * from './slice'; +export * from './journal'; +export * from './eventStore'; +export * from './actor'; diff --git a/packages/agent-core-v2/src/human/eventStore/journal.ts b/packages/agent-core-v2/src/human/eventStore/journal.ts new file mode 100644 index 00000000000..1bee1792124 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/journal.ts @@ -0,0 +1,104 @@ +import type { Branch } from '#/store/branch'; +import type { Tree } from '#/store/tree'; +import type { AppendInput, EntryLine } from '#/store/types'; +import { isOffloadedPayload } from '#/store/types'; + +export interface JournalRecord { + branch: string; + seq: number; + ts: number; + type: string; + kind: string; + data: unknown; +} + +export interface StoreJournal { + readonly ref: { tree: string; branch: string }; + append(input: AppendInput): Promise; + read(): AsyncIterable; + nextSeq(): number; + settled(): Promise; +} + +export interface SyncStoreJournal extends StoreJournal { + readSync(): JournalRecord[]; +} + +export function memoryJournal(ref?: { tree: string; branch: string }): SyncStoreJournal { + const journalRef = ref ?? { tree: 'memory', branch: 'main' }; + const entries: EntryLine[] = []; + const records = (): JournalRecord[] => + entries.map((entry) => ({ + branch: journalRef.branch, + seq: entry.seq, + ts: entry.ts, + type: entry.type, + kind: entry.payload.kind, + data: isOffloadedPayload(entry.payload) ? null : entry.payload.data, + })); + return { + ref: journalRef, + append: (input) => { + const data = input.data ?? null; + const entry: EntryLine = { + kind: 'entry', + seq: entries.length, + ts: Date.now(), + type: input.type, + payload: { kind: input.kind, size: JSON.stringify(data).length, data }, + }; + entries.push(entry); + return Promise.resolve(entry); + }, + read: async function* () { + for (const record of records()) yield record; + }, + readSync: records, + nextSeq: () => entries.length, + settled: () => Promise.resolve(), + }; +} + +export function journalFromBranch(branch: Branch, tree: Tree): StoreJournal { + return { + ref: { tree: branch.tree, branch: branch.name }, + append: (input) => branch.append(input), + settled: () => branch.settled(), + read: () => readBranchChain(branch, tree), + nextSeq: () => branch.nextSeq, + }; +} + +async function* readBranchChain(branch: Branch, tree: Tree): AsyncIterable { + const chain: { name: string; entries: EntryLine[] }[] = []; + let current: Branch | undefined = branch; + let upto: number | null = null; + while (current !== undefined) { + const head = upto ?? current.head; + const entries: EntryLine[] = []; + for (let seq = 0; seq <= (head ?? -1); seq++) { + const entry = current.entryAt(seq); + if (entry !== null) entries.push(entry); + } + chain.push({ name: current.name, entries }); + const parentBranch: string | undefined = current.header.parentBranch; + const parentSeq: number | undefined = current.header.parentSeq; + current = + parentBranch !== undefined && parentSeq !== undefined && tree.has(parentBranch) + ? tree.openBranch(parentBranch) + : undefined; + upto = parentSeq ?? null; + } + for (const segment of chain.reverse()) { + for (const entry of segment.entries) { + yield { + branch: segment.name, + seq: entry.seq, + ts: entry.ts, + type: entry.type, + kind: entry.payload.kind, + data: await tree.resolve(entry), + }; + } + } +} diff --git a/packages/agent-core-v2/src/human/eventStore/slice.ts b/packages/agent-core-v2/src/human/eventStore/slice.ts new file mode 100644 index 00000000000..6a80a01e1e4 --- /dev/null +++ b/packages/agent-core-v2/src/human/eventStore/slice.ts @@ -0,0 +1,27 @@ +import type { BranchRef } from '#/store/types'; + +import type { InternalEvent } from './events'; + +export interface Enqueue { + raise(event: InternalEvent): void; + effect(fn: () => void): void; +} + +export interface FoldContext { + readonly ref: BranchRef; + readonly ts: number; + readonly replaying: boolean; + readonly enqueue: Enqueue; +} + +export interface Slice { + readonly name: Name; + readonly initialState: () => S; + readonly reducers: Record void | S>; + readonly serialize?: (state: S) => unknown; + readonly deserialize?: (raw: unknown) => S; +} + +export function createSlice(def: Slice): Slice { + return def; +} diff --git a/packages/agent-core-v2/src/human/index.ts b/packages/agent-core-v2/src/human/index.ts index 5f908ec4fe3..79e1b9c6362 100644 --- a/packages/agent-core-v2/src/human/index.ts +++ b/packages/agent-core-v2/src/human/index.ts @@ -60,13 +60,16 @@ export * from './agent/machine'; export * from './agent/wait-for'; export * from './agent/turn'; export * from './agent/context-usage'; -export * from './agent/replay'; -export * from './persist/agent'; +export * from './agent/events'; +export * from './agent/slices'; +export * from './agent/historySchema'; export * from './persist/open'; +export * from './eventStore/index'; export * from './persist/v2/migrate'; export * from './session/machine'; -export * from './session/undo'; -export * from './persist/session'; +export * from './session/events'; +export * from './session/slices'; +export * from './session/stores'; export * from './usage/usage'; export * from './usage/machine'; export * from './usage/plugin'; @@ -76,7 +79,7 @@ export * from './interaction/machine'; export * from './interaction/facade'; export * from './kimi/trace'; export * from './todo/todoItem'; -export * from './todo/state'; +export * from './todo/slice'; export * from './todo/tool'; export * from './todo/plugin'; export * from './tool-select/state'; diff --git a/packages/agent-core-v2/src/human/persist/agent.ts b/packages/agent-core-v2/src/human/persist/agent.ts deleted file mode 100644 index fc72fa05482..00000000000 --- a/packages/agent-core-v2/src/human/persist/agent.ts +++ /dev/null @@ -1,134 +0,0 @@ -import type { ActorRefFrom } from 'xstate'; - -import type { HistoryMessage } from '#/agent/turn'; -import type { TurnEntryData, TurnOutcome } from '#/agent/replay'; -import type { Branch } from '#/store/branch'; -import type { AppendInput } from '#/store/types'; - -import type { createAgentMachine } from '#/agent/machine'; - -type AgentActor = ActorRefFrom>; -type AgentSnapshot = ReturnType; - -export interface PersistAgentOptions { - states?: Record unknown>; - persistedMessages?: number; - onError?: (error: unknown) => void; -} - -export interface AgentPersistence { - readonly persistedMessages: number; - flush(): Promise; - dispose(): void; -} - -interface BufferedTurnEvent { - kind: 'start' | 'end'; - turnId?: number; - outcome?: TurnOutcome; - messageCount?: number; -} - -function countPersistedMessages(branch: Branch): number { - let count = 0; - for (const entry of branch.walk()) { - if (entry.type === 'message') count += 1; - } - return count; -} - -export function persistAgent(actor: AgentActor, branch: Branch, opts?: PersistAgentOptions): AgentPersistence { - const report = opts?.onError ?? ((error: unknown) => console.error(error)); - const write = (input: AppendInput): void => { - void branch.append(input).catch(report); - }; - const writeMessage = (message: HistoryMessage): void => { - void branch.append({ type: 'message', kind: 'agent', data: message }).then( - () => { - completed += 1; - }, - report, - ); - }; - - let queued = opts?.persistedMessages ?? countPersistedMessages(branch); - let completed = queued; - let currentTurnId: number | undefined; - const buffered: BufferedTurnEvent[] = []; - const stateValues = new Map(); - for (const [name, get] of Object.entries(opts?.states ?? {})) { - stateValues.set(name, JSON.stringify(get())); - } - - const subscriptions = [ - actor.on('turn.started', (event) => { - if (event.type !== 'turn.started') return; - currentTurnId = event.turnId; - buffered.push({ kind: 'start', turnId: event.turnId }); - }), - actor.on('turn.done', (event) => { - if (event.type !== 'turn.done') return; - buffered.push({ kind: 'end', outcome: 'done', messageCount: event.messages.length }); - }), - actor.on('turn.failed', (event) => { - if (event.type !== 'turn.failed') return; - buffered.push({ kind: 'end', outcome: 'failed', messageCount: event.messages.length }); - }), - actor.on('turn.aborted', (event) => { - if (event.type !== 'turn.aborted') return; - buffered.push({ kind: 'end', outcome: 'aborted', messageCount: event.messages.length }); - }), - ]; - - const flushSnapshot = (snapshot: AgentSnapshot): void => { - const context = snapshot.context; - const writeUpTo = (count: number): void => { - while (queued < count) { - writeMessage(context.messages[queued] as HistoryMessage); - queued += 1; - } - }; - for (const event of buffered) { - if (event.kind === 'end') { - writeUpTo(event.messageCount ?? queued); - const data: TurnEntryData = { - phase: 'end', - turnId: currentTurnId, - outcome: event.outcome as TurnOutcome, - }; - write({ type: 'turn', kind: 'agent', data }); - } else { - writeUpTo(context.messages.length); - const data: TurnEntryData = { phase: 'start', turnId: event.turnId as number }; - write({ type: 'turn', kind: 'agent', data }); - } - } - buffered.length = 0; - writeUpTo(context.messages.length); - for (const [name, get] of Object.entries(opts?.states ?? {})) { - const value = get(); - const json = JSON.stringify(value); - if (json !== stateValues.get(name)) { - stateValues.set(name, json); - write({ type: 'state', kind: 'agent', data: { name, value } }); - } - } - }; - - const subscription = actor.subscribe((snapshot) => { - flushSnapshot(snapshot); - }); - - return { - get persistedMessages() { - return completed; - }, - flush: () => branch.settled(), - dispose: () => { - subscription.unsubscribe(); - for (const sub of subscriptions) { - sub.unsubscribe(); - } - }, - }; -} diff --git a/packages/agent-core-v2/src/human/persist/open.ts b/packages/agent-core-v2/src/human/persist/open.ts index 76d0ae91548..7954cfc96ad 100644 --- a/packages/agent-core-v2/src/human/persist/open.ts +++ b/packages/agent-core-v2/src/human/persist/open.ts @@ -2,6 +2,8 @@ import { NodeBackend } from '#/store/backend/node'; import { TreeStore } from '#/store/store'; import type { Tree } from '#/store/tree'; +import { SessionStores } from '#/session/stores'; + import { isV2SessionDir, migrateV2Session, V2_SESSION_TREE_NAME } from './v2/migrate'; export interface OpenSessionStoreOptions { @@ -12,6 +14,7 @@ export interface OpenSessionStoreOptions { export interface OpenedSessionStore { store: TreeStore; tree: Tree; + stores: SessionStores; migrated: boolean; } @@ -24,7 +27,8 @@ export async function openSessionStore( await migrateV2Session(dir); migrated = true; } - const store = await TreeStore.open(new NodeBackend(dir), { fsync: opts?.fsync ?? false }); + const backend = new NodeBackend(dir); + const store = await TreeStore.open(backend, { fsync: opts?.fsync ?? false }); const tree = await store.tree(opts?.treeName ?? V2_SESSION_TREE_NAME); - return { store, tree, migrated }; + return { store, tree, stores: new SessionStores(tree, backend), migrated }; } diff --git a/packages/agent-core-v2/src/human/persist/session.ts b/packages/agent-core-v2/src/human/persist/session.ts deleted file mode 100644 index 9074d99ad70..00000000000 --- a/packages/agent-core-v2/src/human/persist/session.ts +++ /dev/null @@ -1,207 +0,0 @@ -import type { ActorRefFrom } from '#/xstate2'; - -import { loadAgentState, type LoadedAgentState } from '#/agent/replay'; -import { persistAgent, type AgentPersistence } from '#/persist/agent'; -import type { Branch } from '#/store/branch'; -import type { Tree } from '#/store/tree'; - -import type { AgentActorRef, createSessionMachine } from '#/session/machine'; - -type SessionActor = ActorRefFrom>; - -export interface PersistSessionOptions { - states?: (agentId: string) => Record unknown>; - branches?: ReadonlyMap; - onError?: (error: unknown) => void; -} - -export interface SessionPersistence { - flush(): Promise; - dispose(): void; -} - -export interface LoadedSessionAgent extends LoadedAgentState { - agentId: string; - branch: string; -} - -export interface LoadedSessionState { - agents: LoadedSessionAgent[]; - meta?: unknown; -} - -interface AgentHandle { - persistence: AgentPersistence; - branch: Branch; -} - -export const SESSION_LOG_BRANCH = '_session'; -export const SESSION_AGENT_OPEN_ENTRY_TYPE = 'agent.open'; -export const SESSION_META_ENTRY_TYPE = 'session.meta'; -const CLOSE_ENTRY_TYPE = 'agent.close'; -const SWITCH_ENTRY_TYPE = 'agent.switch'; - -export function persistSession( - session: SessionActor, - tree: Tree, - opts?: PersistSessionOptions, -): SessionPersistence { - const report = opts?.onError ?? ((error: unknown) => console.error(error)); - const handles = new Map(); - let logBranch: Branch | undefined; - - const log = (): Branch => { - logBranch ??= tree.has(SESSION_LOG_BRANCH) - ? tree.openBranch(SESSION_LOG_BRANCH) - : tree.createBranch(SESSION_LOG_BRANCH); - return logBranch; - }; - - const attach = ( - agentId: string, - ref: AgentActorRef, - branch: Branch, - persistedMessages?: number, - ): void => { - const persistence = persistAgent(ref, branch, { - states: opts?.states?.(agentId), - onError: report, - persistedMessages, - }); - handles.set(agentId, { persistence, branch }); - }; - - const attachCreated = (agentId: string, branchId: string, ref: AgentActorRef): void => { - const mapped = opts?.branches?.get(agentId); - if (mapped !== undefined) { - attach(agentId, ref, tree.openBranch(mapped)); - return; - } - if (tree.has(branchId)) { - report(new Error(`branch '${branchId}' already exists for agent '${agentId}'`)); - return; - } - const branch = tree.createBranch(branchId); - void log() - .append({ type: SESSION_AGENT_OPEN_ENTRY_TYPE, kind: 'session', data: { agentId, branch: branch.name } }) - .catch(report); - attach(agentId, ref, branch); - }; - - const attachForked = ( - sourceId: string, - agentId: string, - branchId: string, - ref: AgentActorRef, - ): void => { - const source = handles.get(sourceId); - if (source === undefined) { - report(new Error(`cannot persist fork of unknown agent '${sourceId}'`)); - return; - } - if (tree.has(branchId)) { - report(new Error(`branch '${branchId}' already exists for agent '${agentId}'`)); - return; - } - const head = source.branch.head; - const branch = - head === null - ? tree.createBranch(branchId) - : tree.createBranch(branchId, { from: { branch: source.branch.name, seq: head } }); - void log() - .append({ type: SESSION_AGENT_OPEN_ENTRY_TYPE, kind: 'session', data: { agentId, branch: branch.name } }) - .catch(report); - attach(agentId, ref, branch, source.persistence.persistedMessages); - }; - - const attachSwitched = (agentId: string, branchId: string, reason?: string): void => { - const entry = session.getSnapshot().context.agents[agentId]; - if (entry === undefined) return; - const previous = handles.get(agentId); - const from = - previous === undefined - ? null - : { branch: previous.branch.name, seq: previous.branch.head }; - previous?.persistence.dispose(); - attach(agentId, entry.ref, tree.openBranch(branchId)); - void log() - .append({ - type: SWITCH_ENTRY_TYPE, - kind: 'session', - data: { agentId, branch: branchId, reason, from }, - }) - .catch(report); - }; - - const detach = (agentId: string): void => { - const handle = handles.get(agentId); - if (handle === undefined) return; - handles.delete(agentId); - handle.persistence.dispose(); - void log().append({ type: CLOSE_ENTRY_TYPE, kind: 'session', data: { agentId } }).catch(report); - }; - - for (const [agentId, entry] of Object.entries(session.getSnapshot().context.agents)) { - attachCreated(agentId, entry.ref.getSnapshot().context.branchId, entry.ref); - } - - const subscriptions = [ - session.on('agent.created', (event) => { - if (event.type !== 'agent.created') return; - attachCreated(event.agentId, event.branchId, event.ref); - }), - session.on('agent.forked', (event) => { - if (event.type !== 'agent.forked') return; - attachForked(event.sourceId, event.agentId, event.branchId, event.ref); - }), - session.on('agent.switched', (event) => { - if (event.type !== 'agent.switched') return; - attachSwitched(event.agentId, event.branchId, event.reason); - }), - session.on('agent.stopped', (event) => { - if (event.type !== 'agent.stopped') return; - detach(event.agentId); - }), - ]; - - return { - flush: async () => { - await Promise.all([...handles.values()].map((handle) => handle.persistence.flush())); - await logBranch?.settled(); - }, - dispose: () => { - for (const subscription of subscriptions) { - subscription.unsubscribe(); - } - for (const handle of handles.values()) { - handle.persistence.dispose(); - } - handles.clear(); - }, - }; -} - -export async function loadSessionState(tree: Tree): Promise { - const agents: LoadedSessionAgent[] = []; - if (!tree.has(SESSION_LOG_BRANCH)) return { agents }; - const open = new Map(); - let meta: unknown; - const entries = [...tree.openBranch(SESSION_LOG_BRANCH).walk()].toReversed(); - for (const entry of entries) { - const data = await tree.resolve(entry); - if (entry.type === SESSION_AGENT_OPEN_ENTRY_TYPE || entry.type === SWITCH_ENTRY_TYPE) { - const opened = data as { agentId: string; branch: string }; - open.set(opened.agentId, opened.branch); - } else if (entry.type === CLOSE_ENTRY_TYPE) { - const closed = data as { agentId: string }; - open.delete(closed.agentId); - } else if (entry.type === SESSION_META_ENTRY_TYPE) { - meta = data; - } - } - for (const [agentId, branch] of [...open.entries()].sort(([a], [b]) => a.localeCompare(b))) { - const loaded = await loadAgentState(tree, branch); - agents.push({ agentId, branch, ...loaded }); - } - return meta === undefined ? { agents } : { agents, meta }; -} diff --git a/packages/agent-core-v2/src/human/persist/v2/migrate.ts b/packages/agent-core-v2/src/human/persist/v2/migrate.ts index e8d7d86c4e9..9e5bc238d33 100644 --- a/packages/agent-core-v2/src/human/persist/v2/migrate.ts +++ b/packages/agent-core-v2/src/human/persist/v2/migrate.ts @@ -1,7 +1,8 @@ import { readdir, readFile, rename, rm, stat } from 'node:fs/promises'; import { basename, join } from 'node:path'; -import { SESSION_AGENT_OPEN_ENTRY_TYPE, SESSION_LOG_BRANCH, SESSION_META_ENTRY_TYPE } from '#/persist/session'; +import { messageAppended, stateUpdated, turnEnded } from '#/agent/events'; +import { agentOpened, SESSION_LOG_BRANCH, sessionMetaUpdated } from '#/session/events'; import { NodeBackend } from '#/store/backend/node'; import { TreeStore } from '#/store/store'; import type { Branch } from '#/store/branch'; @@ -174,33 +175,37 @@ export async function migrateV2Session(dir: string): Promise resolveBlob, ); if (converted === null) continue; - await branch.append({ type: 'message', kind: 'agent', data: converted }); + await branch.append({ + type: messageAppended.type, + kind: 'event', + data: messageAppended({ message: converted }), + }); } const lastTurnId = folded.nextTurnId - 1; if (folded.todos.length > 0) { await branch.append({ - type: 'state', - kind: 'agent', - data: { + type: stateUpdated.type, + kind: 'event', + data: stateUpdated({ name: 'todo', value: { todos: folded.todos, lastWriteTurn: Math.max(0, lastTurnId) }, - }, + }), }); } if (folded.nextTurnId > 0) { await branch.append({ - type: 'turn', - kind: 'agent', - data: { phase: 'start', turnId: lastTurnId }, + type: turnEnded.type, + kind: 'event', + data: turnEnded({ turnId: lastTurnId, outcome: 'done' }), }); } await log.append({ - type: SESSION_AGENT_OPEN_ENTRY_TYPE, - kind: 'session', - data: { agentId, branch: agentId }, + type: agentOpened.type, + kind: 'event', + data: agentOpened({ agentId, branch: agentId }), }); } - await log.append({ type: SESSION_META_ENTRY_TYPE, kind: 'session', data: meta }); + await log.append({ type: sessionMetaUpdated.type, kind: 'event', data: sessionMetaUpdated({ meta }) }); for (const branchName of tree.branches()) { await tree.openBranch(branchName).settled(); } diff --git a/packages/agent-core-v2/src/human/session/events.ts b/packages/agent-core-v2/src/human/session/events.ts new file mode 100644 index 00000000000..a40131d2057 --- /dev/null +++ b/packages/agent-core-v2/src/human/session/events.ts @@ -0,0 +1,29 @@ +import { z } from 'zod'; + +import { defineEvent } from '#/eventStore/events'; + +export const SESSION_LOG_BRANCH = '_session'; + +export const agentOpened = defineEvent({ + type: 'agent.opened', + schema: z.object({ agentId: z.string(), branch: z.string() }), +}); +export type AgentOpened = ReturnType; + +export const agentClosed = defineEvent({ + type: 'agent.closed', + schema: z.object({ agentId: z.string() }), +}); +export type AgentClosed = ReturnType; + +export const agentSwitched = defineEvent({ + type: 'agent.switched', + schema: z.object({ agentId: z.string(), branch: z.string(), reason: z.string().optional() }), +}); +export type AgentSwitched = ReturnType; + +export const sessionMetaUpdated = defineEvent({ + type: 'session.meta_updated', + schema: z.object({ meta: z.unknown() }), +}); +export type SessionMetaUpdated = ReturnType; diff --git a/packages/agent-core-v2/src/human/session/machine.ts b/packages/agent-core-v2/src/human/session/machine.ts index 8eec007eceb..8106f5db54f 100644 --- a/packages/agent-core-v2/src/human/session/machine.ts +++ b/packages/agent-core-v2/src/human/session/machine.ts @@ -1,8 +1,9 @@ import { assign, emit, sendTo, setup, stopChild, type ActorRefFrom } from '#/xstate2'; +import type { createAgentMachine, AgentEvent } from '#/agent/machine'; +import type { AgentEventStore } from '#/agent/slices'; import type { LlmRequestConfig } from '#/llm/requester/requester'; -import type { createAgentMachine, AgentEvent, AgentInput } from '#/agent/machine'; -import type { HistoryMessage, TurnLlmEvent, TurnToolEvent } from '#/agent/turn'; +import type { TurnLlmEvent, TurnToolEvent } from '#/agent/turn'; import type { ToolUpdate } from '#/tool/executor'; export interface SessionInput { @@ -19,20 +20,18 @@ export type SessionEvent = | TurnLlmEvent | TurnToolEvent | { type: 'tool.update'; toolCallId: string; update: ToolUpdate } - | { type: 'agent.create'; agentId?: string; input?: Pick } - | { type: 'agent.fork'; sourceId: string; agentId?: string } | { - type: 'agent.switch'; - agentId: string; - input: { branchId: string; history: readonly HistoryMessage[]; reason?: string }; + type: 'agent.create'; + agentId?: string; + input: { request?: LlmRequestConfig; store: AgentEventStore }; } + | { type: 'agent.fork'; sourceId: string; agentId?: string; store: AgentEventStore } | { type: 'agent.send'; agentId: string; event: AgentEvent } | { type: 'agent.stop'; agentId: string }; export type SessionEmitted = | { type: 'agent.created'; agentId: string; branchId: string; ref: AgentActorRef } | { type: 'agent.forked'; sourceId: string; agentId: string; branchId: string; ref: AgentActorRef } - | { type: 'agent.switched'; agentId: string; branchId: string; reason?: string } | { type: 'agent.stopped'; agentId: string } | { type: 'agent.failed'; agentId: string; error: string }; @@ -86,6 +85,8 @@ export function createSessionMachine({ agent }: CreateSessionMachineOptions) { 'tool.failed': {}, 'tool.aborted': {}, 'context.reset': {}, + 'store.reset': {}, + 'store.error': {}, }, states: { active: { @@ -111,10 +112,8 @@ export function createSessionMachine({ agent }: CreateSessionMachineOptions) { const ref = spawn('agentActor', { id: agentId, input: { - request: event.input?.request ?? context.input.request, - history: event.input?.history, - turnId: event.input?.turnId, - branchId: event.input?.branchId ?? agentId, + request: event.input.request ?? context.input.request, + store: event.input.store, }, }); return { @@ -128,7 +127,7 @@ export function createSessionMachine({ agent }: CreateSessionMachineOptions) { return { type: 'agent.created' as const, agentId, - branchId: event.input?.branchId ?? agentId, + branchId: event.input.store.ref.branch, ref: entry.ref, }; }), @@ -162,9 +161,7 @@ export function createSessionMachine({ agent }: CreateSessionMachineOptions) { id: agentId, input: { request: source.context.input.request, - history: [...source.context.messages], - turnId: source.context.turnId, - branchId: agentId, + store: event.store, }, }); return { @@ -179,50 +176,13 @@ export function createSessionMachine({ agent }: CreateSessionMachineOptions) { type: 'agent.forked' as const, sourceId: event.sourceId, agentId, - branchId: agentId, + branchId: event.store.ref.branch, ref: entry.ref, }; }), ], }, ], - 'agent.switch': [ - { - guard: ({ context, event }) => context.agents[event.agentId] === undefined, - actions: emit(({ event }) => ({ - type: 'agent.failed' as const, - agentId: event.agentId, - error: `unknown agent: '${event.agentId}'`, - })), - }, - { - guard: ({ context, event }) => - !(context.agents[event.agentId] as AgentEntry).ref.getSnapshot().matches('idle'), - actions: emit(({ event }) => ({ - type: 'agent.failed' as const, - agentId: event.agentId, - error: `agent is busy: '${event.agentId}'`, - })), - }, - { - actions: [ - sendTo( - ({ context, event }) => (context.agents[event.agentId] as AgentEntry).ref, - ({ event }) => ({ - type: 'context.reset' as const, - history: event.input.history, - branchId: event.input.branchId, - }), - ), - emit(({ event }) => ({ - type: 'agent.switched' as const, - agentId: event.agentId, - branchId: event.input.branchId, - reason: event.input.reason, - })), - ], - }, - ], 'agent.send': [ { guard: ({ context, event }) => context.agents[event.agentId] === undefined, diff --git a/packages/agent-core-v2/src/human/session/slices.ts b/packages/agent-core-v2/src/human/session/slices.ts new file mode 100644 index 00000000000..665bd1b2016 --- /dev/null +++ b/packages/agent-core-v2/src/human/session/slices.ts @@ -0,0 +1,39 @@ +import { createSlice } from '#/eventStore/slice'; + +import type { AgentClosed, AgentOpened, AgentSwitched, SessionMetaUpdated } from './events'; + +export interface RosterState { + agents: Record; +} + +export const rosterSlice = createSlice({ + name: 'roster', + initialState: (): RosterState => ({ agents: {} }), + reducers: { + 'agent.opened': (draft, event: AgentOpened) => { + draft.agents[event.agentId] = event.branch; + }, + 'agent.closed': (draft, event: AgentClosed) => { + delete draft.agents[event.agentId]; + }, + 'agent.switched': (draft, event: AgentSwitched) => { + draft.agents[event.agentId] = event.branch; + }, + }, +}); + +export interface SessionMetaState { + value: unknown; +} + +export const sessionMetaSlice = createSlice({ + name: 'sessionMeta', + initialState: (): SessionMetaState => ({ value: undefined }), + reducers: { + 'session.meta_updated': (draft, event: SessionMetaUpdated) => { + draft.value = event.meta; + }, + }, +}); + +export const sessionSlices = { roster: rosterSlice, sessionMeta: sessionMetaSlice }; diff --git a/packages/agent-core-v2/src/human/session/stores.ts b/packages/agent-core-v2/src/human/session/stores.ts new file mode 100644 index 00000000000..22d3e0441bb --- /dev/null +++ b/packages/agent-core-v2/src/human/session/stores.ts @@ -0,0 +1,155 @@ +import { fromEventStore } from '#/eventStore/actor'; +import { createEventStore, type EventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { agentSlices, type AgentEventStore, type AgentStoreActorLogic } from '#/agent/slices'; +import type { StoreBackend } from '#/store/backend/backend'; +import { StoreError, type BranchRef } from '#/store/types'; +import type { Tree } from '#/store/tree'; + +import { agentClosed, agentOpened, agentSwitched, SESSION_LOG_BRANCH } from './events'; +import { sessionSlices } from './slices'; + +export type SessionStore = EventStore; + +export interface OpenedAgentStore { + store: AgentStoreActorLogic; + engine: AgentEventStore; +} + +export type UndoErrorReason = 'unknown-agent' | 'invalid-count' | 'insufficient'; + +export class UndoError extends Error { + readonly reason: UndoErrorReason; + + constructor(reason: UndoErrorReason, message: string) { + super(message); + this.name = 'UndoError'; + this.reason = reason; + } +} + +export function isValidUndoCount(count: number): boolean { + return Number.isSafeInteger(count) && count > 0; +} + +export function freshBranchName(tree: Tree, agentId: string): string { + if (!tree.has(agentId)) return agentId; + let n = 2; + while (tree.has(`${agentId}~${n}`)) n += 1; + return `${agentId}~${n}`; +} + +function undoForkRef(tree: Tree, start: BranchRef): BranchRef | undefined { + if (start.seq > 0) return { branch: start.branch, seq: start.seq - 1 }; + const header = tree.openBranch(start.branch).header; + if (header.parentBranch !== undefined && header.parentSeq !== undefined) { + return { branch: header.parentBranch, seq: header.parentSeq }; + } + return undefined; +} + +export class SessionStores { + private readonly agents = new Map(); + private sessionStore: SessionStore | undefined; + + constructor( + readonly tree: Tree, + readonly backend: StoreBackend, + ) {} + + get(agentId: string): AgentEventStore | undefined { + return this.agents.get(agentId); + } + + async session(): Promise { + if (this.sessionStore === undefined) { + const branch = this.tree.has(SESSION_LOG_BRANCH) + ? this.tree.openBranch(SESSION_LOG_BRANCH) + : this.tree.createBranch(SESSION_LOG_BRANCH); + this.sessionStore = await createEventStore({ + journal: journalFromBranch(branch, this.tree), + slices: sessionSlices, + }); + } + return this.sessionStore; + } + + async open(agentId: string, opts?: { from?: BranchRef }): Promise { + const existing = this.agents.get(agentId); + if (existing !== undefined) { + return { store: fromEventStore(existing), engine: existing }; + } + const existed = this.tree.has(agentId); + const branch = existed + ? this.tree.openBranch(agentId) + : this.tree.createBranch(agentId, opts?.from !== undefined ? { from: opts.from } : undefined); + const engine = await createEventStore({ + journal: journalFromBranch(branch, this.tree), + slices: agentSlices, + }); + this.agents.set(agentId, engine); + if (!existed) { + await (await this.session()).dispatch(agentOpened({ agentId, branch: branch.name })); + } + return { store: fromEventStore(engine), engine }; + } + + async fork(sourceId: string, agentId: string): Promise { + const source = this.agents.get(sourceId); + if (source === undefined) { + throw new StoreError('unknown-agent', `unknown agent '${sourceId}'`); + } + const sourceBranch = this.tree.openBranch(source.ref.branch); + const head = sourceBranch.head; + return this.open( + agentId, + head === null ? undefined : { from: { branch: sourceBranch.name, seq: head } }, + ); + } + + async close(agentId: string): Promise { + const store = this.agents.get(agentId); + if (store === undefined) return; + this.agents.delete(agentId); + await store.close(); + await (await this.session()).dispatch(agentClosed({ agentId })); + } + + async undo(agentId: string, turns: number): Promise<{ branchId: string }> { + const store = this.agents.get(agentId); + if (store === undefined) { + throw new UndoError('unknown-agent', `unknown agent: '${agentId}'`); + } + if (!isValidUndoCount(turns)) { + throw new UndoError('invalid-count', `invalid undo count: ${turns}`); + } + const index = store.slice('turnIndex').turns; + const cut = index.at(-turns); + if (cut === undefined) { + throw new UndoError('insufficient', `cannot undo ${turns} turn(s): not enough turns`); + } + const from = undoForkRef(this.tree, cut.start); + const branchId = freshBranchName(this.tree, agentId); + const branch = + from === undefined + ? this.tree.createBranch(branchId) + : this.tree.createBranch(branchId, { from }); + await store.reset(journalFromBranch(branch, this.tree)); + await ( + await this.session() + ).dispatch(agentSwitched({ agentId, branch: branchId, reason: 'undo' })); + return { branchId }; + } + + async flush(): Promise { + await Promise.all([...this.agents.values()].map((store) => store.flush())); + await this.sessionStore?.flush(); + } + + async dispose(): Promise { + await Promise.all([...this.agents.values()].map((store) => store.close())); + this.agents.clear(); + await this.sessionStore?.close(); + this.sessionStore = undefined; + } +} diff --git a/packages/agent-core-v2/src/human/session/undo.ts b/packages/agent-core-v2/src/human/session/undo.ts deleted file mode 100644 index f0a24aeb0d0..00000000000 --- a/packages/agent-core-v2/src/human/session/undo.ts +++ /dev/null @@ -1,118 +0,0 @@ -import type { ActorRefFrom } from '#/xstate2'; - -import { loadAgentState, type LoadedAgentState } from '#/agent/replay'; -import type { HistoryMessage, UserMeta } from '#/agent/turn'; -import type { Branch } from '#/store/branch'; -import type { Tree } from '#/store/tree'; -import type { BranchRef, BranchHeader } from '#/store/types'; - -import type { createSessionMachine } from '#/session/machine'; - -type SessionActor = ActorRefFrom>; - -export type UndoErrorReason = 'unknown-agent' | 'invalid-count' | 'busy' | 'insufficient'; - -export class UndoError extends Error { - readonly reason: UndoErrorReason; - - constructor(reason: UndoErrorReason, message: string) { - super(message); - this.name = 'UndoError'; - this.reason = reason; - } -} - -export interface UndoResult extends LoadedAgentState { - agentId: string; - branchId: string; -} - -export function isUndoAnchor(entry: HistoryMessage): boolean { - if (entry.message.role !== 'user') return false; - const source = (entry.meta as UserMeta).source; - return source === undefined || source === 'input'; -} - -export function isValidUndoCount(count: number): boolean { - return Number.isSafeInteger(count) && count > 0; -} - -export function freshBranchName(tree: Tree, agentId: string): string { - if (!tree.has(agentId)) return agentId; - let n = 2; - while (tree.has(`${agentId}~${n}`)) n += 1; - return `${agentId}~${n}`; -} - -export async function findUndoCut( - tree: Tree, - start: Branch, - turns: number, -): Promise { - let remaining = turns; - let branch: Branch | undefined = start; - let from: number | null = start.head; - while (branch !== undefined) { - for (let seq = from ?? -1; seq >= 0; seq--) { - const entry = branch.entryAt(seq); - if (entry === null || entry.type !== 'message') continue; - const data = await tree.resolve(entry); - if (!isUndoAnchor(data as HistoryMessage)) continue; - remaining -= 1; - if (remaining > 0) continue; - if (seq > 0) return { branch: branch.name, seq: seq - 1 }; - const { parentBranch, parentSeq } = branch.header; - if (parentBranch !== undefined && parentSeq !== undefined) { - return { branch: parentBranch, seq: parentSeq }; - } - return null; - } - const header: BranchHeader = branch.header; - const next: Branch | undefined = - header.parentBranch !== undefined && - header.parentSeq !== undefined && - tree.has(header.parentBranch) - ? tree.openBranch(header.parentBranch) - : undefined; - from = header.parentSeq ?? null; - branch = next; - } - throw new UndoError('insufficient', `cannot undo ${turns} turn(s): not enough anchors`); -} - -export async function undoAgentTurns( - session: SessionActor, - tree: Tree, - agentId: string, - turns: number, -): Promise { - const entry = session.getSnapshot().context.agents[agentId]; - if (entry === undefined) { - throw new UndoError('unknown-agent', `unknown agent: '${agentId}'`); - } - if (!isValidUndoCount(turns)) { - throw new UndoError('invalid-count', `invalid undo count: ${turns}`); - } - const snapshot = entry.ref.getSnapshot(); - if (!snapshot.matches('idle')) { - throw new UndoError('busy', `agent is busy: '${agentId}'`); - } - if (!tree.has(snapshot.context.branchId)) { - throw new UndoError('insufficient', `cannot undo ${turns} turn(s): not enough anchors`); - } - const branch = tree.openBranch(snapshot.context.branchId); - await branch.settled(); - const cut = await findUndoCut(tree, branch, turns); - const branchId = freshBranchName(tree, agentId); - tree.createBranch(branchId, cut === null ? undefined : { from: cut }); - const loaded = await loadAgentState(tree, branchId); - session.send({ - type: 'agent.switch', - agentId, - input: { branchId, history: loaded.messages, reason: 'undo' }, - }); - if (entry.ref.getSnapshot().context.branchId !== branchId) { - throw new UndoError('busy', `agent is busy: '${agentId}'`); - } - return { agentId, branchId, ...loaded }; -} diff --git a/packages/agent-core-v2/src/human/test/agent/machine.test.ts b/packages/agent-core-v2/src/human/test/agent/machine.test.ts index ed256ae9b31..f8fa34b7b98 100644 --- a/packages/agent-core-v2/src/human/test/agent/machine.test.ts +++ b/packages/agent-core-v2/src/human/test/agent/machine.test.ts @@ -1,12 +1,6 @@ import { describe, expect, it, vi } from 'vitest'; import { createActor, waitFor } from '#/xstate2'; -const originalWarn = console.warn; -console.warn = (...args: unknown[]) => { - originalWarn(...args); - originalWarn(new Error('warn-trace').stack?.split('\n').slice(1, 16).join('\n')); -}; - import { UNKNOWN_CAPABILITY } from '#/llm/capability'; import { createAssistantMessage, @@ -25,6 +19,8 @@ import { connectPlugins } from '#/plugin'; import { createTimingPlugin } from '#/timing/plugin'; import { createAgentMachine, type AgentEmitted } from '#/agent/machine'; import { estimateMessageTokens, estimateTextTokens } from '#/agent/context-usage'; +import { messageAppended, turnEnded } from '#/agent/events'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import { createTurnMachine, createUserEntry, @@ -33,6 +29,11 @@ import { type HistoryMessage, } from '#/agent/turn'; import { MaxStepsExceededError } from '#/agent/errors'; +import { createEventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; +import type { Tree } from '#/store/tree'; import { waitForTool } from '#/tool/wait-for'; import { defineTool, type ToolDefinition } from '#/tool/tool'; import type { ToolResult } from '#/tool/executor'; @@ -97,22 +98,44 @@ function stubTools( ); } +async function testStore(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('test'); + tree.createBranch('main'); + return createEventStore({ journal: journalFromBranch(tree.openBranch('main'), tree), slices: agentSlices }); +} + +async function seedBranch(tree: Tree, branch: string, texts: readonly string[]): Promise { + tree.createBranch(branch); + const seed = await createEventStore({ + journal: journalFromBranch(tree.openBranch(branch), tree), + slices: agentSlices, + }); + for (const text of texts) { + await seed.dispatch( + messageAppended({ message: createUserEntry(createUserMessage(text), { source: 'input' }) }), + ); + } + await seed.dispatch(turnEnded({ turnId: 0, outcome: 'done' })); + await seed.close(); +} + async function runAgent( requester: LlmRequester, tools: readonly ToolDefinition[], retry?: LlmRetryOptions, ): Promise { - const actor = createActor(createTestAgentMachine(tools, requester, retry), { - input: { request: { model } }, + const store = await testStore(); + const actor = createActor(createTestAgentMachine(tools, requester, retry), { + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - const snapshot = await waitFor( - actor, - (s) => s.matches('idle') && s.context.messages.length > 1, - { timeout: 5000 }, - ); - return snapshot.context.messages; + await waitFor(actor, (s) => s.matches('idle') && store.getState().history.length > 1, { + timeout: 5000, + }); + return store.getState().history; } function rolesAndTexts(messages: readonly HistoryMessage[]): string[] { @@ -358,8 +381,9 @@ describe('agent machine async tools', () => { resolveBg = resolve; }); }, 'bg_tool'); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -373,13 +397,13 @@ describe('agent machine async tools', () => { ).not.toBe('done'); resolveBg?.({ content: [{ type: 'text', text: 'bg-result' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 6, + (s) => s.matches('idle') && store.getState().history.length === 6, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:async running: bg_tool', @@ -405,8 +429,9 @@ describe('agent machine async tools', () => { resolvers.set(call.name, resolve); }); }, 'bg_tool', 'sync_tool'); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -421,13 +446,13 @@ describe('agent machine async tools', () => { }); resolvers.get('bg_tool')?.({ content: [{ type: 'text', text: 'bg-result' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 7, + (s) => s.matches('idle') && store.getState().history.length === 7, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:async running: bg_tool', @@ -457,8 +482,9 @@ describe('agent machine async tools', () => { }, }); const tools = [bgTool, waitForTool]; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -467,13 +493,13 @@ describe('agent machine async tools', () => { expect(actor.getSnapshot().context.turnTools['call-2']).toBeDefined(); }); resolveBg?.({ content: [{ type: 'text', text: 'bg-result' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 7, + (s) => s.matches('idle') && store.getState().history.length === 7, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:async running: bg_tool', @@ -506,8 +532,9 @@ describe('agent machine async tools', () => { }, }); const tools = [bgTool, waitForTool]; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -519,13 +546,13 @@ describe('agent machine async tools', () => { { timeout: 4000 }, ); resolveBg?.({ content: [{ type: 'text', text: 'bg-result' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 8, + (s) => s.matches('idle') && store.getState().history.length === 8, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:async running: bg_tool', @@ -563,26 +590,27 @@ describe('agent machine lifecycle', () => { createAssistantMessage([{ type: 'text', text: 'second' }]), ]); const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const completedTurns: number[] = []; actor.on('turn.done', (event) => completedTurns.push(event.messages.length)); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length === 2, { + await waitFor(actor, (s) => s.matches('idle') && store.getState().history.length === 2, { timeout: 5000, }); actor.send({ type: 'input.submit', message: createUserMessage('again') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 4, + (s) => s.matches('idle') && store.getState().history.length === 4, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:first', 'user:again', @@ -606,8 +634,9 @@ describe('agent machine lifecycle', () => { }), 'slow_tool', ); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -616,17 +645,19 @@ describe('agent machine lifecycle', () => { expect(resolveTool).toBeDefined(); }); actor.send({ type: 'input.submit', message: createUserMessage('mid-turn') }); - expect(actor.getSnapshot().context.queue).toHaveLength(1); - expect(actor.getSnapshot().context.notifications).toHaveLength(0); + await vi.waitFor(() => { + expect(store.getState().queue).toHaveLength(1); + expect(store.getState().notifications).toHaveLength(0); + }); resolveTool?.({ content: [{ type: 'text', text: 'slow' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 6, + (s) => s.matches('idle') && store.getState().history.length === 6, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:slow', @@ -650,30 +681,33 @@ describe('agent machine lifecycle', () => { }, }; const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const failures: unknown[] = []; actor.on('turn.failed', (event) => failures.push(event.error)); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - const failedSnapshot = await waitFor( + await waitFor( actor, (s) => s.matches('idle') && failures.length === 1, { timeout: 5000 }, ); - expect(rolesAndTexts(failedSnapshot.context.messages)).toEqual(['user:hi']); + await vi.waitFor(() => { + expect(rolesAndTexts(store.getState().history)).toEqual(['user:hi']); + }); expect(failures[0]).toMatchObject({ kind: 'unknown', message: 'llm down' }); expect(actor.getSnapshot().status).toBe('active'); actor.send({ type: 'input.submit', message: createUserMessage('retry') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 3, + (s) => s.matches('idle') && store.getState().history.length === 3, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'user:retry', 'assistant:recovered', @@ -695,8 +729,9 @@ describe('agent machine input.notify', () => { }), 'slow_tool', ); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -708,16 +743,18 @@ describe('agent machine input.notify', () => { type: 'input.notify', message: createUserMessage('stale'), }); - expect(actor.getSnapshot().context.notifications).toHaveLength(1); + await vi.waitFor(() => { + expect(store.getState().notifications).toHaveLength(1); + }); resolveTool?.({ content: [{ type: 'text', text: 'slow' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 5, + (s) => s.matches('idle') && store.getState().history.length === 5, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:slow', @@ -744,8 +781,9 @@ describe('agent machine input.notify', () => { }, }; const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const completedTurns: number[] = []; actor.on('turn.done', (event) => completedTurns.push(event.messages.length)); @@ -759,17 +797,19 @@ describe('agent machine input.notify', () => { type: 'input.notify', message: createUserMessage('stale'), }); - expect(actor.getSnapshot().context.notifications).toHaveLength(1); + await vi.waitFor(() => { + expect(store.getState().notifications).toHaveLength(1); + }); streamMessage(createAssistantMessage([{ type: 'text', text: 'first' }]), firstOnEvent); resolveFirst?.(); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 4, + (s) => s.matches('idle') && store.getState().history.length === 4, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:first', 'user:stale', @@ -783,8 +823,9 @@ describe('agent machine input.notify', () => { createAssistantMessage([{ type: 'text', text: 'done' }]), ]); const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const llmDone: AssistantEntry[] = []; actor.on('llm.done', (event) => { @@ -793,13 +834,13 @@ describe('agent machine input.notify', () => { actor.start(); actor.send({ type: 'input.notify', message: createUserMessage('queued') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 2, + (s) => s.matches('idle') && store.getState().history.length === 2, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual(['user:queued', 'assistant:done']); + expect(rolesAndTexts(store.getState().history)).toEqual(['user:queued', 'assistant:done']); expect(llmDone.map((entry) => extractText(entry.message))).toEqual(['done']); expect(llmDone[0]?.meta).toEqual({ model: { provider: 'test', model: 'test-model' }, @@ -826,8 +867,9 @@ describe('agent machine input.remind', () => { }), 'slow_tool', ); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const consumedKeys: (string | undefined)[][] = []; actor.on('turn.reminders_consumed', (event) => { @@ -847,36 +889,37 @@ describe('agent machine input.remind', () => { key: 'todo', message: createUserMessage('\nstale\n'), }); - await new Promise((resolve) => setTimeout(resolve, 50)); - const idleSnapshot = actor.getSnapshot(); - expect(idleSnapshot.matches('idle')).toBe(true); - expect(idleSnapshot.context.messages).toHaveLength(0); - expect(idleSnapshot.context.reminders).toHaveLength(1); - expect(idleSnapshot.context.reminders[0]?.meta).toEqual({ source: 'reminder', key: 'todo' }); - expect(extractText(idleSnapshot.context.reminders[0]?.message ?? createUserMessage(''))).toContain('stale'); + await vi.waitFor(() => { + expect(store.getState().reminders).toHaveLength(1); + }); + expect(actor.getSnapshot().matches('idle')).toBe(true); + expect(store.getState().history).toHaveLength(0); + expect(store.getState().reminders).toHaveLength(1); + expect(store.getState().reminders[0]?.meta).toEqual({ source: 'reminder', key: 'todo' }); + expect(extractText(store.getState().reminders[0]?.message ?? createUserMessage(''))).toContain('stale'); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); await vi.waitFor(() => { expect(resolveTool).toBeDefined(); }); - expect(actor.getSnapshot().context.reminders).toHaveLength(1); + expect(store.getState().reminders).toHaveLength(1); resolveTool?.({ content: [{ type: 'text', text: 'slow' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 5, + (s) => s.matches('idle') && store.getState().history.length === 5, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:slow', 'user:\nstale\n', 'assistant:done', ]); - expect(snapshot.context.messages[3]?.meta).toEqual({ source: 'reminder', key: 'todo' }); - expect(snapshot.context.reminders).toHaveLength(0); + expect(store.getState().history[3]?.meta).toEqual({ source: 'reminder', key: 'todo' }); + expect(store.getState().reminders).toHaveLength(0); expect(consumedKeys).toEqual([['todo']]); }); }); @@ -909,10 +952,11 @@ describe('agent machine llm retry', () => { const ticks = [1000, 1100, 1200, 100000, 100100, 100140, 100200]; const timingPlugin = createTimingPlugin({ now: () => ticks.shift() ?? Number.NaN }); const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor( - createTestAgentMachine(tools, requester, { maxAttemptsPerStep: 3 }), + createTestAgentMachine(tools, requester, { maxAttemptsPerStep: 3 }), { - input: { request: { model } }, + input: { request: { model }, store }, }, ); connectPlugins(actor, [timingPlugin]); @@ -923,9 +967,9 @@ describe('agent machine llm retry', () => { actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 2, + (s) => s.matches('idle') && store.getState().history.length === 2, { timeout: 5000 }, ); @@ -938,7 +982,7 @@ describe('agent machine llm retry', () => { statusCode: 500, }); expect(failures).toHaveLength(0); - expect(rolesAndTexts(snapshot.context.messages)).toEqual(['user:hi', 'assistant:recovered']); + expect(rolesAndTexts(store.getState().history)).toEqual(['user:hi', 'assistant:recovered']); const delayMs = retrying[0]?.delayMs ?? 0; expect(timingPlugin.timing()).toEqual({ @@ -967,8 +1011,9 @@ describe('agent machine input.steer', () => { }), 'slow_tool', ); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -977,24 +1022,30 @@ describe('agent machine input.steer', () => { expect(resolveTool).toBeDefined(); }); actor.send({ type: 'input.submit', id: 'p1', message: createUserMessage('steer me') }); - expect(actor.getSnapshot().context.queue).toHaveLength(1); + await vi.waitFor(() => { + expect(store.getState().queue).toHaveLength(1); + }); actor.send({ type: 'input.steer', id: 'nope' }); - expect(actor.getSnapshot().context.queue).toHaveLength(1); - expect(actor.getSnapshot().context.notifications).toHaveLength(0); + await vi.waitFor(() => { + expect(store.getState().queue).toHaveLength(1); + expect(store.getState().notifications).toHaveLength(0); + }); actor.send({ type: 'input.steer', id: 'p1' }); - expect(actor.getSnapshot().context.queue).toHaveLength(0); - expect(actor.getSnapshot().context.notifications).toHaveLength(1); + await vi.waitFor(() => { + expect(store.getState().queue).toHaveLength(0); + expect(store.getState().notifications).toHaveLength(1); + }); resolveTool?.({ content: [{ type: 'text', text: 'slow' }] }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 5, + (s) => s.matches('idle') && store.getState().history.length === 5, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:slow', @@ -1028,8 +1079,9 @@ describe('agent machine input.abort', () => { }, }; const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const aborted: HistoryMessage[][] = []; const aborting: unknown[] = []; @@ -1042,7 +1094,7 @@ describe('agent machine input.abort', () => { expect(signals).toHaveLength(1); }); actor.send({ type: 'input.abort' }); - const snapshot = await waitFor( + await waitFor( actor, (s) => s.matches('idle') && aborted.length === 1, { timeout: 5000 }, @@ -1050,7 +1102,9 @@ describe('agent machine input.abort', () => { expect(signals[0]?.aborted).toBe(true); expect(aborting).toHaveLength(1); - expect(rolesAndTexts(snapshot.context.messages)).toEqual(['user:hi', 'assistant:hello']); + await vi.waitFor(() => { + expect(rolesAndTexts(store.getState().history)).toEqual(['user:hi', 'assistant:hello']); + }); expect(rolesAndTexts(aborted[0] ?? [])).toEqual(['user:hi', 'assistant:hello']); const salvaged = aborted[0]?.[1]; expect(salvaged?.message.role === 'assistant' && salvaged.message.toolCalls).toEqual([]); @@ -1068,8 +1122,9 @@ describe('agent machine input.abort', () => { signal.addEventListener('abort', () => reject(new Error('tool stopped'))); }); }, 'slow_tool'); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -1079,14 +1134,14 @@ describe('agent machine input.abort', () => { }); actor.send({ type: 'input.abort' }); expect(actor.getSnapshot().value).toEqual({ running: 'aborting' }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 3, + (s) => s.matches('idle') && store.getState().history.length === 3, { timeout: 5000 }, ); expect(signals[0]?.aborted).toBe(true); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:aborted', @@ -1104,8 +1159,9 @@ describe('agent machine input.abort', () => { }), 'slow_tool', ); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); const aborted: HistoryMessage[][] = []; actor.on('turn.aborted', (event) => aborted.push(event.messages)); @@ -1116,17 +1172,19 @@ describe('agent machine input.abort', () => { expect(actor.getSnapshot().context.turnTools['call-1']).toBeDefined(); }); actor.send({ type: 'input.abort' }); - const snapshot = await waitFor( + await waitFor( actor, (s) => s.matches('idle') && aborted.length === 1, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ - 'user:hi', - 'assistant:', - 'tool:partial', - ]); + await vi.waitFor(() => { + expect(rolesAndTexts(store.getState().history)).toEqual([ + 'user:hi', + 'assistant:', + 'tool:partial', + ]); + }); }); it('forces the turn to aborted on a second abort when a tool ignores the signal', async () => { @@ -1138,8 +1196,9 @@ describe('agent machine input.abort', () => { signals.push(signal); return new Promise(() => {}); }, 'slow_tool'); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -1151,14 +1210,14 @@ describe('agent machine input.abort', () => { expect(actor.getSnapshot().value).toEqual({ running: 'aborting' }); actor.send({ type: 'input.abort' }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 3, + (s) => s.matches('idle') && store.getState().history.length === 3, { timeout: 5000 }, ); expect(signals[0]?.aborted).toBe(true); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:aborted', @@ -1170,8 +1229,9 @@ describe('agent machine input.abort', () => { createAssistantMessage([], [toolCall('call-1', 'slow_tool')]), ]); const tools = stubTools(() => new Promise(() => {}), 'slow_tool'); - const actor = createActor(createTestAgentMachine(tools, requester, undefined, 50), { - input: { request: { model } }, + const store = await testStore(); + const actor = createActor(createTestAgentMachine(tools, requester, undefined, 50), { + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -1181,13 +1241,13 @@ describe('agent machine input.abort', () => { }); actor.send({ type: 'input.abort' }); expect(actor.getSnapshot().value).toEqual({ running: 'aborting' }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 3, + (s) => s.matches('idle') && store.getState().history.length === 3, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'assistant:', 'tool:aborted', @@ -1215,8 +1275,9 @@ describe('agent machine input.abort', () => { }, }; const tools: ToolDefinition[] = []; + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -1228,13 +1289,13 @@ describe('agent machine input.abort', () => { actor.send({ type: 'input.submit', id: 'p2', message: createUserMessage('steered') }); actor.send({ type: 'input.steer', id: 'p2' }); actor.send({ type: 'input.abort' }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 4, + (s) => s.matches('idle') && store.getState().history.length === 4, { timeout: 5000 }, ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:hi', 'user:steered', 'user:queued', @@ -1272,8 +1333,9 @@ describe('agent machine input.abort', () => { detach?.({ text: 'async running: bg_tool' }); return new Promise(() => {}); }, 'bg_tool'); + const store = await testStore(); const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); @@ -1311,13 +1373,14 @@ describe('agent machine max steps', () => { }), 'ok_tool', ); + const store = await testStore(); const actor = createActor( createAgentMachine({ tools, turnActor: createTurnMachine(createLlmMachine({ requester })), maxStepsPerTurn: 2, }), - { input: { request: { model } } }, + { input: { request: { model }, store } }, ); const failures: Extract[] = []; actor.on('turn.failed', (event) => failures.push(event)); @@ -1349,7 +1412,7 @@ describe('agent machine max steps', () => { }); resolvers[3]?.(); - const snapshot = await waitFor(actor, (s) => s.matches('idle') && failures.length === 1, { + await waitFor(actor, (s) => s.matches('idle') && failures.length === 1, { timeout: 5000, }); @@ -1358,18 +1421,20 @@ describe('agent machine max steps', () => { expect(failures[0]?.error).toBeInstanceOf(MaxStepsExceededError); expect((failures[0]?.error as MaxStepsExceededError).code).toBe('loop.max_steps_exceeded'); expect((failures[0]?.error as MaxStepsExceededError).details).toEqual({ maxSteps: 2 }); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ - 'user:hi', - 'assistant:', - 'tool:ok', - 'assistant:', - 'tool:ok', - 'user:keep going', - 'assistant:', - 'tool:ok', - 'assistant:', - 'tool:ok', - ]); + await vi.waitFor(() => { + expect(rolesAndTexts(store.getState().history)).toEqual([ + 'user:hi', + 'assistant:', + 'tool:ok', + 'assistant:', + 'tool:ok', + 'user:keep going', + 'assistant:', + 'tool:ok', + 'assistant:', + 'tool:ok', + ]); + }); expect(actor.getSnapshot().status).toBe('active'); }); }); @@ -1379,8 +1444,16 @@ describe('agent machine context reset', () => { const requester = createStubRequester([ createAssistantMessage([{ type: 'text', text: 'reply' }]), ]); + const backend = new MemoryBackend(); + const treeStore = await TreeStore.open(backend, {}); + const tree = await treeStore.tree('test'); + tree.createBranch('main'); + const store = await createEventStore({ + journal: journalFromBranch(tree.openBranch('main'), tree), + slices: agentSlices, + }); const actor = createActor(createTestAgentMachine([], requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); const resets: string[] = []; @@ -1395,40 +1468,41 @@ describe('agent machine context reset', () => { }); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length === 2, { + await waitFor(actor, (s) => s.matches('idle') && store.getState().history.length === 2, { timeout: 5000, }); - actor.send({ - type: 'context.reset', - history: [createUserEntry(createUserMessage('seed'), { source: 'input' })], - branchId: 'main~2', - }); + await seedBranch(tree, 'main~2', ['seed']); + await store.reset(journalFromBranch(tree.openBranch('main~2'), tree)); - expect(actor.getSnapshot().context.messages).toHaveLength(1); - expect(actor.getSnapshot().context.turnId).toBe(1); - expect(actor.getSnapshot().context.branchId).toBe('main~2'); - expect(resets).toEqual(['main~2']); + await vi.waitFor(() => { + expect(resets).toEqual(['main~2']); + }); + expect(store.ref.branch).toBe('main~2'); + expect(store.getState().history).toHaveLength(1); + expect(store.getState().turnIndex.nextTurnId).toBe(1); actor.send({ type: 'input.submit', message: createUserMessage('again') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 3, + (s) => s.matches('idle') && store.getState().history.length === 3, { timeout: 5000 }, ); - expect(snapshot.context.turnId).toBe(2); + expect(store.getState().turnIndex.nextTurnId).toBe(2); expect(turnStarts).toEqual([ { turnId: 0, branchId: 'main' }, { turnId: 1, branchId: 'main~2' }, ]); }); - it('ignores context.reset while running, including a turn driven by notify', async () => { + it('aborts an in-flight turn and applies the new state on store.reset while running', async () => { let calls = 0; + const signals: AbortSignal[] = []; const releases: Array<() => void> = []; const requester: LlmRequester = { - generate: (_config, _content, { onEvent }) => { + generate: (_config, _content, { signal, onEvent }) => { calls += 1; + signals.push(signal); return new Promise((resolve) => { releases.push(() => { onEvent?.({ type: 'llm.streaming.part', part: { type: 'text', text: 'late' } }); @@ -1438,37 +1512,46 @@ describe('agent machine context reset', () => { }); }, }; + const backend = new MemoryBackend(); + const treeStore = await TreeStore.open(backend, {}); + const tree = await treeStore.tree('test'); + tree.createBranch('main'); + const store = await createEventStore({ + journal: journalFromBranch(tree.openBranch('main'), tree), + slices: agentSlices, + }); const actor = createActor(createTestAgentMachine([], requester), { - input: { request: { model } }, + input: { request: { model }, store }, }); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('hi') }); await vi.waitFor(() => expect(calls).toBe(1)); - actor.send({ type: 'context.reset', history: [], branchId: 'other' }); - expect(actor.getSnapshot().context.branchId).toBe('main'); + await seedBranch(tree, 'other', ['seeded']); + await store.reset(journalFromBranch(tree.openBranch('other'), tree)); - releases[0]?.(); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length === 2, { + await waitFor(actor, (s) => s.matches('idle') && store.ref.branch === 'other', { timeout: 5000, }); + expect(signals[0]?.aborted).toBe(true); + expect(rolesAndTexts(store.getState().history)).toEqual(['user:seeded']); + + releases[0]?.(); + await new Promise((resolve) => setTimeout(resolve, 300)); + expect(rolesAndTexts(store.getState().history)).toEqual(['user:seeded']); actor.send({ type: 'input.notify', message: createUserMessage('note') }); await vi.waitFor(() => expect(calls).toBe(2)); - actor.send({ type: 'context.reset', history: [], branchId: 'other' }); - expect(actor.getSnapshot().context.branchId).toBe('main'); + + await seedBranch(tree, 'third', ['third-seed']); + await store.reset(journalFromBranch(tree.openBranch('third'), tree)); + + await waitFor(actor, (s) => s.matches('idle') && store.ref.branch === 'third', { + timeout: 5000, + }); + expect(signals[1]?.aborted).toBe(true); + expect(rolesAndTexts(store.getState().history)).toEqual(['user:third-seed']); releases[1]?.(); - const snapshot = await waitFor( - actor, - (s) => s.matches('idle') && s.context.messages.length === 4, - { timeout: 5000 }, - ); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ - 'user:hi', - 'assistant:late', - 'user:note', - 'assistant:late', - ]); }); }); diff --git a/packages/agent-core-v2/src/human/test/agent/persist.test.ts b/packages/agent-core-v2/src/human/test/agent/persist.test.ts deleted file mode 100644 index 06368b0cf0c..00000000000 --- a/packages/agent-core-v2/src/human/test/agent/persist.test.ts +++ /dev/null @@ -1,264 +0,0 @@ -import { describe, expect, it, vi } from 'vitest'; -import { createActor, waitFor } from 'xstate'; - -import { collectPluginTools, connectPlugins } from '#/plugin'; -import { UNKNOWN_CAPABILITY } from '#/llm/capability'; -import { - createAssistantMessage, - createUserMessage, - type AssistantMessage, - type Message, - type ToolCall, -} from '#/llm/message'; -import type { LlmModel } from '#/llm/model'; -import { createLlmMachine } from '#/llm/requester/machine'; -import type { LlmRequester } from '#/llm/requester/requester'; -import { createAgentMachine } from '#/agent/machine'; -import { loadAgentState, type TurnEntryData } from '#/agent/replay'; -import { persistAgent } from '#/persist/agent'; -import { createTurnMachine } from '#/agent/turn'; -import { MemoryBackend } from '#/store/backend/memory'; -import type { Branch } from '#/store/branch'; -import { TreeStore } from '#/store/store'; -import type { EntryLine } from '#/store/types'; -import { createTodoPlugin } from '#/todo/plugin'; -import { restoreTodoState, snapshotTodoState } from '#/todo/state'; -import { defineTool, type ToolDefinition } from '#/tool/tool'; - -const model: LlmModel = { provider: 'test', model: 'test-model', capability: UNKNOWN_CAPABILITY }; - -function toolCall(id: string, name: string, args: string = '{}'): ToolCall { - return { type: 'function', id, name, arguments: args }; -} - -function createStubRequester(responses: readonly AssistantMessage[]): LlmRequester { - let call = 0; - return { - generate: (_config, _content, { onEvent }) => { - const message = responses[Math.min(call, responses.length - 1)] as AssistantMessage; - call += 1; - for (const part of [...message.content, ...message.toolCalls]) { - onEvent?.({ type: 'llm.streaming.part', part }); - } - onEvent?.({ type: 'llm.done' }); - return Promise.resolve(); - }, - }; -} - -function createTestAgentMachine(tools: readonly ToolDefinition[], requester: LlmRequester) { - return createAgentMachine({ - tools, - turnActor: createTurnMachine(createLlmMachine({ requester })), - }); -} - -const okTool = defineTool({ - name: 'ok_tool', - description: 'stub ok tool', - parameters: { type: 'object', properties: {} }, - execute: () => Promise.resolve({ content: [{ type: 'text', text: 'ok' }] }), -}); - -async function branchOf(store: TreeStore, tree: string, branch: string): Promise { - const loaded = await store.tree(tree); - return loaded.has(branch) ? loaded.openBranch(branch) : loaded.createBranch(branch); -} - -async function entriesOf(store: TreeStore, tree: string, branch: string): Promise { - const loaded = await store.tree(tree); - return [...loaded.openBranch(branch).walk()].toReversed(); -} - -function entryTypes(entries: readonly EntryLine[]): string[] { - return entries.map((entry) => entry.type); -} - -function turnData(entry: EntryLine): TurnEntryData { - return entry.payload.data as TurnEntryData; -} - -async function waitForEntries( - store: TreeStore, - tree: string, - branch: string, - count: number, -): Promise { - let entries: EntryLine[] = []; - await vi.waitFor(async () => { - entries = await entriesOf(store, tree, branch); - expect(entries).toHaveLength(count); - }); - return entries; -} - -describe('persistAgent', () => { - it('persists messages and turn boundaries in order', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const requester = createStubRequester([ - createAssistantMessage([], [toolCall('call-1', 'ok_tool')]), - createAssistantMessage([{ type: 'text', text: 'done' }]), - ]); - const actor = createActor(createTestAgentMachine([okTool], requester), { - input: { request: { model } }, - }); - actor.start(); - const handle = persistAgent(actor, await branchOf(store, 'session', 'main')); - actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - - const entries = await waitForEntries(store, 'session', 'main', 6); - expect(entryTypes(entries)).toEqual([ - 'message', - 'turn', - 'message', - 'message', - 'message', - 'turn', - ]); - expect((entries[0] as EntryLine).payload.data).toMatchObject({ - message: { role: 'user' }, - meta: { source: 'input' }, - }); - expect(turnData(entries[1] as EntryLine)).toEqual({ phase: 'start', turnId: 0 }); - expect(turnData(entries[5] as EntryLine)).toEqual({ - phase: 'end', - turnId: 0, - outcome: 'done', - }); - handle.dispose(); - }); - - it('restores messages and turnId without duplicating entries', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const first = createStubRequester([createAssistantMessage([{ type: 'text', text: 'first' }])]); - const firstActor = createActor(createTestAgentMachine([okTool], first), { - input: { request: { model } }, - }); - firstActor.start(); - persistAgent(firstActor, await branchOf(store, 'session', 'main')); - firstActor.send({ type: 'input.submit', message: createUserMessage('hi') }); - const firstSnapshot = await waitFor( - firstActor, - (s) => s.matches('idle') && s.context.messages.length === 2, - { timeout: 5000 }, - ); - await waitForEntries(store, 'session', 'main', 4); - - const reopened = await TreeStore.open(fs); - const loaded = await loadAgentState(await reopened.tree('session'), 'main'); - expect(loaded.messages).toEqual(firstSnapshot.context.messages); - expect(loaded.turnId).toBe(1); - - const second = createStubRequester([createAssistantMessage([{ type: 'text', text: 'second' }])]); - const secondActor = createActor(createTestAgentMachine([okTool], second), { - input: { request: { model }, history: loaded.messages, turnId: loaded.turnId }, - }); - secondActor.start(); - persistAgent(secondActor, await branchOf(reopened, 'session', 'main')); - secondActor.send({ type: 'input.submit', message: createUserMessage('again') }); - const secondSnapshot = await waitFor( - secondActor, - (s) => s.matches('idle') && s.context.messages.length === 4, - { timeout: 5000 }, - ); - - const entries = await waitForEntries(reopened, 'session', 'main', 8); - const messageEntries = entries.filter((entry) => entry.type === 'message'); - expect(messageEntries.map((entry) => entry.payload.data)).toEqual( - secondSnapshot.context.messages, - ); - const starts = entries - .filter((entry) => entry.type === 'turn') - .map((entry) => turnData(entry)) - .filter((data) => data.phase === 'start'); - expect(starts).toEqual([ - { phase: 'start', turnId: 0 }, - { phase: 'start', turnId: 1 }, - ]); - }); - - it('persists and restores todo state', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const todo = createTodoPlugin(); - const tools = collectPluginTools([todo]); - const requester = createStubRequester([ - createAssistantMessage( - [], - [ - toolCall( - 'call-todo', - 'TodoList', - JSON.stringify({ todos: [{ title: 'task a', status: 'in_progress' }] }), - ), - ], - ), - createAssistantMessage([{ type: 'text', text: 'done' }]), - ]); - const actor = createActor(createTestAgentMachine(tools, requester), { - input: { request: { model } }, - }); - connectPlugins(actor, [todo]); - actor.start(); - persistAgent(actor, await branchOf(store, 'session', 'main'), { - states: { todo: () => snapshotTodoState(todo.state) }, - }); - actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - - const entries = await waitForEntries(store, 'session', 'main', 7); - expect(entryTypes(entries)).toEqual([ - 'message', - 'turn', - 'state', - 'message', - 'message', - 'message', - 'turn', - ]); - const stateEntry = entries.find((entry) => entry.type === 'state') as EntryLine; - expect(stateEntry.payload.data).toEqual({ - name: 'todo', - value: { todos: [{ title: 'task a', status: 'in_progress' }], lastWriteTurn: 1 }, - }); - - const reopened = await TreeStore.open(fs); - const loaded = await loadAgentState(await reopened.tree('session'), 'main'); - const restored = restoreTodoState(loaded.states['todo'], loaded.turnId); - expect(restored.todos).toEqual([{ title: 'task a', status: 'in_progress' }]); - expect(restored.lastWriteTurn).toBe(1); - expect(restored.currentTurn).toBe(1); - }); - - it('restores offloaded messages through blob refs', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs, { offloadThreshold: 1 }); - const requester = createStubRequester([ - createAssistantMessage([{ type: 'text', text: 'a fairly long reply' }]), - ]); - const actor = createActor(createTestAgentMachine([okTool], requester), { - input: { request: { model } }, - }); - actor.start(); - persistAgent(actor, await branchOf(store, 'session', 'main')); - actor.send({ type: 'input.submit', message: createUserMessage('hi') }); - const snapshot = await waitFor( - actor, - (s) => s.matches('idle') && s.context.messages.length === 2, - { timeout: 5000 }, - ); - await waitForEntries(store, 'session', 'main', 4); - - const reopened = await TreeStore.open(fs); - const loaded = await loadAgentState(await reopened.tree('session'), 'main'); - expect(loaded.messages).toEqual(snapshot.context.messages); - expect(loaded.turnId).toBe(1); - }); -}); diff --git a/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts b/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts new file mode 100644 index 00000000000..bfd3cb24b14 --- /dev/null +++ b/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts @@ -0,0 +1,294 @@ +import { describe, expect, it, vi } from 'vitest'; +import { z } from 'zod'; + +import { storeActor } from '#/eventStore/actor'; +import { createEventStore, type Cause } from '#/eventStore/eventStore'; +import { defineEvent } from '#/eventStore/events'; +import { journalFromBranch } from '#/eventStore/journal'; +import { createSlice } from '#/eventStore/slice'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; +import type { Tree } from '#/store/tree'; +import { createActor, waitFor } from '#/xstate2'; + +const counterAdded = defineEvent({ type: 'test.counter_added', schema: z.object({ amount: z.number() }) }); +type CounterAdded = ReturnType; + +const noteTagged = defineEvent({ type: 'test.note_tagged', schema: z.object({ tag: z.string() }) }); +type NoteTagged = ReturnType; + +const counterSlice = createSlice({ + name: 'counter', + initialState: () => 0, + reducers: { + 'test.counter_added': (draft, event: CounterAdded) => draft + event.amount, + 'test.counter_bumped_internal': (draft) => draft + 100, + }, +}); + +const notesSlice = createSlice({ + name: 'notes', + initialState: () => [] as string[], + reducers: { + 'test.note_tagged': (draft, event: NoteTagged) => { + draft.push(event.tag); + }, + }, +}); + +const slices = { counter: counterSlice, notes: notesSlice }; + +async function openTree(backend: MemoryBackend = new MemoryBackend()): Promise { + const store = await TreeStore.open(backend, {}); + return store.tree('test'); +} + +async function openJournal(tree: Tree, branch = 'main') { + if (!tree.has(branch)) tree.createBranch(branch); + return journalFromBranch(tree.openBranch(branch), tree); +} + +async function openStore(tree: Tree, opts?: { snapshot?: { everyEvents?: number } | false; drainLimit?: number; extraSlices?: Record }) { + const journal = await openJournal(tree); + return createEventStore({ journal, slices, snapshot: opts?.snapshot, drainLimit: opts?.drainLimit }); +} + +describe('createEventStore', () => { + it('folds dispatched events and refolds them on reopen', async () => { + const tree = await openTree(); + const store = await openStore(tree); + await store.dispatch(counterAdded({ amount: 3 })); + await store.dispatch(noteTagged({ tag: 'a' })); + await store.flush(); + + expect(store.getState()).toEqual({ counter: 3, notes: ['a'] }); + + const reopened = await openStore(tree); + expect(reopened.getState()).toEqual({ counter: 3, notes: ['a'] }); + }); + + it('folds from the latest snapshot plus its tail', async () => { + const tree = await openTree(); + const store = await openStore(tree, { snapshot: { everyEvents: 2 } }); + await store.dispatch(counterAdded({ amount: 1 })); + await store.dispatch(counterAdded({ amount: 2 })); + await store.dispatch(counterAdded({ amount: 3 })); + await store.flush(); + + const reopened = await openStore(tree); + expect(reopened.getState()).toEqual({ counter: 6, notes: [] }); + }); + + it('seeds slices missing from a snapshot from their initial state', async () => { + const tree = await openTree(); + const journal = await openJournal(tree); + await journal.append({ type: 'snapshot', kind: 'snapshot', data: { slices: { counter: 41 } } }); + await journal.append({ + type: 'test.counter_added', + kind: 'event', + data: { type: 'test.counter_added', time: 1, amount: 1 }, + }); + const store = await createEventStore({ journal, slices }); + expect(store.getState()).toEqual({ counter: 42, notes: [] }); + }); + + it('skips unknown event types when folding', async () => { + const tree = await openTree(); + const journal = await openJournal(tree); + await journal.append({ type: 'test.unknown_event', kind: 'event', data: { type: 'test.unknown_event' } }); + await journal.append({ + type: 'test.counter_added', + kind: 'event', + data: { type: 'test.counter_added', time: 1, amount: 5 }, + }); + const store = await createEventStore({ journal, slices }); + expect(store.slice('counter')).toBe(5); + }); +}); + +describe('dispatch', () => { + it('rejects unregistered events and schema-invalid events', async () => { + const tree = await openTree(); + const store = await openStore(tree); + await expect( + store.dispatch({ type: 'test.ghost_event', time: 1, amount: 1 }), + ).rejects.toMatchObject({ + code: 'unregistered-event', + }); + await expect(store.dispatch(counterAdded({ amount: 'x' as unknown as number }))).rejects.toMatchObject({ + code: 'schema', + }); + expect(store.slice('counter')).toBe(0); + }); + + it('serializes concurrent dispatches in seq order', async () => { + const tree = await openTree(); + const store = await openStore(tree); + const [a, b] = await Promise.all([ + store.dispatch(counterAdded({ amount: 1 })), + store.dispatch(counterAdded({ amount: 2 })), + ]); + expect(a.seq).toBe(0); + expect(b.seq).toBe(1); + expect(store.slice('counter')).toBe(3); + }); + + it('reads back appended entries through the journal', async () => { + const tree = await openTree(); + const store = await openStore(tree); + await store.dispatch(noteTagged({ tag: 'x' })); + await store.flush(); + const records = []; + for await (const record of (await openJournal(tree)).read()) records.push(record); + expect(records).toHaveLength(1); + expect(records[0]).toMatchObject({ branch: 'main', seq: 0, type: 'test.note_tagged', kind: 'event' }); + }); +}); + +describe('internal events', () => { + it('folds raised internal events across slices without persisting them', async () => { + const tree = await openTree(); + const raiserSlice = createSlice({ + name: 'raiser', + initialState: () => 0, + reducers: { + 'test.counter_added': (draft, event: CounterAdded, ctx) => { + ctx.enqueue.raise({ type: 'test.counter_bumped_internal' }); + return draft + event.amount; + }, + }, + }); + const journal = await openJournal(tree); + const store = await createEventStore({ journal, slices: { counter: counterSlice, raiser: raiserSlice } }); + await store.dispatch(counterAdded({ amount: 5 })); + expect(store.getState()).toEqual({ counter: 105, raiser: 5 }); + + const records = []; + for await (const record of journal.read()) records.push(record); + expect(records).toHaveLength(1); + + const reopened = await createEventStore({ journal, slices: { counter: counterSlice, raiser: raiserSlice } }); + expect(reopened.getState()).toEqual({ counter: 105, raiser: 5 }); + }); + + it('enforces the drain limit', async () => { + const tree = await openTree(); + const loopSlice = createSlice({ + name: 'loop', + initialState: () => 0, + reducers: { + 'test.counter_added': (draft, _event, ctx) => { + ctx.enqueue.raise({ type: 'test.counter_bumped_internal' }); + return draft + 1; + }, + 'test.counter_bumped_internal': (draft, _event, ctx) => { + ctx.enqueue.raise({ type: 'test.counter_bumped_internal' }); + return draft + 1; + }, + }, + }); + const journal = await openJournal(tree); + const store = await createEventStore({ journal, slices: { loop: loopSlice }, drainLimit: 10 }); + await expect(store.dispatch(counterAdded({ amount: 1 }))).rejects.toMatchObject({ + code: 'drain-limit', + }); + }); +}); + +describe('registerSlice', () => { + it('folds history for late-joined slices and notifies slice-joined', async () => { + const tree = await openTree(); + const journal = await openJournal(tree); + const store = await createEventStore({ journal, slices: { counter: counterSlice } }); + await store.dispatch(counterAdded({ amount: 7 })); + await store.dispatch(noteTagged({ tag: 'late' })); + + const causes: Cause[] = []; + store.subscribe((_state, cause) => causes.push(cause)); + await store.registerSlice(notesSlice); + expect(store.getState()).toEqual({ counter: 7, notes: ['late'] }); + expect(causes).toEqual([{ kind: 'slice-joined', name: 'notes' }]); + }); +}); + +describe('reset', () => { + it('refolds a forked branch and keeps subscribers attached', async () => { + const tree = await openTree(); + const store = await openStore(tree); + await store.dispatch(counterAdded({ amount: 1 })); + await store.dispatch(counterAdded({ amount: 2 })); + await store.dispatch(counterAdded({ amount: 4 })); + await store.flush(); + + const forked = tree.createBranch('forked', { from: { branch: 'main', seq: 1 } }); + const causes: Cause[] = []; + store.subscribe((_state, cause) => causes.push(cause)); + await store.reset(journalFromBranch(forked, tree)); + + expect(store.ref.branch).toBe('forked'); + expect(store.slice('counter')).toBe(3); + expect(causes).toEqual([{ kind: 'reset', state: { counter: 3, notes: [] } }]); + + await store.dispatch(counterAdded({ amount: 8 })); + expect(store.slice('counter')).toBe(11); + + const reopened = await createEventStore({ journal: journalFromBranch(forked, tree), slices }); + expect(reopened.slice('counter')).toBe(11); + }); +}); + +describe('storeActor', () => { + function once(actor: ReturnType, type: string): Promise { + return new Promise((resolve) => { + const sub = actor.on(type, (event) => { + sub.unsubscribe(); + resolve(event as T); + }); + }); + } + + it('emits store.ready on start and store.changed after store.append', async () => { + const tree = await openTree(); + const store = await openStore(tree); + const actor = createActor(storeActor, { input: { store } }); + const ready = once<{ type: string }>(actor, 'store.ready'); + const changed = once<{ type: string }>(actor, 'store.changed'); + actor.start(); + actor.send({ type: 'store.append', event: counterAdded({ amount: 2 }) }); + expect((await ready).type).toBe('store.ready'); + expect((await changed).type).toBe('store.changed'); + expect(store.slice('counter')).toBe(2); + actor.stop(); + }); + + it('emits store.reset after store.switch to a forked journal', async () => { + const tree = await openTree(); + const store = await openStore(tree); + await store.dispatch(counterAdded({ amount: 1 })); + await store.dispatch(counterAdded({ amount: 2 })); + await store.flush(); + + const actor = createActor(storeActor, { input: { store } }); + const reset = once<{ type: string; branch: string }>(actor, 'store.reset'); + actor.start(); + const forked = tree.createBranch('forked', { from: { branch: 'main', seq: 0 } }); + actor.send({ type: 'store.switch', journal: journalFromBranch(forked, tree) }); + expect(await reset).toMatchObject({ type: 'store.reset', branch: 'forked' }); + expect(store.slice('counter')).toBe(1); + actor.stop(); + }); + + it('emits store.error when a dispatch fails', async () => { + const tree = await openTree(); + const store = await openStore(tree); + const actor = createActor(storeActor, { input: { store } }); + const failure = once<{ type: string }>(actor, 'store.error'); + actor.start(); + actor.send({ + type: 'store.append', + event: { type: 'test.unregistered_event', time: Date.now() }, + }); + expect((await failure).type).toBe('store.error'); + actor.stop(); + }); +}); diff --git a/packages/agent-core-v2/src/human/test/media/tool.test.ts b/packages/agent-core-v2/src/human/test/media/tool.test.ts index b1a5dae458b..1bc4938fa43 100644 --- a/packages/agent-core-v2/src/human/test/media/tool.test.ts +++ b/packages/agent-core-v2/src/human/test/media/tool.test.ts @@ -7,7 +7,12 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import { createActor, waitFor } from '#/xstate2'; import { createAgentMachine } from '#/agent/machine'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import { createTurnMachine } from '#/agent/turn'; +import { createEventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; import type { ModelCapability } from '#/llm/capability'; import { createAssistantMessage, @@ -48,6 +53,14 @@ function toolCall(id: string, name: string, args: string): ToolCall { return { type: 'function', id, name, arguments: args }; } +async function testStore(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('test'); + tree.createBranch('main'); + return createEventStore({ journal: journalFromBranch(tree.openBranch('main'), tree), slices: agentSlices }); +} + afterEach(() => { for (const dir of tmpDirs.splice(0)) { fs.rmSync(dir, { recursive: true, force: true }); @@ -163,6 +176,7 @@ describe('media stack wiring', () => { }, }; + const agentStore = await testStore(); const actor = createActor( createAgentMachine({ tools, @@ -179,11 +193,11 @@ describe('media stack wiring', () => { }), ), }), - { input: { request: { model } } }, + { input: { request: { model }, store: agentStore } }, ); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('watch this') }); - await waitFor(actor, (s) => s.matches('idle') && s.context.messages.length > 1, { + await waitFor(actor, (s) => s.matches('idle') && agentStore.getState().history.length > 1, { timeout: 5000, }); diff --git a/packages/agent-core-v2/src/human/test/session/machine.test.ts b/packages/agent-core-v2/src/human/test/session/machine.test.ts index fa3c6f7c2ba..947c3490aba 100644 --- a/packages/agent-core-v2/src/human/test/session/machine.test.ts +++ b/packages/agent-core-v2/src/human/test/session/machine.test.ts @@ -12,6 +12,8 @@ import { createLlmMachine } from '#/llm/requester/machine'; import type { LlmRequester } from '#/llm/requester/requester'; import { emptyUsage } from '#/llm/usage'; import { createAgentMachine } from '#/agent/machine'; +import { messageAppended, turnEnded } from '#/agent/events'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import { createAssistantEntry, createTurnMachine, @@ -23,6 +25,12 @@ import { createSessionMachine, type AgentActorRef, } from '#/session/machine'; +import { createEventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; +import type { BranchRef } from '#/store/types'; +import type { Tree } from '#/store/tree'; const model: LlmModel = { provider: 'test', model: 'test-model', capability: UNKNOWN_CAPABILITY }; @@ -54,6 +62,32 @@ function createTestSession(requester: LlmRequester): SessionActor { return session; } +interface TestEnv { + tree: Tree; + open(branch: string, from?: BranchRef): Promise; +} + +async function testEnv(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('test'); + return { + tree, + open: (branch, from) => { + if (!tree.has(branch)) { + tree.createBranch(branch, from !== undefined ? { from } : undefined); + } + return createEventStore({ journal: journalFromBranch(tree.openBranch(branch), tree), slices: agentSlices }); + }, + }; +} + +function forkStore(env: TestEnv, source: AgentEventStore, branch: string): Promise { + const sourceBranch = env.tree.openBranch(source.ref.branch); + const head = sourceBranch.head; + return env.open(branch, head === null ? undefined : { branch: sourceBranch.name, seq: head }); +} + function agentRef(session: SessionActor, agentId: string): AgentActorRef { const entry = session.getSnapshot().context.agents[agentId]; expect(entry).toBeDefined(); @@ -68,10 +102,10 @@ function submit(session: SessionActor, agentId: string, text: string): void { }); } -async function waitIdle(ref: AgentActorRef, messageCount: number) { +async function waitIdle(ref: AgentActorRef, store: AgentEventStore, messageCount: number) { return waitFor( ref, - (snapshot) => snapshot.matches('idle') && snapshot.context.messages.length === messageCount, + (snapshot) => snapshot.matches('idle') && store.getState().history.length === messageCount, { timeout: 5000 }, ); } @@ -85,9 +119,10 @@ describe('session machine agent lifecycle', () => { const session = createTestSession(createEchoRequester()); const created: string[] = []; session.on('agent.created', (event) => created.push(event.agentId)); + const env = await testEnv(); - session.send({ type: 'agent.create' }); - session.send({ type: 'agent.create' }); + session.send({ type: 'agent.create', input: { store: await env.open('agent-1') } }); + session.send({ type: 'agent.create', input: { store: await env.open('agent-2') } }); expect(created).toEqual(['agent-1', 'agent-2']); expect(Object.keys(session.getSnapshot().context.agents).toSorted()).toEqual(['agent-1', 'agent-2']); @@ -95,22 +130,29 @@ describe('session machine agent lifecycle', () => { it('creates a agent with restored messages and turnId', async () => { const session = createTestSession(createEchoRequester()); - const history: HistoryMessage[] = [ - createUserEntry(createUserMessage('old'), { source: 'input' }), - createAssistantEntry(createAssistantMessage([{ type: 'text', text: 'echo:old' }]), { - source: 'llm', - usage: emptyUsage(), + const env = await testEnv(); + const store = await env.open('restored'); + await store.dispatch( + messageAppended({ message: createUserEntry(createUserMessage('old'), { source: 'input' }) }), + ); + await store.dispatch( + messageAppended({ + message: createAssistantEntry(createAssistantMessage([{ type: 'text', text: 'echo:old' }]), { + source: 'llm', + usage: emptyUsage(), + }), }), - ]; - session.send({ type: 'agent.create', agentId: 'restored', input: { history, turnId: 7 } }); + ); + await store.dispatch(turnEnded({ turnId: 6, outcome: 'done' })); + session.send({ type: 'agent.create', agentId: 'restored', input: { store } }); const ref = agentRef(session, 'restored'); - expect(ref.getSnapshot().context.turnId).toBe(7); + expect(store.getState().turnIndex.nextTurnId).toBe(7); submit(session, 'restored', 'new'); - const snapshot = await waitIdle(ref, 4); + await waitIdle(ref, store, 4); - expect(snapshot.context.turnId).toBe(8); - expect(rolesAndTexts(snapshot.context.messages)).toEqual([ + expect(store.getState().turnIndex.nextTurnId).toBe(8); + expect(rolesAndTexts(store.getState().history)).toEqual([ 'user:old', 'assistant:echo:old', 'user:new', @@ -122,10 +164,11 @@ describe('session machine agent lifecycle', () => { const session = createTestSession(createEchoRequester()); const errors: string[] = []; session.on('agent.failed', (event) => errors.push(event.error)); + const env = await testEnv(); - session.send({ type: 'agent.create', agentId: 'a' }); + session.send({ type: 'agent.create', agentId: 'a', input: { store: await env.open('a') } }); const first = agentRef(session, 'a'); - session.send({ type: 'agent.create', agentId: 'a' }); + session.send({ type: 'agent.create', agentId: 'a', input: { store: await env.open('a-dup') } }); expect(errors).toEqual([`duplicate agent id: 'a'`]); expect(agentRef(session, 'a')).toBe(first); @@ -133,7 +176,8 @@ describe('session machine agent lifecycle', () => { it('stops a agent and removes it from the registry', async () => { const session = createTestSession(createEchoRequester()); - session.send({ type: 'agent.create', agentId: 'a' }); + const env = await testEnv(); + session.send({ type: 'agent.create', agentId: 'a', input: { store: await env.open('a') } }); const ref = agentRef(session, 'a'); const stopped: string[] = []; session.on('agent.stopped', (event) => stopped.push(event.agentId)); @@ -176,8 +220,11 @@ describe('session machine concurrent agents', () => { }, }; const session = createTestSession(requester); - session.send({ type: 'agent.create', agentId: 'a' }); - session.send({ type: 'agent.create', agentId: 'b' }); + const env = await testEnv(); + const storeA = await env.open('a'); + const storeB = await env.open('b'); + session.send({ type: 'agent.create', agentId: 'a', input: { store: storeA } }); + session.send({ type: 'agent.create', agentId: 'b', input: { store: storeB } }); submit(session, 'a', 'hello-a'); submit(session, 'b', 'hello-b'); @@ -190,16 +237,16 @@ describe('session machine concurrent agents', () => { resolvers.get('hello-a')?.(); resolvers.get('hello-b')?.(); - const [snapshotA, snapshotB] = await Promise.all([ - waitIdle(agentRef(session, 'a'), 2), - waitIdle(agentRef(session, 'b'), 2), + await Promise.all([ + waitIdle(agentRef(session, 'a'), storeA, 2), + waitIdle(agentRef(session, 'b'), storeB, 2), ]); - expect(rolesAndTexts(snapshotA.context.messages)).toEqual([ + expect(rolesAndTexts(storeA.getState().history)).toEqual([ 'user:hello-a', 'assistant:echo:hello-a', ]); - expect(rolesAndTexts(snapshotB.context.messages)).toEqual([ + expect(rolesAndTexts(storeB.getState().history)).toEqual([ 'user:hello-b', 'assistant:echo:hello-b', ]); @@ -209,112 +256,51 @@ describe('session machine concurrent agents', () => { describe('session machine agent fork', () => { it('forks a agent with the source context and diverges afterwards', async () => { const session = createTestSession(createEchoRequester()); - session.send({ type: 'agent.create', agentId: 'a' }); + const env = await testEnv(); + const storeA = await env.open('a'); + session.send({ type: 'agent.create', agentId: 'a', input: { store: storeA } }); submit(session, 'a', 'hi'); - const snapshotA = await waitIdle(agentRef(session, 'a'), 2); - expect(snapshotA.context.turnId).toBe(1); + await waitIdle(agentRef(session, 'a'), storeA, 2); + expect(storeA.getState().turnIndex.nextTurnId).toBe(1); + await storeA.flush(); const forked: string[] = []; session.on('agent.forked', (event) => forked.push(event.agentId)); - session.send({ type: 'agent.fork', sourceId: 'a', agentId: 'b' }); + const storeB = await forkStore(env, storeA, 'b'); + session.send({ type: 'agent.fork', sourceId: 'a', agentId: 'b', store: storeB }); expect(forked).toEqual(['b']); const refB = agentRef(session, 'b'); - const forkSnapshot = refB.getSnapshot(); - expect(forkSnapshot.value).toEqual({ idle: 'ready' }); - expect(forkSnapshot.context.turnId).toBe(1); - expect(rolesAndTexts(forkSnapshot.context.messages)).toEqual([ + expect(refB.getSnapshot().value).toEqual({ idle: 'ready' }); + expect(storeB.getState().turnIndex.nextTurnId).toBe(1); + expect(rolesAndTexts(storeB.getState().history)).toEqual([ 'user:hi', 'assistant:echo:hi', ]); submit(session, 'b', 'fork-hi'); - const snapshotB = await waitIdle(refB, 4); + await waitIdle(refB, storeB, 4); - expect(snapshotB.context.turnId).toBe(2); - expect(rolesAndTexts(snapshotB.context.messages)).toEqual([ + expect(storeB.getState().turnIndex.nextTurnId).toBe(2); + expect(rolesAndTexts(storeB.getState().history)).toEqual([ 'user:hi', 'assistant:echo:hi', 'user:fork-hi', 'assistant:echo:fork-hi', ]); - expect(agentRef(session, 'a').getSnapshot().context.messages).toHaveLength(2); - expect(agentRef(session, 'a').getSnapshot().context.turnId).toBe(1); + expect(storeA.getState().history).toHaveLength(2); + expect(storeA.getState().turnIndex.nextTurnId).toBe(1); }); it('emits agent.failed when forking an unknown source', async () => { const session = createTestSession(createEchoRequester()); const errors: string[] = []; session.on('agent.failed', (event) => errors.push(event.error)); + const env = await testEnv(); - session.send({ type: 'agent.fork', sourceId: 'nope', agentId: 'b' }); + session.send({ type: 'agent.fork', sourceId: 'nope', agentId: 'b', store: await env.open('b') }); expect(errors).toEqual([`unknown agent: 'nope'`]); expect(session.getSnapshot().context.agents['b']).toBeUndefined(); }); }); - - -describe('session machine agent switch', () => { - it('switches an idle agent to a new branch context and emits agent.switched', async () => { - const session = createTestSession(createEchoRequester()); - session.send({ type: 'agent.create', agentId: 'main' }); - submit(session, 'main', 'hi'); - await waitIdle(agentRef(session, 'main'), 2); - - const switched: Array<{ agentId: string; branchId: string; reason?: string }> = []; - session.on('agent.switched', (event) => { - if (event.type === 'agent.switched') { - switched.push({ agentId: event.agentId, branchId: event.branchId, reason: event.reason }); - } - }); - - session.send({ - type: 'agent.switch', - agentId: 'main', - input: { branchId: 'main~2', history: [], reason: 'undo' }, - }); - - expect(switched).toEqual([{ agentId: 'main', branchId: 'main~2', reason: 'undo' }]); - const snapshot = agentRef(session, 'main').getSnapshot(); - expect(snapshot.context.branchId).toBe('main~2'); - expect(snapshot.context.messages).toEqual([]); - expect(snapshot.context.turnId).toBe(1); - }); - - it('emits agent.failed when switching an unknown or busy agent', async () => { - let release: (() => void) | undefined; - const requester: LlmRequester = { - generate: (_config, _content, { onEvent }) => - new Promise((resolve) => { - release = () => { - onEvent?.({ type: 'llm.streaming.part', part: { type: 'text', text: 'late' } }); - onEvent?.({ type: 'llm.done' }); - resolve(); - }; - }), - }; - const session = createTestSession(requester); - const errors: string[] = []; - session.on('agent.failed', (event) => errors.push(event.error)); - session.send({ type: 'agent.create', agentId: 'main' }); - - session.send({ - type: 'agent.switch', - agentId: 'nope', - input: { branchId: 'x', history: [] }, - }); - - submit(session, 'main', 'hi'); - await vi.waitFor(() => expect(release).toBeDefined()); - session.send({ - type: 'agent.switch', - agentId: 'main', - input: { branchId: 'main~2', history: [] }, - }); - - expect(errors).toEqual([`unknown agent: 'nope'`, `agent is busy: 'main'`]); - expect(agentRef(session, 'main').getSnapshot().context.branchId).toBe('main'); - release?.(); - }); -}); diff --git a/packages/agent-core-v2/src/human/test/session/migrate-v2.test.ts b/packages/agent-core-v2/src/human/test/session/migrate-v2.test.ts index 12efd30fa58..54592a6902e 100644 --- a/packages/agent-core-v2/src/human/test/session/migrate-v2.test.ts +++ b/packages/agent-core-v2/src/human/test/session/migrate-v2.test.ts @@ -12,10 +12,12 @@ import type { LlmModel } from '#/llm/model'; import { createLlmMachine } from '#/llm/requester/machine'; import type { LlmRequester } from '#/llm/requester/requester'; import { createAgentMachine } from '#/agent/machine'; -import { createTurnMachine } from '#/agent/turn'; +import type { StateUpdated } from '#/agent/events'; +import type { AgentEventStore, TurnIndexState } from '#/agent/slices'; +import { createTurnMachine, type HistoryMessage } from '#/agent/turn'; import { createSessionMachine, type AgentActorRef } from '#/session/machine'; -import { undoAgentTurns } from '#/session/undo'; -import { loadSessionState, persistSession } from '#/persist/session'; +import type { SessionStores } from '#/session/stores'; +import { createSlice } from '#/eventStore/slice'; import { openSessionStore } from '#/persist/open'; import { migrateV2Session } from '#/persist/v2/migrate'; @@ -153,10 +155,55 @@ function assistantStep(uuid: string, text: string, agentId = MAIN): Record>> { +const statesSlice = createSlice({ + name: 'states', + initialState: () => ({}) as Record, + reducers: { + 'state.updated': (draft, event: StateUpdated) => { + draft[event.name] = event.value; + }, + }, +}); + +interface LoadedAgent { + agentId: string; + messages: HistoryMessage[]; + turnIndex: TurnIndexState; + states: Record; +} + +function readStates(store: AgentEventStore): Record { + return (store.getState() as unknown as Record>)['states'] ?? {}; +} + +async function loadAgent(stores: SessionStores, agentId: string): Promise { + const store = await stores.open(agentId); + if ((store.getState() as unknown as Record)['states'] === undefined) { + await store.registerSlice(statesSlice); + } + return { + agentId, + messages: store.getState().history, + turnIndex: store.getState().turnIndex, + states: readStates(store), + }; +} + +async function loadAgents(stores: SessionStores): Promise { + const roster = (await stores.session()).getState().roster.agents; + const agents: LoadedAgent[] = []; + for (const agentId of Object.keys(roster).sort()) { + agents.push(await loadAgent(stores, agentId)); + } + return agents; +} + +async function loadMigrated(dir: string) { await migrateV2Session(dir); const opened = await openSessionStore(dir); - return loadSessionState(opened.tree); + const agents = await loadAgents(opened.stores); + const meta = (await opened.stores.session()).getState().sessionMeta.value; + return { agents, meta }; } describe('migrateV2Session', () => { @@ -236,7 +283,7 @@ describe('migrateV2Session', () => { if (second?.message.role === 'assistant') { expect(second.meta.finish).toEqual({ finishReason: 'completed', rawFinishReason: null }); } - expect(agent.turnId).toBe(1); + expect(agent.turnIndex.nextTurnId).toBe(1); expect(loaded.meta).toMatchObject({ id: 'session_basic', version: 2, @@ -523,37 +570,41 @@ describe('migrateV2Session', () => { const first = await openSessionStore(dir); expect(first.migrated).toBe(true); - const loaded = await loadSessionState(first.tree); + const agentStore = await first.stores.open(MAIN); + expect(agentStore.getState().history.map((entry) => extractText(entry.message))).toEqual([ + 'first', + 'first-reply', + ]); + expect(agentStore.getState().turnIndex.nextTurnId).toBe(1); + const session = createTestSession(); - const persistence = persistSession(session, first.tree, { - branches: new Map(loaded.agents.map((agent) => [agent.agentId, agent.branch])), - }); - for (const agent of loaded.agents) { - session.send({ - type: 'agent.create', - agentId: agent.agentId, - input: { history: agent.messages, turnId: agent.turnId, branchId: agent.branch }, - }); - } + session.send({ type: 'agent.create', agentId: MAIN, input: { store: agentStore } }); submit(session, MAIN, 'again'); - await waitFor(agentRef(session, MAIN), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); + await waitFor( + agentRef(session, MAIN), + (s) => s.matches('idle') && agentStore.getState().history.length === 4, + { timeout: 5000 }, + ); + await first.stores.flush(); - const undone = await undoAgentTurns(session, first.tree, MAIN, 1); - expect(undone.messages.map((entry) => extractText(entry.message))).toEqual(['first', 'first-reply']); - await persistence.flush(); - persistence.dispose(); + const undone = await first.stores.undo(MAIN, 1); + expect(undone.branchId).toBe('main~2'); + expect(agentStore.ref.branch).toBe('main~2'); + expect(agentStore.getState().history.map((entry) => extractText(entry.message))).toEqual([ + 'first', + 'first-reply', + 'again', + ]); + expect((await first.stores.session()).getState().roster.agents[MAIN]).toBe('main~2'); + await first.stores.flush(); session.stop(); + await first.stores.dispose(); const second = await openSessionStore(dir); expect(second.migrated).toBe(false); - const reloaded = await loadSessionState(second.tree); - expect(reloaded.agents.map((agent) => agent.agentId)).toEqual([MAIN]); - expect(reloaded.agents[0]!.messages.map((entry) => extractText(entry.message))).toEqual([ - 'first', - 'first-reply', - ]); + const roster = (await second.stores.session()).getState().roster.agents; + expect(Object.keys(roster)).toEqual([MAIN]); + expect(roster[MAIN]).toBe('main~2'); const names = await readdir(dir); expect(names.filter((name) => name.startsWith('.migrate'))).toEqual([]); expect(names).toContain('state.json'); diff --git a/packages/agent-core-v2/src/human/test/session/persist.test.ts b/packages/agent-core-v2/src/human/test/session/persist.test.ts deleted file mode 100644 index ca4a3f4db0d..00000000000 --- a/packages/agent-core-v2/src/human/test/session/persist.test.ts +++ /dev/null @@ -1,500 +0,0 @@ -import { describe, expect, it, vi } from 'vitest'; - -import { createActor, waitFor, type ActorRefFrom } from '#/xstate2'; - -import { UNKNOWN_CAPABILITY } from '#/llm/capability'; -import { createAssistantMessage, createUserMessage, extractText, type AssistantMessage, type Message, type ToolCall } from '#/llm/message'; -import type { LlmModel } from '#/llm/model'; -import { createLlmMachine } from '#/llm/requester/machine'; -import type { LlmRequester } from '#/llm/requester/requester'; -import { createAgentMachine } from '#/agent/machine'; -import { createTurnMachine, toInputMessages, type HistoryMessage } from '#/agent/turn'; -import { - createSessionMachine, - type AgentActorRef, -} from '#/session/machine'; -import { undoAgentTurns } from '#/session/undo'; -import { loadSessionState, persistSession } from '#/persist/session'; -import { collectPluginTools, connectPlugins } from '#/plugin'; -import { createTodoPlugin } from '#/todo/plugin'; -import { restoreTodoState, snapshotTodoState } from '#/todo/state'; -import { MemoryBackend } from '#/store/backend/memory'; -import type { Branch } from '#/store/branch'; -import type { Tree } from '#/store/tree'; -import { TreeStore } from '#/store/store'; - -const model: LlmModel = { provider: 'test', model: 'test-model', capability: UNKNOWN_CAPABILITY }; - -type SessionActor = ActorRefFrom>; - -function toolCall(id: string, name: string, args: string = '{}'): ToolCall { - return { type: 'function', id, name, arguments: args }; -} - -function createStubRequester(responses: readonly AssistantMessage[]): LlmRequester { - let call = 0; - return { - generate: (_config, _content, { onEvent }) => { - const message = responses[Math.min(call, responses.length - 1)]; - call += 1; - for (const part of [...message.content, ...message.toolCalls]) { - onEvent?.({ type: 'llm.streaming.part', part }); - } - onEvent?.({ type: 'llm.done' }); - return Promise.resolve(); - }, - }; -} - -function createEchoRequester(): LlmRequester { - return { - generate: (_config, { messages }, { onEvent }) => { - const last = messages.at(-1); - const text = last !== undefined && last.role === 'user' ? extractText(last) : ''; - onEvent?.({ type: 'llm.streaming.part', part: { type: 'text', text: `echo:${text}` } }); - onEvent?.({ type: 'llm.done' }); - return Promise.resolve(); - }, - }; -} - -function createTestSession(requester: LlmRequester): SessionActor { - const session = createActor( - createSessionMachine({ - agent: createAgentMachine({ - tools: [], - turnActor: createTurnMachine(createLlmMachine({ requester })), - }), - }), - { input: { request: { model } } }, - ); - session.start(); - return session; -} - -function agentRef(session: SessionActor, agentId: string): AgentActorRef { - const entry = session.getSnapshot().context.agents[agentId]; - expect(entry).toBeDefined(); - return (entry as { ref: AgentActorRef }).ref; -} - -function submit(session: SessionActor, agentId: string, text: string): void { - session.send({ - type: 'agent.send', - agentId, - event: { type: 'input.submit', message: createUserMessage(text) }, - }); -} - -function chainMessages(branch: Branch): Message[] { - return toInputMessages( - [...branch.walk()] - .toReversed() - .filter((entry) => entry.type === 'message') - .map((entry) => entry.payload.data as HistoryMessage), - ); -} - -async function restoreSession( - tree: Tree, -): Promise<{ - session: SessionActor; - loaded: Awaited>; - persistence: ReturnType; -}> { - const loaded = await loadSessionState(tree); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree, { - branches: new Map(loaded.agents.map((agent) => [agent.agentId, agent.branch])), - }); - for (const agent of loaded.agents) { - session.send({ - type: 'agent.create', - agentId: agent.agentId, - input: { history: agent.messages, turnId: agent.turnId, branchId: agent.branch }, - }); - } - return { session, loaded, persistence }; -} - -describe('persistSession', () => { - it('persists and restores a full session with a forked agent', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree); - - session.send({ type: 'agent.create', agentId: 'main' }); - submit(session, 'main', 'hi'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - await persistence.flush(); - - session.send({ type: 'agent.fork', sourceId: 'main', agentId: 'fork' }); - submit(session, 'fork', 'fork-hi'); - submit(session, 'main', 'main-hi'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - await waitFor(agentRef(session, 'fork'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - await persistence.flush(); - - expect(tree.branches()).toEqual(['_session', 'fork', 'main']); - const mainBranch = tree.openBranch('main'); - const forkBranch = tree.openBranch('fork'); - expect(forkBranch.header.parentBranch).toBe('main'); - expect(forkBranch.header.parentSeq).toBe(3); - expect(chainMessages(mainBranch).map(extractText)).toEqual([ - 'hi', - 'echo:hi', - 'main-hi', - 'echo:main-hi', - ]); - expect(chainMessages(forkBranch).map(extractText)).toEqual([ - 'hi', - 'echo:hi', - 'fork-hi', - 'echo:fork-hi', - ]); - expect(forkBranch.nextSeq).toBe(4); - - await persistence.flush(); - persistence.dispose(); - session.stop(); - - const reopened = await TreeStore.open(fs); - const { session: restored, loaded, persistence: restoredPersistence } = await restoreSession( - await reopened.tree('sess'), - ); - expect(loaded.agents.map((agent) => agent.agentId).sort()).toEqual(['fork', 'main']); - for (const agent of loaded.agents) { - expect(agent.turnId).toBe(2); - expect(agentRef(restored, agent.agentId).getSnapshot().context.messages).toEqual(agent.messages); - } - - submit(restored, 'main', 'again'); - await waitFor(agentRef(restored, 'main'), (s) => s.matches('idle') && s.context.messages.length === 6, { - timeout: 5000, - }); - await restoredPersistence.flush(); - const mainMessages = chainMessages((await reopened.tree('sess')).openBranch('main')); - expect(mainMessages.map(extractText)).toEqual([ - 'hi', - 'echo:hi', - 'main-hi', - 'echo:main-hi', - 'again', - 'echo:again', - ]); - }); - - it('skips cleanly stopped agents on restore', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree); - - session.send({ type: 'agent.create', agentId: 'main' }); - session.send({ type: 'agent.create', agentId: 'temp' }); - submit(session, 'main', 'hi'); - submit(session, 'temp', 'temp-hi'); - await waitFor(agentRef(session, 'temp'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - session.send({ type: 'agent.stop', agentId: 'temp' }); - await persistence.flush(); - - const logEntries = [...tree.openBranch('_session').walk()] - .toReversed() - .map((entry) => `${entry.type}:${(entry.payload.data as { agentId: string }).agentId}`); - expect(logEntries).toEqual(['agent.open:main', 'agent.open:temp', 'agent.close:temp']); - - persistence.dispose(); - session.stop(); - - const reopened = await TreeStore.open(fs); - const loaded = await loadSessionState(await reopened.tree('sess')); - expect(loaded.agents.map((agent) => agent.agentId)).toEqual(['main']); - }); - - it('gives a recreated anonymous agent a fresh branch', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree); - - session.send({ type: 'agent.create' }); - submit(session, 'agent-1', 'hi'); - await waitFor(agentRef(session, 'agent-1'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - session.send({ type: 'agent.stop', agentId: 'agent-1' }); - await persistence.flush(); - persistence.dispose(); - session.stop(); - - const reopened = await TreeStore.open(fs); - const loaded = await loadSessionState(await reopened.tree('sess')); - expect(loaded.agents).toEqual([]); - - const restored = createTestSession(createEchoRequester()); - const restoredPersistence = persistSession(restored, await reopened.tree('sess')); - restored.send({ type: 'agent.create', input: { branchId: 'agent-1~2' } }); - submit(restored, 'agent-1', 'fresh'); - await waitFor(agentRef(restored, 'agent-1'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - await restoredPersistence.flush(); - - const branches = (await reopened.tree('sess')).branches(); - expect(branches).toEqual(['_session', 'agent-1', 'agent-1~2']); - expect(chainMessages((await reopened.tree('sess')).openBranch('agent-1~2')).map(extractText)).toEqual([ - 'fresh', - 'echo:fresh', - ]); - }); -}); - - -describe('undoAgentTurns', () => { - it('undoes the last turn by switching to a forked branch and restores it after reopen', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree); - - session.send({ type: 'agent.create', agentId: 'main' }); - submit(session, 'main', 'first'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - submit(session, 'main', 'second'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - await persistence.flush(); - - const result = await undoAgentTurns(session, tree, 'main', 1); - expect(result.branchId).toBe('main~2'); - expect(result.turnId).toBe(1); - expect(toInputMessages(result.messages).map(extractText)).toEqual(['first', 'echo:first']); - - const snapshot = agentRef(session, 'main').getSnapshot(); - expect(snapshot.context.branchId).toBe('main~2'); - expect(snapshot.context.turnId).toBe(2); - expect(snapshot.context.messages).toHaveLength(2); - - const undoBranch = tree.openBranch('main~2'); - expect(undoBranch.header.parentBranch).toBe('main'); - expect(undoBranch.header.parentSeq).toBe(3); - expect(chainMessages(tree.openBranch('main')).map(extractText)).toEqual([ - 'first', - 'echo:first', - 'second', - 'echo:second', - ]); - - submit(session, 'main', 'third'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - await persistence.flush(); - expect(chainMessages(undoBranch).map(extractText)).toEqual([ - 'first', - 'echo:first', - 'third', - 'echo:third', - ]); - expect(chainMessages(tree.openBranch('main')).map(extractText)).toHaveLength(4); - - const logEntries = [...tree.openBranch('_session').walk()].toReversed(); - expect(logEntries.map((entry) => entry.type)).toEqual(['agent.open', 'agent.switch']); - expect(logEntries[1]?.payload.data).toMatchObject({ - agentId: 'main', - branch: 'main~2', - reason: 'undo', - from: { branch: 'main', seq: 7 }, - }); - - persistence.dispose(); - session.stop(); - - const reopened = await TreeStore.open(fs); - const { session: restored, loaded } = await restoreSession(await reopened.tree('sess')); - expect(loaded.agents).toHaveLength(1); - expect(loaded.agents[0]?.branch).toBe('main~2'); - expect(loaded.agents[0]?.turnId).toBe(3); - expect(toInputMessages(loaded.agents[0]?.messages ?? []).map(extractText)).toEqual([ - 'first', - 'echo:first', - 'third', - 'echo:third', - ]); - expect(agentRef(restored, 'main').getSnapshot().context.branchId).toBe('main~2'); - }); - - it('rejects undo when anchors are insufficient or the agent is busy', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - let release: (() => void) | undefined; - const requester: LlmRequester = { - generate: (_config, _content, { onEvent }) => - new Promise((resolve) => { - release = () => { - onEvent?.({ type: 'llm.streaming.part', part: { type: 'text', text: 'late' } }); - onEvent?.({ type: 'llm.done' }); - resolve(); - }; - }), - }; - const session = createTestSession(requester); - persistSession(session, tree); - session.send({ type: 'agent.create', agentId: 'main' }); - - await expect(undoAgentTurns(session, tree, 'main', 1)).rejects.toMatchObject({ - reason: 'insufficient', - }); - - submit(session, 'main', 'hi'); - await vi.waitFor(() => expect(release).toBeDefined()); - await expect(undoAgentTurns(session, tree, 'main', 1)).rejects.toMatchObject({ reason: 'busy' }); - release?.(); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle'), { timeout: 5000 }); - - await expect(undoAgentTurns(session, tree, 'main', 5)).rejects.toMatchObject({ - reason: 'insufficient', - }); - await expect(undoAgentTurns(session, tree, 'nope', 1)).rejects.toMatchObject({ - reason: 'unknown-agent', - }); - await expect(undoAgentTurns(session, tree, 'main', 0)).rejects.toMatchObject({ - reason: 'invalid-count', - }); - }); - - it('rolls back plugin state from the undo result', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const todo = createTodoPlugin(); - const tools = collectPluginTools([todo]); - const requester = createStubRequester([ - createAssistantMessage( - [], - [ - toolCall( - 'call-todo-1', - 'TodoList', - JSON.stringify({ todos: [{ title: 'task a', status: 'in_progress' }] }), - ), - ], - ), - createAssistantMessage([{ type: 'text', text: 'noted' }]), - createAssistantMessage( - [], - [ - toolCall( - 'call-todo-2', - 'TodoList', - JSON.stringify({ todos: [{ title: 'task a', status: 'done' }] }), - ), - ], - ), - createAssistantMessage([{ type: 'text', text: 'done' }]), - ]); - const session = createActor( - createSessionMachine({ - agent: createAgentMachine({ - tools, - turnActor: createTurnMachine(createLlmMachine({ requester })), - }), - }), - { input: { request: { model } } }, - ); - session.start(); - persistSession(session, tree, { - states: () => ({ todo: () => snapshotTodoState(todo.state) }), - }); - session.send({ type: 'agent.create', agentId: 'main' }); - connectPlugins(agentRef(session, 'main'), [todo]); - - submit(session, 'main', 'hi'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - expect(todo.state.todos).toEqual([{ title: 'task a', status: 'in_progress' }]); - - submit(session, 'main', 'finish'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 8, { - timeout: 5000, - }); - expect(todo.state.todos).toEqual([{ title: 'task a', status: 'done' }]); - - const result = await undoAgentTurns(session, tree, 'main', 1); - expect(result.states['todo']).toMatchObject({ - todos: [{ title: 'task a', status: 'in_progress' }], - lastWriteTurn: 1, - }); - - Object.assign(todo.state, restoreTodoState(result.states['todo'], result.turnId)); - expect(todo.state.todos).toEqual([{ title: 'task a', status: 'in_progress' }]); - expect(todo.state.currentTurn).toBe(1); - }); - - it('undoes twice by descending the parent chain', async () => { - const fs = new MemoryBackend(); - const store = await TreeStore.open(fs); - const tree = await store.tree('sess'); - const session = createTestSession(createEchoRequester()); - const persistence = persistSession(session, tree); - - session.send({ type: 'agent.create', agentId: 'main' }); - submit(session, 'main', 'one'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 2, { - timeout: 5000, - }); - submit(session, 'main', 'two'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 4, { - timeout: 5000, - }); - submit(session, 'main', 'three'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 6, { - timeout: 5000, - }); - await persistence.flush(); - - const first = await undoAgentTurns(session, tree, 'main', 1); - expect(first.branchId).toBe('main~2'); - expect(toInputMessages(first.messages).map(extractText)).toEqual([ - 'one', - 'echo:one', - 'two', - 'echo:two', - ]); - - submit(session, 'main', 'three-alt'); - await waitFor(agentRef(session, 'main'), (s) => s.matches('idle') && s.context.messages.length === 6, { - timeout: 5000, - }); - await persistence.flush(); - - const second = await undoAgentTurns(session, tree, 'main', 2); - expect(second.branchId).toBe('main~3'); - expect(second.turnId).toBe(1); - expect(toInputMessages(second.messages).map(extractText)).toEqual(['one', 'echo:one']); - const secondBranch = tree.openBranch('main~3'); - expect(secondBranch.header.parentBranch).toBe('main'); - expect(secondBranch.header.parentSeq).toBe(3); - }); -}); diff --git a/packages/agent-core-v2/src/human/test/session/stores.test.ts b/packages/agent-core-v2/src/human/test/session/stores.test.ts new file mode 100644 index 00000000000..35c7674caf3 --- /dev/null +++ b/packages/agent-core-v2/src/human/test/session/stores.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from 'vitest'; +import { createActor, waitFor, type ActorRefFrom } from '#/xstate2'; + +import { UNKNOWN_CAPABILITY } from '#/llm/capability'; +import { createUserMessage, extractText } from '#/llm/message'; +import type { LlmModel } from '#/llm/model'; +import { createLlmMachine } from '#/llm/requester/machine'; +import type { LlmRequester } from '#/llm/requester/requester'; +import { createAgentMachine } from '#/agent/machine'; +import { createTurnMachine } from '#/agent/turn'; +import type { AgentEventStore } from '#/agent/slices'; +import { SessionStores } from '#/session/stores'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; +import type { Tree } from '#/store/tree'; + +const model: LlmModel = { provider: 'test', model: 'test-model', capability: UNKNOWN_CAPABILITY }; + +type AgentActor = ActorRefFrom>; + +function createEchoRequester(): LlmRequester { + return { + generate: (_config, { messages }, { onEvent }) => { + const last = messages.at(-1); + const text = last !== undefined && last.role === 'user' ? extractText(last) : ''; + onEvent?.({ type: 'llm.streaming.part', part: { type: 'text', text: `echo:${text}` } }); + onEvent?.({ type: 'llm.done' }); + return Promise.resolve(); + }, + }; +} + +interface TestEnv { + backend: MemoryBackend; + tree: Tree; + stores: SessionStores; +} + +async function testEnv(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('sess'); + return { backend, tree, stores: new SessionStores(tree, backend) }; +} + +async function reopen(env: TestEnv): Promise { + await env.stores.flush(); + await env.stores.dispose(); + const store = await TreeStore.open(env.backend, {}); + const tree = await store.tree('sess'); + return { backend: env.backend, tree, stores: new SessionStores(tree, env.backend) }; +} + +function startAgent(store: AgentEventStore, requester: LlmRequester = createEchoRequester()): AgentActor { + const actor = createActor( + createAgentMachine({ + tools: [], + turnActor: createTurnMachine(createLlmMachine({ requester })), + }), + { input: { request: { model }, store } }, + ); + actor.start(); + return actor; +} + +async function runTurn(actor: AgentActor, store: AgentEventStore, text: string, historyLength: number): Promise { + actor.send({ type: 'input.submit', message: createUserMessage(text) }); + await waitFor(actor, (s) => s.matches('idle') && store.getState().history.length === historyLength, { + timeout: 5000, + }); +} + +function historyTexts(store: AgentEventStore): string[] { + return store.getState().history.map((entry) => extractText(entry.message)); +} + +describe('SessionStores open/fork', () => { + it('folds history and turnIndex for opened and forked agents, then diverges', async () => { + const env = await testEnv(); + const main = await env.stores.open('main'); + const actor = startAgent(main); + await runTurn(actor, main, 'hi', 2); + await env.stores.flush(); + + const fork = await env.stores.fork('main', 'fork'); + expect(fork.ref.branch).toBe('fork'); + expect(historyTexts(fork)).toEqual(['hi', 'echo:hi']); + expect(fork.getState().turnIndex.nextTurnId).toBe(1); + const forkHeader = env.tree.openBranch('fork').header; + expect(forkHeader.parentBranch).toBe('main'); + expect(forkHeader.parentSeq).toBe(env.tree.openBranch('main').head); + + expect((await env.stores.session()).getState().roster.agents).toEqual({ + fork: 'fork', + main: 'main', + }); + + const forkActor = startAgent(fork); + await runTurn(forkActor, fork, 'fork-hi', 4); + await runTurn(actor, main, 'main-hi', 4); + + expect(historyTexts(fork)).toEqual(['hi', 'echo:hi', 'fork-hi', 'echo:fork-hi']); + expect(historyTexts(main)).toEqual(['hi', 'echo:hi', 'main-hi', 'echo:main-hi']); + expect(fork.getState().turnIndex.nextTurnId).toBe(2); + expect(main.getState().turnIndex.nextTurnId).toBe(2); + + forkActor.stop(); + actor.stop(); + }); + + it('removes the agent from the roster on close', async () => { + const env = await testEnv(); + await env.stores.open('main'); + await env.stores.open('temp'); + expect((await env.stores.session()).getState().roster.agents).toEqual({ + main: 'main', + temp: 'temp', + }); + + await env.stores.close('temp'); + + expect(env.stores.get('temp')).toBeUndefined(); + expect((await env.stores.session()).getState().roster.agents).toEqual({ main: 'main' }); + }); +}); + +describe('SessionStores undo', () => { + it('rolls back to the turn boundary, forks with a parent ref, and updates the roster', async () => { + const env = await testEnv(); + const main = await env.stores.open('main'); + const actor = startAgent(main); + await runTurn(actor, main, 'first', 2); + await runTurn(actor, main, 'second', 4); + await env.stores.flush(); + const cutStart = main.getState().turnIndex.turns.at(-1)?.start; + expect(cutStart).toBeDefined(); + + const result = await env.stores.undo('main', 1); + + expect(result.branchId).toBe('main~2'); + expect(main.ref.branch).toBe('main~2'); + expect(historyTexts(main)).toEqual(['first', 'echo:first', 'second']); + expect(main.getState().queue).toEqual([]); + expect(main.getState().turnIndex.turns).toHaveLength(1); + expect(main.getState().turnIndex.nextTurnId).toBe(1); + const header = env.tree.openBranch('main~2').header; + expect(header.parentBranch).toBe('main'); + expect(header.parentSeq).toBe((cutStart as { seq: number }).seq - 1); + expect((await env.stores.session()).getState().roster.agents['main']).toBe('main~2'); + expect(env.tree.openBranch('main').head).toBe(11); + + await waitFor(actor, (s) => s.matches('idle'), { timeout: 5000 }); + await runTurn(actor, main, 'third', 5); + expect(historyTexts(main)).toEqual(['first', 'echo:first', 'second', 'third', 'echo:third']); + expect(env.tree.openBranch('main').head).toBe(11); + expect(main.getState().turnIndex.nextTurnId).toBe(2); + + actor.stop(); + }); + + it('rejects invalid counts, unknown agents, and insufficient turns', async () => { + const env = await testEnv(); + const main = await env.stores.open('main'); + const actor = startAgent(main); + await runTurn(actor, main, 'hi', 2); + + await expect(env.stores.undo('main', 2)).rejects.toMatchObject({ reason: 'insufficient' }); + await expect(env.stores.undo('nope', 1)).rejects.toMatchObject({ reason: 'unknown-agent' }); + await expect(env.stores.undo('main', 0)).rejects.toMatchObject({ reason: 'invalid-count' }); + + actor.stop(); + }); +}); + +describe('SessionStores reopen', () => { + it('restores agent state from the branch after reopen', async () => { + const env = await testEnv(); + const main = await env.stores.open('main'); + const actor = startAgent(main); + await runTurn(actor, main, 'first', 2); + await runTurn(actor, main, 'second', 4); + await env.stores.flush(); + const fork = await env.stores.fork('main', 'fork'); + const forkActor = startAgent(fork); + await runTurn(forkActor, fork, 'fork-hi', 6); + forkActor.stop(); + actor.stop(); + + const restored = await reopen(env); + + expect((await restored.stores.session()).getState().roster.agents).toEqual({ + fork: 'fork', + main: 'main', + }); + const restoredMain = await restored.stores.open('main'); + expect(historyTexts(restoredMain)).toEqual(['first', 'echo:first', 'second', 'echo:second']); + expect(restoredMain.getState().turnIndex.nextTurnId).toBe(2); + const restoredFork = await restored.stores.open('fork'); + expect(historyTexts(restoredFork)).toEqual([ + 'first', + 'echo:first', + 'second', + 'echo:second', + 'fork-hi', + 'echo:fork-hi', + ]); + expect(restoredFork.getState().turnIndex.nextTurnId).toBe(3); + + const actor2 = startAgent(restoredMain); + await runTurn(actor2, restoredMain, 'again', 6); + expect(historyTexts(restoredMain)).toEqual([ + 'first', + 'echo:first', + 'second', + 'echo:second', + 'again', + 'echo:again', + ]); + expect(restoredMain.getState().turnIndex.nextTurnId).toBe(3); + + actor2.stop(); + }); +}); diff --git a/packages/agent-core-v2/src/human/test/todo/plugin.test.ts b/packages/agent-core-v2/src/human/test/todo/plugin.test.ts index c74febeab19..34a687f0ec2 100644 --- a/packages/agent-core-v2/src/human/test/todo/plugin.test.ts +++ b/packages/agent-core-v2/src/human/test/todo/plugin.test.ts @@ -1,11 +1,23 @@ import { describe, expect, it } from 'vitest'; +import { turnStarted } from '#/agent/events'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; +import { createEventStoreSync } from '#/eventStore/eventStore'; +import { memoryJournal } from '#/eventStore/journal'; import { extractText, type SystemMessage, type ToolCall, type UserMessage } from '#/llm/message'; import type { AgentEmitted } from '#/agent/machine'; import { connectPlugins, type AgentPluginTarget, type Plugin } from '#/plugin'; import type { ToolExecuteInput } from '#/tool/executor'; import type { ToolDefinition } from '#/tool/tool'; import { createTodoPlugin, type TodoPlugin } from '#/todo/plugin'; +import { todoSlice } from '#/todo/slice'; + +function testStore(): AgentEventStore { + return createEventStoreSync({ + journal: memoryJournal(), + slices: { ...agentSlices, todo: todoSlice }, + }) as AgentEventStore; +} function toolCall(args: unknown): ToolCall { return { type: 'function', id: 'call-1', name: 'TodoList', arguments: JSON.stringify(args) }; @@ -21,7 +33,7 @@ function pluginTool(plugin: TodoPlugin): ToolDefinition { return tool; } -function createTarget() { +function createTarget(store: AgentEventStore) { const handlers: ((event: AgentEmitted) => void)[] = []; const notified: UserMessage[] = []; const reminded: { key: string; message: UserMessage | SystemMessage }[] = []; @@ -37,9 +49,12 @@ function createTarget() { reminded.push({ key, message }); }, }; - const turnStart = () => { + let turnId = 0; + const turnStart = async () => { + await store.dispatch(turnStarted({ turnId })); + turnId += 1; for (const handler of handlers) { - handler({ type: 'turn.started', turnId: handlers.length, branchId: 'main' }); + handler({ type: 'turn.started', turnId, branchId: 'main' }); } }; return { target, notified, reminded, turnStart }; @@ -47,13 +62,13 @@ function createTarget() { describe('todo plugin tool', () => { it('reads an empty list', async () => { - const plugin = createTodoPlugin(); + const plugin = createTodoPlugin(testStore()); const result = await pluginTool(plugin).execute(executeInput({})); expect(result.content).toEqual([{ type: 'text', text: 'Todo list is empty.' }]); }); it('replaces the list and reads it back', async () => { - const plugin = createTodoPlugin(); + const plugin = createTodoPlugin(testStore()); const result = await pluginTool(plugin).execute( executeInput({ todos: [ @@ -76,7 +91,7 @@ describe('todo plugin tool', () => { }); it('clears the list with an empty array', async () => { - const plugin = createTodoPlugin(); + const plugin = createTodoPlugin(testStore()); await pluginTool(plugin).execute(executeInput({ todos: [{ title: 'task a', status: 'pending' }] })); const result = await pluginTool(plugin).execute(executeInput({ todos: [] })); expect(result.content).toEqual([{ type: 'text', text: 'Todo list cleared.' }]); @@ -86,7 +101,7 @@ describe('todo plugin tool', () => { }); it('drops malformed items on write', async () => { - const plugin = createTodoPlugin(); + const plugin = createTodoPlugin(testStore()); await pluginTool(plugin).execute( executeInput({ todos: [{ title: 'task a' }, { title: 'task b', status: 'done' }, 'junk'], @@ -101,46 +116,49 @@ describe('todo plugin tool', () => { describe('todo plugin reminder', () => { it('notifies once when the list goes stale', async () => { - const plugin = createTodoPlugin(); - const { target, notified, turnStart } = createTarget(); + const store = testStore(); + const plugin = createTodoPlugin(store); + const { target, notified, turnStart } = createTarget(store); plugin.connect?.(target); - turnStart(); + await turnStart(); await pluginTool(plugin).execute(executeInput({ todos: [{ title: 'task a', status: 'pending' }] })); - turnStart(); + await turnStart(); expect(notified).toHaveLength(0); - turnStart(); + await turnStart(); expect(notified).toHaveLength(1); const text = extractText(notified[0]); expect(text).toContain(''); expect(text).toContain('[pending] task a'); - turnStart(); + await turnStart(); expect(notified).toHaveLength(1); }); it('stays silent when every item is done', async () => { - const plugin = createTodoPlugin(); - const { target, notified, turnStart } = createTarget(); + const store = testStore(); + const plugin = createTodoPlugin(store); + const { target, notified, turnStart } = createTarget(store); plugin.connect?.(target); - turnStart(); + await turnStart(); await pluginTool(plugin).execute(executeInput({ todos: [{ title: 'task a', status: 'done' }] })); - turnStart(); - turnStart(); + await turnStart(); + await turnStart(); expect(notified).toHaveLength(0); }); - it('stays silent while the list is empty', () => { - const plugin = createTodoPlugin(); - const { target, notified, turnStart } = createTarget(); + it('stays silent while the list is empty', async () => { + const store = testStore(); + const plugin = createTodoPlugin(store); + const { target, notified, turnStart } = createTarget(store); plugin.connect?.(target); - turnStart(); - turnStart(); - turnStart(); + await turnStart(); + await turnStart(); + await turnStart(); expect(notified).toHaveLength(0); }); }); @@ -148,7 +166,7 @@ describe('todo plugin reminder', () => { describe('connectPlugins notify channel', () => { it('maps target.notify to an input.notify event on the actor', () => { const sent: unknown[] = []; - const plugin = createTodoPlugin(); + const plugin = createTodoPlugin(testStore()); const probe: Plugin = { name: 'probe', tools: () => [], diff --git a/packages/agent-core-v2/src/human/test/tool-select/plugin.test.ts b/packages/agent-core-v2/src/human/test/tool-select/plugin.test.ts index ddfdb609b14..731bca0d5e2 100644 --- a/packages/agent-core-v2/src/human/test/tool-select/plugin.test.ts +++ b/packages/agent-core-v2/src/human/test/tool-select/plugin.test.ts @@ -15,7 +15,12 @@ import { createLlmMachine } from '#/llm/requester/machine'; import type { LlmRequestConfig, LlmRequester, LlmRequestEvent } from '#/llm/requester/requester'; import { connectPlugins, type AgentPluginTarget } from '#/plugin'; import { createAgentMachine, type AgentEmitted } from '#/agent/machine'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import { createTurnMachine, type HistoryMessage } from '#/agent/turn'; +import { createEventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; import type { ToolExecuteInput } from '#/tool/executor'; import { defineTool, type ToolDefinition } from '#/tool/tool'; import { @@ -41,6 +46,14 @@ function executeInput(name: string, args: unknown): ToolExecuteInput { return { toolCall: toolCall(name, args), signal: new AbortController().signal }; } +async function testStore(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('test'); + tree.createBranch('main'); + return createEventStore({ journal: journalFromBranch(tree.openBranch('main'), tree), slices: agentSlices }); +} + function weatherTool(execute?: ToolDefinition['execute']): ToolDefinition { return defineTool({ name: 'get_weather', @@ -288,19 +301,20 @@ describe('tool select agent flow', () => { }), state, ); + const store = await testStore(); const actor = createActor( createAgentMachine({ tools: [createSelectToolsTool(state), deferred], turnActor: createTurnMachine(createLlmMachine({ requester })), }), - { input: { request: { model } } }, + { input: { request: { model }, store } }, ); connectPlugins(actor, [plugin]); actor.start(); actor.send({ type: 'input.submit', message: createUserMessage('weather?') }); - const snapshot = await waitFor( + await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length > 1, + (s) => s.matches('idle') && store.getState().history.length > 1, { timeout: 5000 }, ); @@ -309,7 +323,7 @@ describe('tool select agent flow', () => { expect(firstTools).not.toContain('get_weather'); expect(executed).toEqual(['get_weather']); - const schemaEntry = snapshot.context.messages.find( + const schemaEntry = store.getState().history.find( (entry: HistoryMessage) => entry.message.role === 'system', ); expect(schemaEntry?.meta.key).toBe(DYNAMIC_TOOL_SCHEMA_REMINDER_KEY); diff --git a/packages/agent-core-v2/src/human/test/usage/machine.test.ts b/packages/agent-core-v2/src/human/test/usage/machine.test.ts index 598d2f3197a..1e5271fd565 100644 --- a/packages/agent-core-v2/src/human/test/usage/machine.test.ts +++ b/packages/agent-core-v2/src/human/test/usage/machine.test.ts @@ -9,7 +9,12 @@ import { createLlmMachine } from '#/llm/requester/machine'; import type { LlmRequester } from '#/llm/requester/requester'; import type { TokenUsage } from '#/llm/usage'; import { createAgentMachine } from '#/agent/machine'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import { createTurnMachine } from '#/agent/turn'; +import { createEventStore } from '#/eventStore/eventStore'; +import { journalFromBranch } from '#/eventStore/journal'; +import { MemoryBackend } from '#/store/backend/memory'; +import { TreeStore } from '#/store/store'; import { createUsageMachine } from '#/usage/machine'; import type { UsageEmitted } from '#/usage/machine'; import { createUsagePlugin } from '#/usage/plugin'; @@ -34,6 +39,14 @@ function record( return { usage: usage(inputOther, output), model: extra?.model, turnId: extra?.turnId, at: 0 }; } +async function testStore(): Promise { + const backend = new MemoryBackend(); + const store = await TreeStore.open(backend, {}); + const tree = await store.tree('test'); + tree.createBranch('main'); + return createEventStore({ journal: journalFromBranch(tree.openBranch('main'), tree), slices: agentSlices }); +} + describe('xstate inspection collector', () => { it('publishes JSON-safe scalar envelopes with no machine context', () => { const envelopes: XstateInspectionEnvelope[] = []; @@ -117,11 +130,12 @@ describe('usage plugin', () => { }; const plugin = createUsagePlugin({ model }); const timingPlugin = createTimingPlugin({ now: () => ticks.shift() ?? Number.NaN }); + const store = await testStore(); const actor = createActor( createAgentMachine({ turnActor: createTurnMachine(createLlmMachine({ requester })), }), - { input: { request: { model } } }, + { input: { request: { model }, store } }, ); connectPlugins(actor, [plugin, timingPlugin]); actor.start(); @@ -129,7 +143,7 @@ describe('usage plugin', () => { actor.send({ type: 'input.submit', message: createUserMessage('again') }); await waitFor( actor, - (s) => s.matches('idle') && s.context.messages.length === 4, + (s) => s.matches('idle') && store.getState().history.length === 4, { timeout: 5000 }, ); diff --git a/packages/agent-core-v2/src/human/todo/plugin.ts b/packages/agent-core-v2/src/human/todo/plugin.ts index 0847ab14055..139c9d250ab 100644 --- a/packages/agent-core-v2/src/human/todo/plugin.ts +++ b/packages/agent-core-v2/src/human/todo/plugin.ts @@ -1,7 +1,8 @@ +import type { AgentEventStore } from '#/agent/slices'; import { createUserMessage } from '#/llm/message'; import type { Plugin } from '#/plugin'; -import { createTodoState, type TodoState } from './state'; +import { readTodoState } from './slice'; import { createTodoListTool } from './tool'; import { renderTodoList } from './todoItem'; @@ -9,26 +10,24 @@ const STALE_TURNS = 2; export interface TodoPlugin extends Plugin { readonly name: 'todo'; - readonly state: TodoState; } -export function createTodoPlugin(state: TodoState = createTodoState()): TodoPlugin { - const tool = createTodoListTool(state); +export function createTodoPlugin(store: AgentEventStore): TodoPlugin { + const tool = createTodoListTool(store); return { name: 'todo', - state, tools: () => [tool], connect(target) { if (target.kind !== 'agent') return; target.on('turn.started', (event) => { if (event.type !== 'turn.started') return; - state.currentTurn += 1; - if (state.todos.length === 0) return; - if (state.todos.every((todo) => todo.status === 'done')) return; - if (state.currentTurn - state.lastWriteTurn !== STALE_TURNS) return; + const { todos, currentTurn, lastWriteTurn } = readTodoState(store); + if (todos.length === 0) return; + if (todos.every((todo) => todo.status === 'done')) return; + if (currentTurn - lastWriteTurn !== STALE_TURNS) return; target.notify( createUserMessage( - `\nThe todo list has not been updated recently. If the work is still in progress, update the list to reflect the current progress.\n${renderTodoList(state.todos)}\n`, + `\nThe todo list has not been updated recently. If the work is still in progress, update the list to reflect the current progress.\n${renderTodoList(todos)}\n`, ), ); }); diff --git a/packages/agent-core-v2/src/human/todo/slice.ts b/packages/agent-core-v2/src/human/todo/slice.ts new file mode 100644 index 00000000000..b7113e77c40 --- /dev/null +++ b/packages/agent-core-v2/src/human/todo/slice.ts @@ -0,0 +1,34 @@ +import { createSlice } from '#/eventStore/slice'; + +import type { StateUpdated, TurnStarted } from '#/agent/events'; + +import { readTodoItems, type TodoItem } from './todoItem'; + +export interface TodoSliceState { + todos: readonly TodoItem[]; + currentTurn: number; + lastWriteTurn: number; +} + +export function readTodoState(store: { slice(name: string): unknown }): TodoSliceState { + const state = store.slice('todo') as TodoSliceState | undefined; + return state ?? { todos: [], currentTurn: 0, lastWriteTurn: 0 }; +} + +export const todoSlice = createSlice({ + name: 'todo', + initialState: (): TodoSliceState => ({ todos: [], currentTurn: 0, lastWriteTurn: 0 }), + reducers: { + 'state.updated': (draft, event: StateUpdated) => { + if (event.name !== 'todo') return; + const value = event.value as { todos?: unknown; lastWriteTurn?: unknown }; + draft.todos = readTodoItems(value.todos); + if (typeof value.lastWriteTurn === 'number') { + draft.lastWriteTurn = value.lastWriteTurn; + } + }, + 'turn.started': (draft, _event: TurnStarted) => { + draft.currentTurn += 1; + }, + }, +}); diff --git a/packages/agent-core-v2/src/human/todo/state.ts b/packages/agent-core-v2/src/human/todo/state.ts deleted file mode 100644 index dae96de98bb..00000000000 --- a/packages/agent-core-v2/src/human/todo/state.ts +++ /dev/null @@ -1,33 +0,0 @@ -import { readTodoItems, type TodoItem } from './todoItem'; - -export interface TodoState { - todos: readonly TodoItem[]; - currentTurn: number; - lastWriteTurn: number; -} - -export function createTodoState(): TodoState { - return { todos: [], currentTurn: 0, lastWriteTurn: 0 }; -} - -export interface PersistedTodoState { - todos: readonly TodoItem[]; - lastWriteTurn: number; -} - -export function snapshotTodoState(state: TodoState): PersistedTodoState { - return { todos: state.todos, lastWriteTurn: state.lastWriteTurn }; -} - -export function restoreTodoState(value: unknown, currentTurn: number): TodoState { - const state = createTodoState(); - state.currentTurn = currentTurn; - if (typeof value === 'object' && value !== null) { - const record = value as Record; - state.todos = readTodoItems(record['todos']); - if (typeof record['lastWriteTurn'] === 'number') { - state.lastWriteTurn = record['lastWriteTurn']; - } - } - return state; -} diff --git a/packages/agent-core-v2/src/human/todo/tool.ts b/packages/agent-core-v2/src/human/todo/tool.ts index df232d5f255..b042a40ce20 100644 --- a/packages/agent-core-v2/src/human/todo/tool.ts +++ b/packages/agent-core-v2/src/human/todo/tool.ts @@ -1,11 +1,13 @@ +import { stateUpdated } from '#/agent/events'; +import type { AgentEventStore } from '#/agent/slices'; import { defineTool, type ToolDefinition } from '#/tool/tool'; -import type { TodoState } from './state'; import { readTodoItems, renderTodoList, TODO_LIST_TOOL_NAME } from './todoItem'; +import { readTodoState } from './slice'; import DESCRIPTION from './todo-list.md?raw'; import TODO_LIST_WRITE_REMINDER from './todo-list-write-reminder.md?raw'; -export function createTodoListTool(state: TodoState): ToolDefinition { +export function createTodoListTool(store: AgentEventStore): ToolDefinition { return defineTool({ name: TODO_LIST_TOOL_NAME, description: DESCRIPTION, @@ -34,11 +36,15 @@ export function createTodoListTool(state: TodoState): ToolDefinition { async execute({ toolCall }) { const args = JSON.parse(toolCall.arguments ?? '{}') as { todos?: unknown }; if (args.todos === undefined) { - return { content: [{ type: 'text', text: renderTodoList(state.todos) }] }; + return { content: [{ type: 'text', text: renderTodoList(readTodoState(store).todos) }] }; } const next = readTodoItems(args.todos); - state.todos = next; - state.lastWriteTurn = state.currentTurn; + await store.dispatch( + stateUpdated({ + name: 'todo', + value: { todos: next, lastWriteTurn: readTodoState(store).currentTurn }, + }), + ); if (next.length === 0) { return { content: [{ type: 'text', text: 'Todo list cleared.' }] }; } From 6df3cee8f3f2f601eaf6888ae91833e431371fff Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Wed, 9 Sep 2026 23:50:02 +0800 Subject: [PATCH 2/3] refactor(agent-core-v2): make the agent machine self-consistent and the store actor static Phantom-turn root cause: local patches plus the async store.changed mirror were two consistency models for the same data. The machine now owns its state (every transition is a local assign plus a store.append dual write) and consumes only store.ready/reset/error. - store actor becomes static storeActor logic with the engine passed via invoke input; the session machine registers agentActor and spawns it by name - SessionStores.open/fork return the engine directly - store.reset aborts the old abort scope (cascades to in-flight turn/tools/background) - no-op drain events are no longer journaled - engine bridge seeds the initial turn id into the journal before folding, and resetHistory is awaited before the retry notify - tests pass the store via machine input and wait for the log to catch up before store assertions; drop debug residue --- .../src/agent/loop/loopService.ts | 2 +- .../src/agent/loop/machine/engine.ts | 26 ++-- .../agent-core-v2/src/human/agent/machine.ts | 131 ++++++++++++------ .../agent-core-v2/src/human/agent/slices.ts | 2 - .../src/human/eventStore/actor.ts | 64 +++++---- .../agent-core-v2/src/human/session/stores.ts | 16 +-- .../agent-core-v2/src/human/todo/slice.ts | 6 +- 7 files changed, 145 insertions(+), 102 deletions(-) diff --git a/packages/agent-core-v2/src/agent/loop/loopService.ts b/packages/agent-core-v2/src/agent/loop/loopService.ts index 2b33fe150d8..0b4ec99c1a2 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -1314,7 +1314,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { turn.interruptStep = undefined; if (turn.retryRequested) { turn.retryRequested = false; - this.machineEngine().resetHistory(historyFromContext(this.context.get())); + await this.machineEngine().resetHistory(historyFromContext(this.context.get())); this.machineEngine().notify(EMPTY_MACHINE_PROMPT); } return; diff --git a/packages/agent-core-v2/src/agent/loop/machine/engine.ts b/packages/agent-core-v2/src/agent/loop/machine/engine.ts index f110df1b399..629f2af3ce3 100644 --- a/packages/agent-core-v2/src/agent/loop/machine/engine.ts +++ b/packages/agent-core-v2/src/agent/loop/machine/engine.ts @@ -161,7 +161,7 @@ export interface MachineEngine { remind(key: string, message: UserMessage): void; cancelQueueItem(id: string): void; abort(): void; - resetHistory(history: readonly HistoryMessage[]): void; + resetHistory(history: readonly HistoryMessage[]): Promise; stop(): void; snapshot(): MachineEngineSnapshot; currentStep(): number; @@ -268,11 +268,16 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin publish({ type: 'toolBatchFailed', error }); }, }); - const store: AgentEventStore = createEventStoreSync({ journal: memoryJournal(), slices: agentSlices }); + const journal = memoryJournal(); const initialTurnId = options.initialTurnId ?? 0; if (initialTurnId > 0) { - void store.dispatch(turnEnded({ turnId: initialTurnId - 1, outcome: 'done' })); + void journal.append({ + type: turnEnded.type, + kind: 'event', + data: turnEnded({ turnId: initialTurnId - 1, outcome: 'done' }), + }); } + const store: AgentEventStore = createEventStoreSync({ journal, slices: agentSlices }); const actor = createActor( createAgentMachine({ tools: tools.tools, @@ -446,7 +451,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin for (const message of history) { void journal.append({ type: messageAppended.type, kind: 'event', data: messageAppended({ message }) }); } - const nextTurnId = store.slice('turnIndex').nextTurnId; + const nextTurnId = (actor.getSnapshot() as unknown as MachineSnapshotLike).context.turnId; if (nextTurnId > 0) { void journal.append({ type: turnEnded.type, @@ -454,7 +459,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin data: turnEnded({ turnId: nextTurnId - 1, outcome: 'done' }), }); } - void store.reset(journal); + return store.reset(journal); }, stop: () => { for (const subscription of subscriptions) subscription.unsubscribe(); @@ -463,7 +468,6 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin snapshot: () => { const snapshot = actor.getSnapshot() as unknown as MachineSnapshotLike; const value = snapshot.value; - const storeState = store.getState(); const turnRef = snapshot.children['turn']; let turn: MachineEngineTurnSnapshot | undefined; if (turnRef !== undefined) { @@ -477,7 +481,7 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin : ('running' as const); const context = turnSnapshot.context; turn = { - turnId: storeState.turnIndex.nextTurnId, + turnId: snapshot.context.turnId, phase, step: context.steps, retry: @@ -501,10 +505,10 @@ export function createMachineEngine(options: CreateMachineEngineOptions): Machin waitingForBackground: typeof value === 'object' && value !== null && 'idle' in value && (value as { idle?: unknown }).idle === 'waiting', - queueLength: storeState.queue.length, - queueIds: storeState.queue.map((entry) => entry.id), - notificationCount: storeState.notifications.length, - reminderCount: storeState.reminders.length, + queueLength: snapshot.context.queue.length, + queueIds: snapshot.context.queue.map((entry) => entry.id), + notificationCount: snapshot.context.notifications.length, + reminderCount: snapshot.context.reminders.length, backgroundCount: Object.keys(snapshot.context.background).length, turn, }; diff --git a/packages/agent-core-v2/src/human/agent/machine.ts b/packages/agent-core-v2/src/human/agent/machine.ts index 35eb923acad..3570a17fc6c 100644 --- a/packages/agent-core-v2/src/human/agent/machine.ts +++ b/packages/agent-core-v2/src/human/agent/machine.ts @@ -21,15 +21,17 @@ import { turnEnded, turnStarted, } from './events'; -import { createUserEntry } from './turn'; +import { createSystemEntry, createUserEntry } from './turn'; import { createAbortScope, withAbort, type AbortScope } from '#/utils/abort'; import type { createTurnMachine, HistoryMessage, TurnLlmEvent, TurnOutput, UserEntry } from './turn'; -import type { AgentStoreActorLogic, AgentStoreState, QueuedPrompt } from './slices'; +import { storeActor } from '#/eventStore/actor'; +import type { AgentEventStore, AgentStoreState, QueuedPrompt } from './slices'; export type { QueuedPrompt } from './slices'; export interface AgentInput { request: LlmRequestConfig; + store: AgentEventStore; } export type AgentEvent = @@ -202,7 +204,6 @@ function mirrorPatch(state: AgentStoreState): Pick< export interface CreateAgentMachineOptions { tools?: readonly ToolDefinition[]; turnActor: ReturnType; - store: AgentStoreActorLogic; abortTimeoutMs?: number; maxStepsPerTurn?: number; } @@ -232,7 +233,6 @@ function dispatchTools(tools: readonly ToolDefinition[]): ToolExecutor { export function createAgentMachine({ tools, turnActor, - store, abortTimeoutMs, maxStepsPerTurn, }: CreateAgentMachineOptions) { @@ -247,7 +247,7 @@ export function createAgentMachine({ actors: { turnActor, toolActor: createToolMachine(executor), - storeActor: store, + storeActor, controllerGuard: fromCallback( ({ input }) => () => @@ -271,11 +271,9 @@ export function createAgentMachine({ }; }), emitReset: emit(({ context }) => ({ type: 'context.reset' as const, branchId: context.branchId })), - abortBackgroundScopes: enqueueActions(({ context }) => { - for (const entry of Object.values(context.background)) { - entry.scope.abort(); - } - }), + abortScope: ({ context }) => { + context.scope.abort(); + }, spawnTurnTools: assign(({ context, spawn, self, event }) => { if (event.type !== 'turn.spawn_tools') { return {}; @@ -347,31 +345,69 @@ export function createAgentMachine({ { id: 'store', src: 'storeActor', + input: ({ context }) => ({ store: context.input.store }), }, ], on: { 'input.submit': { - actions: sendTo('store', ({ event }) => ({ - type: 'store.append' as const, - event: inputSubmitted({ id: event.id, message: event.message }), - })), + actions: [ + assign(({ context, event }) => { + if (event.type !== 'input.submit') return {}; + return { queue: [...context.queue, { id: event.id, message: event.message }] }; + }), + sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputSubmitted({ id: event.id, message: event.message }), + })), + ], }, 'input.notify': { - actions: sendTo('store', ({ event }) => ({ - type: 'store.append' as const, - event: inputNotified({ message: event.message }), - })), + actions: [ + assign(({ context, event }) => { + if (event.type !== 'input.notify') return {}; + return { + notifications: [ + ...context.notifications, + createUserEntry(event.message, { source: 'notify' }), + ], + }; + }), + sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputNotified({ message: event.message }), + })), + ], }, 'input.remind': { - actions: sendTo('store', ({ event }) => ({ - type: 'store.append' as const, - event: inputReminded({ key: event.key, message: event.message }), - })), + actions: [ + assign(({ context, event }) => { + if (event.type !== 'input.remind') return {}; + const kept = context.reminders.filter((entry) => entry.meta.key !== event.key); + kept.push( + event.message.role === 'system' + ? createSystemEntry(event.message, { source: 'reminder', key: event.key }) + : createUserEntry(event.message, { source: 'reminder', key: event.key }), + ); + return { reminders: kept }; + }), + sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputReminded({ key: event.key, message: event.message }), + })), + ], }, 'input.steer': { actions: enqueueActions(({ context, event, enqueue }) => { + if (event.type !== 'input.steer') return; const entry = context.queue.find((item) => item.id === event.id); if (entry === undefined) return; + enqueue.assign({ + queue: context.queue.filter((item) => item.id !== event.id), + notifications: [ + ...context.notifications, + createUserEntry(entry.message, { source: 'input' }), + ], + }); enqueue.sendTo('store', { type: 'store.append' as const, event: inputSteered({ id: event.id, message: entry.message }), @@ -379,17 +415,20 @@ export function createAgentMachine({ }), }, 'input.cancel': { - actions: sendTo('store', ({ event }) => ({ - type: 'store.append' as const, - event: inputCancelled({ id: event.id }), - })), - }, - 'store.changed': { - actions: assign(({ event }) => mirrorPatch(event.state)), + actions: [ + assign(({ context, event }) => { + if (event.type !== 'input.cancel') return {}; + return { queue: context.queue.filter((item) => item.id !== event.id) }; + }), + sendTo('store', ({ event }) => ({ + type: 'store.append' as const, + event: inputCancelled({ id: event.id }), + })), + ], }, 'store.reset': { target: '.idle', - actions: ['abortBackgroundScopes', 'resetMirror', 'emitReset', 'forwardToParent'], + actions: ['abortScope', 'resetMirror', 'emitReset', 'forwardToParent'], }, 'store.error': { actions: 'forwardToParent', @@ -441,9 +480,11 @@ export function createAgentMachine({ ...context.notifications.map((entry) => messageAppended({ message: entry })), ...(head === undefined ? [] - : [messageAppended({ message: createUserEntry(head.message, { source: 'input' }) })]), - queueDrained({ id: head?.id }), - notificationsDrained({}), + : [ + messageAppended({ message: createUserEntry(head.message, { source: 'input' }) }), + queueDrained({ id: head.id }), + ]), + ...(context.notifications.length === 0 ? [] : [notificationsDrained({})]), ], }; }), @@ -524,6 +565,8 @@ export function createAgentMachine({ })), ], exit: [ + 'abortTurnTools', + 'stopTurnTools', assign({ turnTools: {} }), assign({ turnId: ({ context }) => context.turnId + 1 }), ], @@ -532,27 +575,23 @@ export function createAgentMachine({ 'store.reset': { target: '#agent.idle', actions: [ - 'abortTurn', - 'abortTurnTools', - 'stopTurnTools', - 'abortBackgroundScopes', + 'abortScope', 'resetMirror', 'emitReset', 'forwardToParent', ], }, 'turn.drain': { - actions: [ - sendTo('turn', ({ context }) => ({ - type: 'turn.notify' as const, - messages: [...context.notifications, ...context.reminders], - })), - assign({ notifications: [], reminders: [] }), - sendTo('store', { + actions: enqueueActions(({ context, enqueue }) => { + const messages = [...context.notifications, ...context.reminders]; + enqueue.sendTo('turn', { type: 'turn.notify' as const, messages }); + if (messages.length === 0) return; + enqueue.assign({ notifications: [], reminders: [] }); + enqueue.sendTo('store', { type: 'store.append' as const, event: inputDrained({}), - }), - ], + }); + }), }, 'tool.detached': { guard: ({ context, event }) => context.turnTools[event.toolCallId] !== undefined, diff --git a/packages/agent-core-v2/src/human/agent/slices.ts b/packages/agent-core-v2/src/human/agent/slices.ts index 82b211f2e3b..7e10d1910a4 100644 --- a/packages/agent-core-v2/src/human/agent/slices.ts +++ b/packages/agent-core-v2/src/human/agent/slices.ts @@ -1,4 +1,3 @@ -import type { StoreActorLogic } from '#/eventStore/actor'; import type { CombinedState, EventStore } from '#/eventStore/eventStore'; import { createSlice } from '#/eventStore/slice'; import type { BranchRef } from '#/store/types'; @@ -128,4 +127,3 @@ export const agentSlices = { export type AgentSlices = typeof agentSlices; export type AgentEventStore = EventStore; export type AgentStoreState = CombinedState; -export type AgentStoreActorLogic = StoreActorLogic; diff --git a/packages/agent-core-v2/src/human/eventStore/actor.ts b/packages/agent-core-v2/src/human/eventStore/actor.ts index 79b28d8cfe1..7aeaf2e93c9 100644 --- a/packages/agent-core-v2/src/human/eventStore/actor.ts +++ b/packages/agent-core-v2/src/human/eventStore/actor.ts @@ -1,11 +1,12 @@ import { fromCallback } from '#/xstate2'; import type { CombinedState, EventStore, SliceMap } from './eventStore'; +import type { ExternalEvent } from './events'; +import type { StoreJournal } from './journal'; -export type StoreActorEvent = { - type: 'store.append'; - event: { type: string } & Record | readonly ({ type: string } & Record)[]; -}; +export type StoreActorEvent = + | { type: 'store.append'; event: ExternalEvent | readonly ExternalEvent[] } + | { type: 'store.switch'; journal: StoreJournal }; export type StoreActorEmitted = | { type: 'store.ready'; state: CombinedState; branch: string } @@ -13,29 +14,36 @@ export type StoreActorEmitted = | { type: 'store.reset'; state: CombinedState; branch: string } | { type: 'store.error'; error: unknown }; -export function fromEventStore(store: EventStore) { - return fromCallback>(({ emit, sendBack, receive }) => { - const publish = (event: StoreActorEmitted): void => { - sendBack(event); - emit(event); - }; - publish({ type: 'store.ready', state: store.getState(), branch: store.ref.branch }); - const unsubscribe = store.subscribe((state, cause) => { - if (cause.kind === 'reset') { - publish({ type: 'store.reset', state, branch: store.ref.branch }); - } else { - publish({ type: 'store.changed', state }); - } - }); - receive((event) => { - if (event.type === 'store.append') { - void store.dispatch(event.event).catch((error: unknown) => { - publish({ type: 'store.error', error }); - }); - } - }); - return unsubscribe; +export const storeActor = fromCallback< + StoreActorEvent, + { store: EventStore }, + StoreActorEmitted +>(({ input, emit, sendBack, receive }) => { + const store = input.store; + const publish = (event: StoreActorEmitted): void => { + sendBack(event); + emit(event); + }; + publish({ type: 'store.ready', state: store.getState(), branch: store.ref.branch }); + const unsubscribe = store.subscribe((state, cause) => { + if (cause.kind === 'reset') { + publish({ type: 'store.reset', state, branch: store.ref.branch }); + } else { + publish({ type: 'store.changed', state }); + } }); -} + receive((event) => { + if (event.type === 'store.append') { + void store.dispatch(event.event).catch((error: unknown) => { + publish({ type: 'store.error', error }); + }); + } else if (event.type === 'store.switch') { + void store.reset(event.journal).catch((error: unknown) => { + publish({ type: 'store.error', error }); + }); + } + }); + return unsubscribe; +}); -export type StoreActorLogic = ReturnType>; +export type StoreActorLogic = typeof storeActor; diff --git a/packages/agent-core-v2/src/human/session/stores.ts b/packages/agent-core-v2/src/human/session/stores.ts index 22d3e0441bb..aa3e7da82ac 100644 --- a/packages/agent-core-v2/src/human/session/stores.ts +++ b/packages/agent-core-v2/src/human/session/stores.ts @@ -1,7 +1,6 @@ -import { fromEventStore } from '#/eventStore/actor'; import { createEventStore, type EventStore } from '#/eventStore/eventStore'; import { journalFromBranch } from '#/eventStore/journal'; -import { agentSlices, type AgentEventStore, type AgentStoreActorLogic } from '#/agent/slices'; +import { agentSlices, type AgentEventStore } from '#/agent/slices'; import type { StoreBackend } from '#/store/backend/backend'; import { StoreError, type BranchRef } from '#/store/types'; import type { Tree } from '#/store/tree'; @@ -11,11 +10,6 @@ import { sessionSlices } from './slices'; export type SessionStore = EventStore; -export interface OpenedAgentStore { - store: AgentStoreActorLogic; - engine: AgentEventStore; -} - export type UndoErrorReason = 'unknown-agent' | 'invalid-count' | 'insufficient'; export class UndoError extends Error { @@ -74,10 +68,10 @@ export class SessionStores { return this.sessionStore; } - async open(agentId: string, opts?: { from?: BranchRef }): Promise { + async open(agentId: string, opts?: { from?: BranchRef }): Promise { const existing = this.agents.get(agentId); if (existing !== undefined) { - return { store: fromEventStore(existing), engine: existing }; + return existing; } const existed = this.tree.has(agentId); const branch = existed @@ -91,10 +85,10 @@ export class SessionStores { if (!existed) { await (await this.session()).dispatch(agentOpened({ agentId, branch: branch.name })); } - return { store: fromEventStore(engine), engine }; + return engine; } - async fork(sourceId: string, agentId: string): Promise { + async fork(sourceId: string, agentId: string): Promise { const source = this.agents.get(sourceId); if (source === undefined) { throw new StoreError('unknown-agent', `unknown agent '${sourceId}'`); diff --git a/packages/agent-core-v2/src/human/todo/slice.ts b/packages/agent-core-v2/src/human/todo/slice.ts index b7113e77c40..57a2a718b9d 100644 --- a/packages/agent-core-v2/src/human/todo/slice.ts +++ b/packages/agent-core-v2/src/human/todo/slice.ts @@ -1,6 +1,6 @@ import { createSlice } from '#/eventStore/slice'; -import type { StateUpdated, TurnStarted } from '#/agent/events'; +import { stateUpdated, turnStarted, type StateUpdated, type TurnStarted } from '#/agent/events'; import { readTodoItems, type TodoItem } from './todoItem'; @@ -19,7 +19,7 @@ export const todoSlice = createSlice({ name: 'todo', initialState: (): TodoSliceState => ({ todos: [], currentTurn: 0, lastWriteTurn: 0 }), reducers: { - 'state.updated': (draft, event: StateUpdated) => { + [stateUpdated.type]: (draft, event: StateUpdated) => { if (event.name !== 'todo') return; const value = event.value as { todos?: unknown; lastWriteTurn?: unknown }; draft.todos = readTodoItems(value.todos); @@ -27,7 +27,7 @@ export const todoSlice = createSlice({ draft.lastWriteTurn = value.lastWriteTurn; } }, - 'turn.started': (draft, _event: TurnStarted) => { + [turnStarted.type]: (draft, _event: TurnStarted) => { draft.currentTurn += 1; }, }, From 2fc619b5f4f259f317faef5014da0f4fe3219868 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Wed, 9 Sep 2026 23:50:02 +0800 Subject: [PATCH 3/3] refactor(agent-core-v2): drop journal snapshots for plain full-replay event sourcing Snapshots serialized the whole slice state every 500 events and, on an asynchronous backend, their fire-and-forget append could still be queued when the next dispatch read journal.nextSeq(), making the recorded FoldContext.ref.seq one less than the persisted event's sequence and shifting turnIndex undo boundaries. Removing the mechanism eliminates the race and the full-state rewrite; refold now always seeds from initial state and replays every event record (legacy snapshot entries are ignored). --- .../src/human/eventStore/eventStore.ts | 66 ++++++------------- .../src/human/eventStore/slice.ts | 2 - .../human/test/eventStore/eventStore.test.ts | 20 ++---- 3 files changed, 23 insertions(+), 65 deletions(-) diff --git a/packages/agent-core-v2/src/human/eventStore/eventStore.ts b/packages/agent-core-v2/src/human/eventStore/eventStore.ts index af87db77dc3..62e402dc73e 100644 --- a/packages/agent-core-v2/src/human/eventStore/eventStore.ts +++ b/packages/agent-core-v2/src/human/eventStore/eventStore.ts @@ -5,7 +5,7 @@ import { StoreError } from '#/store/types'; import type { ExternalEvent, InternalEvent } from './events'; import { eventSchemaFor, parseEvent, validateEvent } from './events'; -import type { JournalRecord, StoreJournal } from './journal'; +import type { JournalRecord, StoreJournal, SyncStoreJournal } from './journal'; import type { FoldContext, Slice } from './slice'; export type SliceMap = Record>; @@ -23,7 +23,6 @@ export type Cause = export interface EventStoreOptions { journal: StoreJournal; slices: SM; - snapshot?: { everyEvents?: number } | false; drainLimit?: number; onError?: (error: unknown) => void; } @@ -44,10 +43,7 @@ export interface EventStore { close(): Promise; } -const SNAPSHOT_ENTRY_TYPE = 'snapshot'; -const SNAPSHOT_ENTRY_KIND = 'snapshot'; const EVENT_ENTRY_KIND = 'event'; -const DEFAULT_SNAPSHOT_EVERY_EVENTS = 500; const DEFAULT_DRAIN_LIMIT = 100; type Listener = (state: CombinedState, cause: Cause) => void; @@ -60,14 +56,20 @@ export async function createEventStore( return store; } +export function createEventStoreSync( + opts: EventStoreOptions & { journal: SyncStoreJournal }, +): EventStore { + const store = new EventStoreImpl(opts); + store.refoldSync(opts.journal); + return store; +} + class EventStoreImpl implements EventStore { private journal: StoreJournal; private slices: SliceMap; private state: Record; private phaseValue: 'open' | 'closed' = 'open'; private tail: Promise = Promise.resolve(); - private eventsSinceSnapshot = 0; - private readonly snapshotEvery: number | false; private readonly drainLimit: number; private readonly report: (error: unknown) => void; private readonly listeners = new Set>(); @@ -76,8 +78,6 @@ class EventStoreImpl implements EventStore { this.journal = opts.journal; this.slices = { ...opts.slices }; this.state = {}; - this.snapshotEvery = - opts.snapshot === false ? false : (opts.snapshot?.everyEvents ?? DEFAULT_SNAPSHOT_EVERY_EVENTS); this.drainLimit = opts.drainLimit ?? DEFAULT_DRAIN_LIMIT; this.report = opts.onError ?? ((error) => console.error(error)); } @@ -179,36 +179,23 @@ class EventStoreImpl implements EventStore { for await (const record of journal.read()) { records.push(record); } - let snapshotIndex = -1; - for (let i = records.length - 1; i >= 0; i--) { - if (records[i]?.kind === SNAPSHOT_ENTRY_KIND) { - snapshotIndex = i; - break; - } - } - const saved = - snapshotIndex >= 0 - ? (((records[snapshotIndex] as JournalRecord).data as { slices?: Record }) - .slices ?? {}) - : {}; + this.foldRecords(records); + } + + refoldSync(journal: SyncStoreJournal): void { + this.foldRecords(journal.readSync()); + } + + private foldRecords(records: JournalRecord[]): void { const seeded: Record = {}; for (const [name, slice] of Object.entries(this.slices)) { - seeded[name] = - name in saved - ? slice.deserialize !== undefined - ? slice.deserialize(saved[name]) - : saved[name] - : slice.initialState(); + seeded[name] = slice.initialState(); } this.state = seeded; - let sinceSnapshot = 0; - for (let i = snapshotIndex + 1; i < records.length; i++) { - const record = records[i] as JournalRecord; + for (const record of records) { if (record.kind !== EVENT_ENTRY_KIND) continue; - sinceSnapshot += 1; this.replayRecord(record); } - this.eventsSinceSnapshot = sinceSnapshot; } private replayRecord(record: JournalRecord): void { @@ -243,7 +230,6 @@ class EventStoreImpl implements EventStore { const entry = await this.journal.append({ type: event.type, kind: EVENT_ENTRY_KIND, data: event }); entries.push(entry); causes.push({ kind: 'event', event, entry }, ...internalCauses); - this.maybeSnapshot(); } this.notify(causes); return entries[entries.length - 1] as EntryLine; @@ -315,20 +301,6 @@ class EventStoreImpl implements EventStore { return causes; } - private maybeSnapshot(): void { - if (this.snapshotEvery === false) return; - this.eventsSinceSnapshot += 1; - if (this.eventsSinceSnapshot < this.snapshotEvery) return; - this.eventsSinceSnapshot = 0; - const slices: Record = {}; - for (const [name, slice] of Object.entries(this.slices)) { - slices[name] = slice.serialize !== undefined ? slice.serialize(this.state[name]) : this.state[name]; - } - void this.journal - .append({ type: SNAPSHOT_ENTRY_TYPE, kind: SNAPSHOT_ENTRY_KIND, data: { slices } }) - .catch((error) => this.report(error)); - } - private notify(causes: Cause[]): void { const state = this.getState(); for (const cause of causes) { diff --git a/packages/agent-core-v2/src/human/eventStore/slice.ts b/packages/agent-core-v2/src/human/eventStore/slice.ts index 6a80a01e1e4..3479db4d71a 100644 --- a/packages/agent-core-v2/src/human/eventStore/slice.ts +++ b/packages/agent-core-v2/src/human/eventStore/slice.ts @@ -18,8 +18,6 @@ export interface Slice { readonly name: Name; readonly initialState: () => S; readonly reducers: Record void | S>; - readonly serialize?: (state: S) => unknown; - readonly deserialize?: (raw: unknown) => S; } export function createSlice(def: Slice): Slice { diff --git a/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts b/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts index bfd3cb24b14..cb7b505ad84 100644 --- a/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts +++ b/packages/agent-core-v2/src/human/test/eventStore/eventStore.test.ts @@ -48,9 +48,9 @@ async function openJournal(tree: Tree, branch = 'main') { return journalFromBranch(tree.openBranch(branch), tree); } -async function openStore(tree: Tree, opts?: { snapshot?: { everyEvents?: number } | false; drainLimit?: number; extraSlices?: Record }) { +async function openStore(tree: Tree, opts?: { drainLimit?: number; extraSlices?: Record }) { const journal = await openJournal(tree); - return createEventStore({ journal, slices, snapshot: opts?.snapshot, drainLimit: opts?.drainLimit }); + return createEventStore({ journal, slices, drainLimit: opts?.drainLimit }); } describe('createEventStore', () => { @@ -67,19 +67,7 @@ describe('createEventStore', () => { expect(reopened.getState()).toEqual({ counter: 3, notes: ['a'] }); }); - it('folds from the latest snapshot plus its tail', async () => { - const tree = await openTree(); - const store = await openStore(tree, { snapshot: { everyEvents: 2 } }); - await store.dispatch(counterAdded({ amount: 1 })); - await store.dispatch(counterAdded({ amount: 2 })); - await store.dispatch(counterAdded({ amount: 3 })); - await store.flush(); - - const reopened = await openStore(tree); - expect(reopened.getState()).toEqual({ counter: 6, notes: [] }); - }); - - it('seeds slices missing from a snapshot from their initial state', async () => { + it('ignores legacy snapshot entries when folding', async () => { const tree = await openTree(); const journal = await openJournal(tree); await journal.append({ type: 'snapshot', kind: 'snapshot', data: { slices: { counter: 41 } } }); @@ -89,7 +77,7 @@ describe('createEventStore', () => { data: { type: 'test.counter_added', time: 1, amount: 1 }, }); const store = await createEventStore({ journal, slices }); - expect(store.getState()).toEqual({ counter: 42, notes: [] }); + expect(store.getState()).toEqual({ counter: 1, notes: [] }); }); it('skips unknown event types when folding', async () => {