Add session forking capability (#5882)

Co-authored-by: Zane Staggs <zane@squareup.com>

ui/desktop/openapi.json

@@ -2472,64 +2472,6 @@
         ]
       }
     },
-    "/sessions/{session_id}/edit_message": {
-      "post": {
-        "tags": [
-          "Session Management"
-        ],
-        "operationId": "edit_message",
-        "parameters": [
-          {
-            "name": "session_id",
-            "in": "path",
-            "description": "Unique identifier for the session",
-            "required": true,
-            "schema": {
-              "type": "string"
-            }
-          }
-        ],
-        "requestBody": {
-          "content": {
-            "application/json": {
-              "schema": {
-                "$ref": "#/components/schemas/EditMessageRequest"
-              }
-            }
-          },
-          "required": true
-        },
-        "responses": {
-          "200": {
-            "description": "Session prepared for editing - frontend should submit the edited message",
-            "content": {
-              "application/json": {
-                "schema": {
-                  "$ref": "#/components/schemas/EditMessageResponse"
-                }
-              }
-            }
-          },
-          "400": {
-            "description": "Bad request - Invalid message timestamp"
-          },
-          "401": {
-            "description": "Unauthorized - Invalid or missing API key"
-          },
-          "404": {
-            "description": "Session or message not found"
-          },
-          "500": {
-            "description": "Internal server error"
-          }
-        },
-        "security": [
-          {
-            "api_key": []
-          }
-        ]
-      }
-    },
     "/sessions/{session_id}/export": {
       "get": {
         "tags": [
@@ -2620,6 +2562,64 @@
         ]
       }
     },
+    "/sessions/{session_id}/fork": {
+      "post": {
+        "tags": [
+          "Session Management"
+        ],
+        "operationId": "fork_session",
+        "parameters": [
+          {
+            "name": "session_id",
+            "in": "path",
+            "description": "Unique identifier for the session",
+            "required": true,
+            "schema": {
+              "type": "string"
+            }
+          }
+        ],
+        "requestBody": {
+          "content": {
+            "application/json": {
+              "schema": {
+                "$ref": "#/components/schemas/ForkRequest"
+              }
+            }
+          },
+          "required": true
+        },
+        "responses": {
+          "200": {
+            "description": "Session forked successfully",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ForkResponse"
+                }
+              }
+            }
+          },
+          "400": {
+            "description": "Bad request - truncate=true requires timestamp"
+          },
+          "401": {
+            "description": "Unauthorized - Invalid or missing API key"
+          },
+          "404": {
+            "description": "Session not found"
+          },
+          "500": {
+            "description": "Internal server error"
+          }
+        },
+        "security": [
+          {
+            "api_key": []
+          }
+        ]
+      }
+    },
     "/sessions/{session_id}/name": {
       "put": {
         "tags": [
@@ -3453,39 +3453,6 @@
           }
         }
       },
-      "EditMessageRequest": {
-        "type": "object",
-        "required": [
-          "timestamp"
-        ],
-        "properties": {
-          "editType": {
-            "$ref": "#/components/schemas/EditType"
-          },
-          "timestamp": {
-            "type": "integer",
-            "format": "int64"
-          }
-        }
-      },
-      "EditMessageResponse": {
-        "type": "object",
-        "required": [
-          "sessionId"
-        ],
-        "properties": {
-          "sessionId": {
-            "type": "string"
-          }
-        }
-      },
-      "EditType": {
-        "type": "string",
-        "enum": [
-          "fork",
-          "edit"
-        ]
-      },
       "EmbeddedResource": {
         "type": "object",
         "required": [
@@ -3961,6 +3928,37 @@
           }
         }
       },
+      "ForkRequest": {
+        "type": "object",
+        "required": [
+          "truncate",
+          "copy"
+        ],
+        "properties": {
+          "copy": {
+            "type": "boolean"
+          },
+          "timestamp": {
+            "type": "integer",
+            "format": "int64",
+            "nullable": true
+          },
+          "truncate": {
+            "type": "boolean"
+          }
+        }
+      },
+      "ForkResponse": {
+        "type": "object",
+        "required": [
+          "sessionId"
+        ],
+        "properties": {
+          "sessionId": {
+            "type": "string"
+          }
+        }
+      },
       "FrontendToolRequest": {
         "type": "object",
         "required": [

ui/desktop/src/api/types.gen.ts

@@ -189,17 +189,6 @@ export type DetectProviderResponse = {
     provider_name: string;
 };
 
-export type EditMessageRequest = {
-    editType?: EditType;
-    timestamp: number;
-};
-
-export type EditMessageResponse = {
-    sessionId: string;
-};
-
-export type EditType = 'fork' | 'edit';
-
 export type EmbeddedResource = {
     _meta?: {
         [key: string]: unknown;
@@ -352,6 +341,16 @@ export type ExtensionResponse = {
     warnings?: Array<string>;
 };
 
+export type ForkRequest = {
+    copy: boolean;
+    timestamp?: number | null;
+    truncate: boolean;
+};
+
+export type ForkResponse = {
+    sessionId: string;
+};
+
 export type FrontendToolRequest = {
     id: string;
     toolCall: {
@@ -3178,46 +3177,6 @@ export type GetSessionResponses = {
 
 export type GetSessionResponse = GetSessionResponses[keyof GetSessionResponses];
 
-export type EditMessageData = {
-    body: EditMessageRequest;
-    path: {
-        /**
-         * Unique identifier for the session
-         */
-        session_id: string;
-    };
-    query?: never;
-    url: '/sessions/{session_id}/edit_message';
-};
-
-export type EditMessageErrors = {
-    /**
-     * Bad request - Invalid message timestamp
-     */
-    400: unknown;
-    /**
-     * Unauthorized - Invalid or missing API key
-     */
-    401: unknown;
-    /**
-     * Session or message not found
-     */
-    404: unknown;
-    /**
-     * Internal server error
-     */
-    500: unknown;
-};
-
-export type EditMessageResponses = {
-    /**
-     * Session prepared for editing - frontend should submit the edited message
-     */
-    200: EditMessageResponse;
-};
-
-export type EditMessageResponse2 = EditMessageResponses[keyof EditMessageResponses];
-
 export type ExportSessionData = {
     body?: never;
     path: {
@@ -3290,6 +3249,46 @@ export type GetSessionExtensionsResponses = {
 
 export type GetSessionExtensionsResponse = GetSessionExtensionsResponses[keyof GetSessionExtensionsResponses];
 
+export type ForkSessionData = {
+    body: ForkRequest;
+    path: {
+        /**
+         * Unique identifier for the session
+         */
+        session_id: string;
+    };
+    query?: never;
+    url: '/sessions/{session_id}/fork';
+};
+
+export type ForkSessionErrors = {
+    /**
+     * Bad request - truncate=true requires timestamp
+     */
+    400: unknown;
+    /**
+     * Unauthorized - Invalid or missing API key
+     */
+    401: unknown;
+    /**
+     * Session not found
+     */
+    404: unknown;
+    /**
+     * Internal server error
+     */
+    500: unknown;
+};
+
+export type ForkSessionResponses = {
+    /**
+     * Session forked successfully
+     */
+    200: ForkResponse;
+};
+
+export type ForkSessionResponse = ForkSessionResponses[keyof ForkSessionResponses];
+
 export type UpdateSessionNameData = {
     body: UpdateSessionNameRequest;
     path: {

ui/desktop/src/components/ChatInput.tsx

@@ -1,6 +1,6 @@
 import { AppEvents } from '../constants/events';
 import React, { useRef, useState, useEffect, useMemo, useCallback } from 'react';
-import { Bug, ScrollText, ChefHat } from 'lucide-react';
+import { Bug, ChefHat, ScrollText } from 'lucide-react';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/Tooltip';
 import { Button } from './ui/button';
 import type { View } from '../utils/navigationUtils';

ui/desktop/src/components/sessions/SessionListView.tsx

@@ -11,6 +11,7 @@ import {
   Download,
   Upload,
   ExternalLink,
+  Copy,
   Puzzle,
 } from 'lucide-react';
 import { Card } from '../ui/card';
@@ -28,6 +29,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '../ui/
 import {
   deleteSession,
   exportSession,
+  forkSession,
   importSession,
   listSessions,
   Session,
@@ -436,6 +438,25 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
       setShowDeleteConfirmation(true);
     }, []);
 
+    const handleDuplicateSession = useCallback(
+      async (session: Session) => {
+        try {
+          await forkSession({
+            path: { session_id: session.id },
+            body: { truncate: false, copy: true },
+            throwOnError: true,
+          });
+          toast.success(`Session "${session.name}" duplicated successfully`);
+          await loadSessions();
+        } catch (error) {
+          console.error('Error duplicating session:', error);
+          const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+          toast.error(`Failed to duplicate session: ${errorMessage}`);
+        }
+      },
+      [loadSessions]
+    );
+
     const handleConfirmDelete = useCallback(async () => {
       if (!sessionToDelete) return;
 
@@ -530,27 +551,37 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
     const SessionItem = React.memo(function SessionItem({
       session,
       onEditClick,
+      onDuplicateClick,
       onDeleteClick,
       onExportClick,
       onOpenInNewWindow,
     }: {
       session: Session;
       onEditClick: (session: Session) => void;
+      onDuplicateClick: (session: Session) => void;
       onDeleteClick: (session: Session) => void;
       onExportClick: (session: Session, e: React.MouseEvent) => void;
       onOpenInNewWindow: (session: Session, e: React.MouseEvent) => void;
     }) {
       const handleEditClick = useCallback(
         (e: React.MouseEvent) => {
-          e.stopPropagation(); // Prevent card click
+          e.stopPropagation();
           onEditClick(session);
         },
         [onEditClick, session]
       );
 
+      const handleDuplicateClick = useCallback(
+        (e: React.MouseEvent) => {
+          e.stopPropagation();
+          onDuplicateClick(session);
+        },
+        [onDuplicateClick, session]
+      );
+
       const handleDeleteClick = useCallback(
         (e: React.MouseEvent) => {
-          e.stopPropagation(); // Prevent card click
+          e.stopPropagation();
           onDeleteClick(session);
         },
         [onDeleteClick, session]
@@ -605,6 +636,13 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
               >
                 <Edit2 className="w-3 h-3 text-textSubtle hover:text-textStandard" />
               </button>
+              <button
+                onClick={handleDuplicateClick}
+                className="p-2 rounded hover:bg-gray-100 dark:hover:bg-gray-700 cursor-pointer"
+                title="Duplicate session"
+              >
+                <Copy className="w-3 h-3 text-textSubtle hover:text-textStandard" />
+              </button>
               <button
                 onClick={handleDeleteClick}
                 className="p-2 rounded hover:bg-red-50 dark:hover:bg-red-900/20 cursor-pointer transition-colors"
@@ -757,6 +795,7 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
                     key={session.id}
                     session={session}
                     onEditClick={handleEditSession}
+                    onDuplicateClick={handleDuplicateSession}
                     onDeleteClick={handleDeleteSession}
                     onExportClick={handleExportSession}
                     onOpenInNewWindow={handleOpenInNewWindow}

ui/desktop/src/hooks/useChatStream.ts

@@ -650,27 +650,28 @@ export function useChatStream({
       const currentState = stateRef.current;
 
       try {
-        const { editMessage } = await import('../api');
+        const { forkSession } = await import('../api');
         const message = currentState.messages.find((m) => m.id === messageId);
 
         if (!message) {
           throw new Error(`Message with id ${messageId} not found in current messages`);
         }
 
-        const response = await editMessage({
+        const response = await forkSession({
           path: {
             session_id: sessionId,
           },
           body: {
             timestamp: message.created,
-            editType,
+            truncate: true,
+            copy: editType === 'fork',
           },
           throwOnError: true,
         });
 
         const targetSessionId = response.data?.sessionId;
         if (!targetSessionId) {
-          throw new Error('No session ID returned from edit_message');
+          throw new Error('No session ID returned from fork');
         }
 
         if (editType === 'fork') {