From 4726e2acead44ef1d19ce7638b048fcf72bab0f6 Mon Sep 17 00:00:00 2001
From: David Kaya <david@kaya.sk>
Date: Sun, 29 Mar 2026 21:43:58 +0200
Subject: [PATCH] feat: allow branching from both user and assistant messages

- Relax branchSessionRecord validation to accept user or assistant roles
- Show "Branch from here" hover button on all conversation messages
- Contextual tooltip: "starting from this message" vs "continuing from this response"
- Add test for branching from assistant message
- Replace non-user rejection test with non-conversation rejection test

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/renderer/components/ChatPane.tsx |  6 ++-
 src/shared/domain/sessionLibrary.ts  |  4 +-
 tests/shared/sessionLibrary.test.ts  | 60 +++++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/renderer/components/ChatPane.tsx b/src/renderer/components/ChatPane.tsx
index 40eeea1..c3dec7a 100644
--- a/src/renderer/components/ChatPane.tsx
+++ b/src/renderer/components/ChatPane.tsx
@@ -380,11 +380,13 @@ export function ChatPane({
                               {phaseLabel}
                             </span>
                           )}
-                          {isUser && onBranchFromMessage && (
+                          {onBranchFromMessage && (
                             <button
                               className="ml-auto flex items-center gap-1 rounded-md px-1.5 py-0.5 text-[11px] text-[var(--color-text-muted)] opacity-0 transition-all duration-150 hover:bg-[var(--color-surface-2)] hover:text-[var(--color-accent)] group-hover:opacity-100"
                               onClick={() => onBranchFromMessage(message.id)}
-                              title="Branch from here — create a new session starting from this message"
+                              title={isUser
+                                ? 'Branch from here — create a new session starting from this message'
+                                : 'Branch from here — create a new session continuing from this response'}
                               type="button"
                             >
                               <GitBranch className="size-3" />
diff --git a/src/shared/domain/sessionLibrary.ts b/src/shared/domain/sessionLibrary.ts
index d2e29d0..43e2e2f 100644
--- a/src/shared/domain/sessionLibrary.ts
+++ b/src/shared/domain/sessionLibrary.ts
@@ -237,8 +237,8 @@ export function branchSessionRecord(
     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.');
+  if (sourceMessage.role !== 'user' && sourceMessage.role !== 'assistant') {
+    throw new Error('Only user or assistant messages can be used as a branch point.');
   }
 
   const branchedMessages = session.messages.slice(0, sourceMessageIndex + 1).map(cloneChatMessageRecord);
diff --git a/tests/shared/sessionLibrary.test.ts b/tests/shared/sessionLibrary.test.ts
index aead541..1c1168a 100644
--- a/tests/shared/sessionLibrary.test.ts
+++ b/tests/shared/sessionLibrary.test.ts
@@ -279,14 +279,14 @@ describe('session library helpers', () => {
     expect(sourceSession.messages[1]?.pending).toBe(true);
   });
 
-  test('rejects branching from non-user messages', () => {
+  test('rejects branching from non-conversation messages', () => {
     const sourceSession = createSession({
       messages: [
         {
           id: 'msg-1',
-          role: 'assistant',
-          authorName: 'Reviewer',
-          content: 'I can help with that.',
+          role: 'system' as 'user',
+          authorName: 'System',
+          content: 'System prompt.',
           createdAt: '2026-03-23T00:00:00.000Z',
         },
       ],
@@ -299,7 +299,57 @@ describe('session library helpers', () => {
         'session-branch',
         'msg-1',
         '2026-03-23T00:04:00.000Z',
-      )).toThrow('Only user messages can be used as a branch point.');
+      )).toThrow('Only user or assistant messages can be used as a branch point.');
+  });
+
+  test('branches from an assistant message and retains transcript up to that response', () => {
+    const sourceSession = createSession({
+      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',
+        },
+        {
+          id: 'msg-3',
+          role: 'user',
+          authorName: 'You',
+          content: 'Try a different approach.',
+          createdAt: '2026-03-23T00:02:00.000Z',
+        },
+        {
+          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-2',
+      '2026-03-23T00:05:00.000Z',
+    );
+
+    expect(branch.messages.map((m) => m.id)).toEqual(['msg-1', 'msg-2']);
+    expect(branch.branchOrigin).toMatchObject({
+      sourceSessionId: 'session-1',
+      sourceMessageId: 'msg-2',
+      sourceMessageIndex: 1,
+    });
   });
 
   test('searches across session title, messages, projects, and patterns', () => {