diff --git a/src/main/AryxAppService.ts b/src/main/AryxAppService.ts index 6dc3e0a..6530616 100644 --- a/src/main/AryxAppService.ts +++ b/src/main/AryxAppService.ts @@ -62,6 +62,7 @@ import { } from '@shared/domain/approval'; import { isScratchpadProject, type ProjectRecord } from '@shared/domain/project'; import { + branchSessionRecord, duplicateSessionRecord, querySessions as queryWorkspaceSessions, renameSessionRecord, @@ -703,6 +704,23 @@ export class AryxAppService extends EventEmitter { return this.persistAndBroadcast(workspace); } + async branchSession(sessionId: string, messageId: string): Promise { + const workspace = await this.loadWorkspace(); + const session = this.requireSession(workspace, sessionId); + const pattern = this.requirePattern(workspace, session.patternId); + const branch = branchSessionRecord(session, pattern, createId('session'), messageId, nowIso()); + if (isScratchpadProject(branch.projectId)) { + branch.cwd = undefined; + } + + await this.ensureScratchpadSessionDirectory(branch); + workspace.sessions.unshift(branch); + workspace.selectedProjectId = branch.projectId; + workspace.selectedPatternId = branch.patternId; + workspace.selectedSessionId = branch.id; + return this.persistAndBroadcast(workspace); + } + async renameSession(sessionId: string, title: string): Promise { const workspace = await this.loadWorkspace(); const session = this.requireSession(workspace, sessionId); diff --git a/src/main/ipc/registerIpcHandlers.ts b/src/main/ipc/registerIpcHandlers.ts index 4f16972..e565183 100644 --- a/src/main/ipc/registerIpcHandlers.ts +++ b/src/main/ipc/registerIpcHandlers.ts @@ -3,6 +3,7 @@ import type { BrowserWindow } from 'electron'; import { ipcChannels } from '@shared/contracts/channels'; import type { + BranchSessionInput, CancelSessionTurnInput, CreateSessionInput, DismissSessionMcpAuthInput, @@ -135,6 +136,9 @@ export function registerIpcHandlers(window: BrowserWindow, service: AryxAppServi ipcMain.handle(ipcChannels.duplicateSession, (_event, input: DuplicateSessionInput) => service.duplicateSession(input.sessionId), ); + ipcMain.handle(ipcChannels.branchSession, (_event, input: BranchSessionInput) => + service.branchSession(input.sessionId, input.messageId), + ); ipcMain.handle(ipcChannels.renameSession, (_event, input: RenameSessionInput) => service.renameSession(input.sessionId, input.title), ); diff --git a/src/main/persistence/workspaceRepository.ts b/src/main/persistence/workspaceRepository.ts index c20c627..c8d5573 100644 --- a/src/main/persistence/workspaceRepository.ts +++ b/src/main/persistence/workspaceRepository.ts @@ -6,7 +6,7 @@ import { isScratchpadProject, mergeScratchpadProject } from '@shared/domain/proj import { normalizeDiscoveredToolingState } from '@shared/domain/discoveredTooling'; import { normalizeProjectCustomizationState } from '@shared/domain/projectCustomization'; import { normalizeSessionRunRecords } from '@shared/domain/runTimeline'; -import type { SessionRecord } from '@shared/domain/session'; +import { normalizeSessionBranchOrigin, type SessionRecord } from '@shared/domain/session'; import { normalizeSessionToolingSelection, normalizeWorkspaceSettings, @@ -81,6 +81,7 @@ export class WorkspaceRepository { const sessions = await Promise.all((stored.sessions ?? []).map(async (session): Promise => { const normalizedSession: SessionRecord = { ...session, + branchOrigin: normalizeSessionBranchOrigin(session.branchOrigin), runs: normalizeSessionRunRecords(session.runs), tooling: normalizeSessionToolingSelection(session.tooling), approvalSettings: normalizeSessionApprovalSettings(session.approvalSettings), diff --git a/src/preload/index.ts b/src/preload/index.ts index 3d804f4..097dc46 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -47,6 +47,7 @@ const api: ElectronApi = { ipcRenderer.invoke(ipcChannels.updateSessionApprovalSettings, input), createSession: (input) => ipcRenderer.invoke(ipcChannels.createSession, input), duplicateSession: (input) => ipcRenderer.invoke(ipcChannels.duplicateSession, input), + branchSession: (input) => ipcRenderer.invoke(ipcChannels.branchSession, input), renameSession: (input) => ipcRenderer.invoke(ipcChannels.renameSession, input), setSessionPinned: (input) => ipcRenderer.invoke(ipcChannels.setSessionPinned, input), setSessionArchived: (input) => ipcRenderer.invoke(ipcChannels.setSessionArchived, input), diff --git a/src/shared/contracts/channels.ts b/src/shared/contracts/channels.ts index e45e199..a720868 100644 --- a/src/shared/contracts/channels.ts +++ b/src/shared/contracts/channels.ts @@ -31,6 +31,7 @@ export const ipcChannels = { updateSessionApprovalSettings: 'sessions:update-approval-settings', createSession: 'sessions:create', duplicateSession: 'sessions:duplicate', + branchSession: 'sessions:branch', renameSession: 'sessions:rename', setSessionPinned: 'sessions:set-pinned', setSessionArchived: 'sessions:set-archived', diff --git a/src/shared/contracts/ipc.ts b/src/shared/contracts/ipc.ts index 1669c5c..5eded8b 100644 --- a/src/shared/contracts/ipc.ts +++ b/src/shared/contracts/ipc.ts @@ -58,6 +58,11 @@ export interface DuplicateSessionInput { sessionId: string; } +export interface BranchSessionInput { + sessionId: string; + messageId: string; +} + export interface RenameSessionInput { sessionId: string; title: string; @@ -174,6 +179,7 @@ export interface ElectronApi { updateSessionApprovalSettings(input: UpdateSessionApprovalSettingsInput): Promise; createSession(input: CreateSessionInput): Promise; duplicateSession(input: DuplicateSessionInput): Promise; + branchSession(input: BranchSessionInput): Promise; renameSession(input: RenameSessionInput): Promise; setSessionPinned(input: SetSessionPinnedInput): Promise; setSessionArchived(input: SetSessionArchivedInput): Promise; diff --git a/src/shared/domain/session.ts b/src/shared/domain/session.ts index 95a4f9c..243fae5 100644 --- a/src/shared/domain/session.ts +++ b/src/shared/domain/session.ts @@ -36,6 +36,13 @@ export interface ChatMessageRecord { attachments?: ChatMessageAttachment[]; } +export interface SessionBranchOrigin { + sourceSessionId: string; + sourceMessageId: string; + sourceMessageIndex: number; + branchedAt: string; +} + export interface SessionRecord { id: string; projectId: string; @@ -47,6 +54,7 @@ export interface SessionRecord { status: SessionStatus; isPinned?: boolean; isArchived?: boolean; + branchOrigin?: SessionBranchOrigin; interactionMode?: InteractionMode; cwd?: string; messages: ChatMessageRecord[]; @@ -62,6 +70,38 @@ export interface SessionRecord { runs: SessionRunRecord[]; } +function normalizeOptionalString(value?: string): string | undefined { + const trimmed = value?.trim(); + return trimmed ? trimmed : undefined; +} + +export function normalizeSessionBranchOrigin( + branchOrigin?: Partial, +): SessionBranchOrigin | undefined { + const sourceSessionId = normalizeOptionalString(branchOrigin?.sourceSessionId); + const sourceMessageId = normalizeOptionalString(branchOrigin?.sourceMessageId); + const branchedAt = normalizeOptionalString(branchOrigin?.branchedAt); + const sourceMessageIndex = branchOrigin?.sourceMessageIndex; + + if ( + !sourceSessionId + || !sourceMessageId + || !branchedAt + || typeof sourceMessageIndex !== 'number' + || !Number.isInteger(sourceMessageIndex) + || sourceMessageIndex < 0 + ) { + return undefined; + } + + return { + sourceSessionId, + sourceMessageId, + sourceMessageIndex, + branchedAt, + }; +} + export function resolveSessionTitle( session: Pick, pattern: PatternDefinition, diff --git a/src/shared/domain/sessionLibrary.ts b/src/shared/domain/sessionLibrary.ts index a65df12..d2e29d0 100644 --- a/src/shared/domain/sessionLibrary.ts +++ b/src/shared/domain/sessionLibrary.ts @@ -1,6 +1,6 @@ import type { PatternDefinition } from '@shared/domain/pattern'; import { isScratchpadProject, type ProjectRecord } from '@shared/domain/project'; -import type { ChatMessageRecord, SessionRecord, SessionStatus } from '@shared/domain/session'; +import { resolveSessionTitle, type ChatMessageRecord, type SessionBranchOrigin, type SessionRecord, type SessionStatus } from '@shared/domain/session'; import type { WorkspaceState } from '@shared/domain/workspace'; export type SessionQueryMatchField = 'title' | 'message' | 'project' | 'pattern'; @@ -158,21 +158,32 @@ export function renameSessionRecord(session: SessionRecord, title: string, updat }; } -export function duplicateSessionRecord( +function cloneBranchOrigin(branchOrigin?: SessionBranchOrigin): SessionBranchOrigin | undefined { + return branchOrigin ? { ...branchOrigin } : undefined; +} + +function cloneChatMessageRecord(message: ChatMessageRecord): ChatMessageRecord { + return { + ...message, + pending: false, + attachments: message.attachments?.map((attachment) => ({ ...attachment })), + }; +} + +function createDerivedSessionRecord( session: SessionRecord, sessionId: string, - duplicatedAt: string, + createdAt: string, ): SessionRecord { return { ...session, id: sessionId, - title: `${session.title} (Copy)`, - titleSource: 'manual', - createdAt: duplicatedAt, - updatedAt: duplicatedAt, + createdAt, + updatedAt: createdAt, status: 'idle', isPinned: false, isArchived: false, + branchOrigin: cloneBranchOrigin(session.branchOrigin), lastError: undefined, sessionModelConfig: session.sessionModelConfig ? { ...session.sessionModelConfig } : undefined, tooling: session.tooling @@ -188,11 +199,60 @@ export function duplicateSessionRecord( : undefined, pendingApproval: undefined, pendingApprovalQueue: undefined, + pendingUserInput: undefined, + pendingPlanReview: undefined, + pendingMcpAuth: undefined, runs: [], - messages: session.messages.map((message): ChatMessageRecord => ({ - ...message, - pending: false, - })), + messages: [], + }; +} + +export function duplicateSessionRecord( + session: SessionRecord, + sessionId: string, + duplicatedAt: string, +): SessionRecord { + return { + ...createDerivedSessionRecord(session, sessionId, duplicatedAt), + title: `${session.title} (Copy)`, + titleSource: 'manual', + messages: session.messages.map(cloneChatMessageRecord), + }; +} + +export function branchSessionRecord( + session: SessionRecord, + pattern: PatternDefinition, + sessionId: string, + messageId: string, + branchedAt: string, +): SessionRecord { + const sourceMessageIndex = session.messages.findIndex((message) => message.id === messageId); + if (sourceMessageIndex < 0) { + throw new Error(`Message ${messageId} not found in session ${session.id}.`); + } + + const sourceMessage = session.messages[sourceMessageIndex]; + if (!sourceMessage) { + throw new Error(`Message ${messageId} not found in session ${session.id}.`); + } + + if (sourceMessage.role !== 'user') { + throw new Error('Only user messages can be used as a branch point.'); + } + + const branchedMessages = session.messages.slice(0, sourceMessageIndex + 1).map(cloneChatMessageRecord); + + return { + ...createDerivedSessionRecord(session, sessionId, branchedAt), + title: resolveSessionTitle(session, pattern, branchedMessages), + messages: branchedMessages, + branchOrigin: { + sourceSessionId: session.id, + sourceMessageId: messageId, + sourceMessageIndex, + branchedAt, + }, }; } diff --git a/tests/main/appServiceScratchpadDirectories.test.ts b/tests/main/appServiceScratchpadDirectories.test.ts index b3761ed..e41abaa 100644 --- a/tests/main/appServiceScratchpadDirectories.test.ts +++ b/tests/main/appServiceScratchpadDirectories.test.ts @@ -169,6 +169,73 @@ describe('AryxAppService scratchpad directories', () => { expect(await pathExists(duplicate.cwd!)).toBe(true); }); + test('creates a new directory when branching a scratchpad session', async () => { + const workspace = createWorkspaceSeed(); + const pattern = requireSinglePattern(workspace); + const scratchpadProject = createScratchpadProject(join(USER_DATA_PATH, 'scratchpad'), TIMESTAMP); + const originalDirectory = join(USER_DATA_PATH, 'scratchpad', 'session-original'); + const originalSession = createScratchpadSession(pattern.id, { + id: 'session-original', + cwd: originalDirectory, + messages: [ + { + id: 'msg-1', + role: 'user', + authorName: 'You', + content: 'Draft a plan.', + createdAt: TIMESTAMP, + }, + { + id: 'msg-2', + role: 'assistant', + authorName: 'Primary Agent', + content: 'Here is the first approach.', + createdAt: TIMESTAMP, + }, + { + id: 'msg-3', + role: 'user', + authorName: 'You', + content: 'Try again from here with a stricter checklist.', + createdAt: TIMESTAMP, + }, + { + id: 'msg-4', + role: 'assistant', + authorName: 'Primary Agent', + content: 'Here is the revised approach.', + createdAt: TIMESTAMP, + }, + ], + }); + + workspace.projects = [scratchpadProject]; + workspace.sessions = [originalSession]; + workspace.selectedProjectId = scratchpadProject.id; + workspace.selectedPatternId = pattern.id; + workspace.selectedSessionId = originalSession.id; + + const service = createService(workspace); + + const result = await service.branchSession(originalSession.id, 'msg-3'); + const branch = result.sessions[0]; + if (!branch) { + throw new Error('Expected branchSession to prepend the branch.'); + } + + expect(branch.id).not.toBe(originalSession.id); + expect(branch.cwd).toBe(join(USER_DATA_PATH, 'scratchpad', branch.id)); + expect(branch.cwd).not.toBe(originalSession.cwd); + expect(branch.branchOrigin).toMatchObject({ + sourceSessionId: originalSession.id, + sourceMessageId: 'msg-3', + sourceMessageIndex: 2, + }); + expect(branch.messages.map((message) => message.id)).toEqual(['msg-1', 'msg-2', 'msg-3']); + expect(result.selectedSessionId).toBe(branch.id); + expect(await pathExists(branch.cwd!)).toBe(true); + }); + test('removes the scratchpad directory when deleting a session', async () => { const workspace = createWorkspaceSeed(); const pattern = requireSinglePattern(workspace); diff --git a/tests/shared/sessionLibrary.test.ts b/tests/shared/sessionLibrary.test.ts index 8022a98..aead541 100644 --- a/tests/shared/sessionLibrary.test.ts +++ b/tests/shared/sessionLibrary.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'bun:test'; -import { querySessions, duplicateSessionRecord, renameSessionRecord } from '@shared/domain/sessionLibrary'; +import { branchSessionRecord, querySessions, duplicateSessionRecord, renameSessionRecord } from '@shared/domain/sessionLibrary'; import type { PatternDefinition } from '@shared/domain/pattern'; import type { ProjectRecord } from '@shared/domain/project'; import type { SessionRecord } from '@shared/domain/session'; @@ -176,6 +176,132 @@ describe('session library helpers', () => { expect(session.runs).toEqual([]); }); + test('branches sessions from a user message and retains only the prior transcript', () => { + const sourceSession = createSession({ + title: 'Manual branch source', + titleSource: 'manual', + status: 'error', + isPinned: true, + isArchived: true, + lastError: 'sidecar crashed', + approvalSettings: { + autoApprovedToolNames: ['git.status'], + }, + pendingApproval: { + id: 'approval-1', + kind: 'tool-call', + status: 'pending', + requestedAt: '2026-03-23T00:01:00.000Z', + title: 'Approve tool access', + }, + pendingApprovalQueue: [ + { + id: 'approval-2', + kind: 'final-response', + status: 'pending', + requestedAt: '2026-03-23T00:02:00.000Z', + title: 'Approve final response', + }, + ], + messages: [ + { + id: 'msg-1', + role: 'user', + authorName: 'You', + content: 'Investigate the refresh bug.', + createdAt: '2026-03-23T00:00:00.000Z', + }, + { + id: 'msg-2', + role: 'assistant', + authorName: 'Reviewer', + content: 'I found two likely causes.', + createdAt: '2026-03-23T00:01:00.000Z', + pending: true, + }, + { + id: 'msg-3', + role: 'user', + authorName: 'You', + content: 'Try a different approach focused on session state.', + createdAt: '2026-03-23T00:02:00.000Z', + attachments: [ + { + type: 'file', + path: 'C:\\workspace\\alpha\\notes.txt', + displayName: 'notes.txt', + }, + ], + }, + { + id: 'msg-4', + role: 'assistant', + authorName: 'Reviewer', + content: 'Here is the alternate plan.', + createdAt: '2026-03-23T00:03:00.000Z', + }, + ], + }); + + const branch = branchSessionRecord( + sourceSession, + createPattern(), + 'session-branch', + 'msg-3', + '2026-03-23T00:04:00.000Z', + ); + + expect(branch).toMatchObject({ + id: 'session-branch', + title: 'Manual branch source', + titleSource: 'manual', + status: 'idle', + isPinned: false, + isArchived: false, + lastError: undefined, + createdAt: '2026-03-23T00:04:00.000Z', + updatedAt: '2026-03-23T00:04:00.000Z', + branchOrigin: { + sourceSessionId: 'session-1', + sourceMessageId: 'msg-3', + sourceMessageIndex: 2, + branchedAt: '2026-03-23T00:04:00.000Z', + }, + }); + expect(branch.messages.map((message) => message.id)).toEqual(['msg-1', 'msg-2', 'msg-3']); + expect(branch.messages[1]?.pending).toBe(false); + expect(branch.messages[2]?.attachments).toEqual(sourceSession.messages[2]?.attachments); + expect(branch.messages[2]?.attachments).not.toBe(sourceSession.messages[2]?.attachments); + expect(branch.messages[1]).not.toBe(sourceSession.messages[1]); + expect(branch.pendingApproval).toBeUndefined(); + expect(branch.pendingApprovalQueue).toBeUndefined(); + expect(branch.runs).toEqual([]); + expect(sourceSession.messages[1]?.pending).toBe(true); + }); + + test('rejects branching from non-user messages', () => { + const sourceSession = createSession({ + messages: [ + { + id: 'msg-1', + role: 'assistant', + authorName: 'Reviewer', + content: 'I can help with that.', + createdAt: '2026-03-23T00:00:00.000Z', + }, + ], + }); + + expect(() => + branchSessionRecord( + sourceSession, + createPattern(), + 'session-branch', + 'msg-1', + '2026-03-23T00:04:00.000Z', + )).toThrow('Only user messages can be used as a branch point.'); + }); + test('searches across session title, messages, projects, and patterns', () => { const workspace = createWorkspace();