Add session forking capability (#5882)

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

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') {