diff --git a/packages/agent-core-v2/src/human/session/stores.ts b/packages/agent-core-v2/src/human/session/stores.ts index d998d3c5db..c17071288e 100644 --- a/packages/agent-core-v2/src/human/session/stores.ts +++ b/packages/agent-core-v2/src/human/session/stores.ts @@ -3,6 +3,7 @@ import type { ExternalEvent } from '#/eventStore/events'; import { journalFromBranch } from '#/eventStore/journal'; import { agentSlices, type AgentEventStore } from '#/agent/slices'; import type { StoreBackend } from '#/store/backend/backend'; +import type { Branch } from '#/store/branch'; import { StoreError, type BranchRef } from '#/store/types'; import type { Tree } from '#/store/tree'; @@ -74,17 +75,26 @@ export class SessionStores { if (existing !== undefined) { return 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 session = await this.session(); + const registered = session.getState().roster.agents[agentId]; + let branch: Branch; + if (registered !== undefined) { + branch = this.tree.openBranch(registered); + } else if (this.tree.has(agentId)) { + branch = this.tree.openBranch(agentId); + } else { + branch = 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 })); + if (registered === undefined) { + await session.dispatch(agentOpened({ agentId, branch: branch.name })); } return engine; } 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 a0e3d4c74f..36cd67f02b 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 @@ -604,6 +604,13 @@ describe('migrateV2Session', () => { const roster = (await second.stores.session()).getState().roster.agents; expect(Object.keys(roster)).toEqual([MAIN]); expect(roster[MAIN]).toBe('main~2'); + const reopenedAgent = await second.stores.open(MAIN); + expect(reopenedAgent.ref.branch).toBe('main~2'); + expect(reopenedAgent.getState().history.map((entry) => extractText(entry.message))).toEqual([ + 'first', + 'first-reply', + 'again', + ]); 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/stores.test.ts b/packages/agent-core-v2/src/human/test/session/stores.test.ts index e076ca6787..29342f5391 100644 --- a/packages/agent-core-v2/src/human/test/session/stores.test.ts +++ b/packages/agent-core-v2/src/human/test/session/stores.test.ts @@ -185,16 +185,18 @@ describe('SessionStores reopen', () => { await runTurn(forkActor, fork, 'fork-hi', 6); forkActor.stop(); actor.stop(); + await env.stores.undo('main', 1); const restored = await reopen(env); expect((await restored.stores.session()).getState().roster.agents).toEqual({ fork: 'fork', - main: 'main', + main: 'main~2', }); const restoredMain = await restored.stores.open('main'); - expect(historyTexts(restoredMain)).toEqual(['first', 'echo:first', 'second', 'echo:second']); - expect(restoredMain.getState().turnIndex.nextTurnId).toBe(2); + expect(restoredMain.ref.branch).toBe('main~2'); + expect(historyTexts(restoredMain)).toEqual(['first', 'echo:first', 'second']); + expect(restoredMain.getState().turnIndex.nextTurnId).toBe(1); const restoredFork = await restored.stores.open('fork'); expect(historyTexts(restoredFork)).toEqual([ 'first', @@ -207,16 +209,15 @@ describe('SessionStores reopen', () => { expect(restoredFork.getState().turnIndex.nextTurnId).toBe(3); const actor2 = startAgent(restoredMain); - await runTurn(actor2, restoredMain, 'again', 6); + await runTurn(actor2, restoredMain, 'again', 5); expect(historyTexts(restoredMain)).toEqual([ 'first', 'echo:first', 'second', - 'echo:second', 'again', 'echo:again', ]); - expect(restoredMain.getState().turnIndex.nextTurnId).toBe(3); + expect(restoredMain.getState().turnIndex.nextTurnId).toBe(2); actor2.stop(); });