Merge pull request #110 from will-lp1/persist-tool-state

will-lp1 · web-flow · commit 1195fd06f13c · 2025-11-18T18:56:51.000Z
Persist tool state
diff --git a/apps/saru/app/api/messages/route.ts b/apps/saru/app/api/messages/route.ts
@@ -1,7 +1,7 @@
 import { NextResponse } from 'next/server';
 import { auth } from "@/lib/auth";
 import { headers } from 'next/headers';
-import { getMessagesByChatId, getChatById } from '@/lib/db/queries';
+import { getMessagesByChatId, getChatById, getMessageById, updateToolMetadata } from '@/lib/db/queries';
 
 export async function GET(request: Request) {
   try {
@@ -45,4 +45,54 @@ export async function GET(request: Request) {
     console.error('Error fetching messages:', error);
     return NextResponse.json({ error: 'Error fetching messages' }, { status: 500 });
   }
-} 
+}
+
+export async function PATCH(request: Request) {
+  try {
+    const readonlyHeaders = await headers();
+    const requestHeaders = new Headers(readonlyHeaders);
+    const session = await auth.api.getSession({ headers: requestHeaders });
+
+    if (!session?.user?.id) {
+      console.error('Session error in PATCH /api/messages');
+      return NextResponse.json({ error: 'Authentication error' }, { status: 401 });
+    }
+    const userId = session.user.id;
+
+    const body = await request.json();
+    const { messageId, chatId, toolCallId, applied, rejected } = body;
+
+    if (!messageId || !chatId || !toolCallId) {
+      return NextResponse.json(
+        { error: 'messageId, chatId, and toolCallId are required' },
+        { status: 400 }
+      );
+    }
+
+    // Verify the chat belongs to the user
+    const chat = await getChatById({ id: chatId });
+    if (!chat) {
+      return NextResponse.json({ error: 'Chat not found' }, { status: 404 });
+    }
+    if (chat.userId !== userId) {
+      return NextResponse.json({ error: 'Unauthorized' }, { status: 403 });
+    }
+
+    // Verify the message exists and belongs to the chat
+    const message = await getMessageById({ id: messageId });
+    if (!message) {
+      return NextResponse.json({ error: 'Message not found' }, { status: 404 });
+    }
+    if (message.chatId !== chatId) {
+      return NextResponse.json({ error: 'Message does not belong to this chat' }, { status: 403 });
+    }
+
+    // Update the tool metadata
+    await updateToolMetadata({ messageId, toolCallId, applied, rejected });
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    console.error('Error updating message:', error);
+    return NextResponse.json({ error: 'Error updating message' }, { status: 500 });
+  }
+}
diff --git a/apps/saru/components/chat/message.tsx b/apps/saru/components/chat/message.tsx
@@ -232,6 +232,8 @@ const PurePreviewMessage = ({
                         type={actionType}
                         result={{ ...result, toolCallId: p.toolCallId } as DocumentToolResultProps['result']}
                         isReadonly={isReadonly}
+                        chatId={chatId}
+                        messageId={message.id}
                       />
                     );
                   }
diff --git a/apps/saru/components/document/document-tool.tsx b/apps/saru/components/document/document-tool.tsx
@@ -106,24 +106,25 @@ export interface DocumentToolResultProps {
     content?: string;
     message?: string;
     toolCallId?: string;
+    applied?: boolean;
+    rejected?: boolean;
   };
   isReadonly: boolean;
+  chatId?: string;
+  messageId?: string;
 }
 
 function PureDocumentToolResult({
   type,
   result,
   isReadonly,
+  chatId,
+  messageId,
 }: DocumentToolResultProps) {
   const { document, setDocument } = useDocument();
   const [isSaving, setIsSaving] = useState(false);
-  const [isApplied, setIsApplied] = useState(() => {
-    if (type === 'update' && result.id && document.documentId === result.id) {
-      return document.content === result.newContent;
-    }
-    return false;
-  });
-  const [isRejected, setIsRejected] = useState(false);
+  const [isApplied, setIsApplied] = useState(() => result.applied ?? false);
+  const [isRejected, setIsRejected] = useState(() => result.rejected ?? false);
 
   const isUpdateProposal =
     type === 'update' && 
@@ -132,7 +133,7 @@ function PureDocumentToolResult({
     result.originalContent !== result.newContent;
 
   useEffect(() => {
-    if (isUpdateProposal && result.id && result.newContent) {
+    if (isUpdateProposal && result.id && result.newContent && !isApplied && !isRejected) {
       const event = new CustomEvent('preview-document-update', {
         detail: {
           documentId: result.id,
@@ -142,7 +143,7 @@ function PureDocumentToolResult({
       });
       window.dispatchEvent(event);
     }
-  }, [isUpdateProposal, result.id, result.newContent, result.originalContent]);
+  }, [isUpdateProposal, result.id, result.newContent, result.originalContent, isApplied, isRejected]);
 
   const handleApplyUpdate = useCallback(() => {
     if (type !== 'update' || !result.newContent || !result.id || isSaving) return;
@@ -171,6 +172,24 @@ function PureDocumentToolResult({
       detail: { documentId: result.id, newContent: result.newContent, transient: false },
     }));
     setIsApplied(true);
+
+    // Persist applied state to chat metadata
+    if (chatId && messageId && result.toolCallId) {
+      fetch('/api/messages', {
+        method: 'PATCH',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+          messageId,
+          chatId,
+          toolCallId: result.toolCallId,
+          applied: true,
+          rejected: false,
+        }),
+      }).catch(err => {
+        console.error('[DocumentToolResult] Failed to persist applied state:', err);
+      });
+    }
+
     fetch('/api/document', {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
@@ -196,6 +215,23 @@ function PureDocumentToolResult({
 
     setIsRejected(true);
 
+    // Persist rejected state to chat metadata
+    if (chatId && messageId && result.toolCallId) {
+      fetch('/api/messages', {
+        method: 'PATCH',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+          messageId,
+          chatId,
+          toolCallId: result.toolCallId,
+          applied: false,
+          rejected: true,
+        }),
+      }).catch(err => {
+        console.error('[DocumentToolResult] Failed to persist rejected state:', err);
+      });
+    }
+
     window.dispatchEvent(new CustomEvent('tool-result', {
       detail: {
         toolCallId: result.toolCallId,
@@ -219,7 +255,7 @@ function PureDocumentToolResult({
     });
     window.dispatchEvent(event);
     toast.info('Update proposal rejected.');
-  }, [result.id, result.title, result.originalContent, result.newContent, result.toolCallId, type]);
+  }, [result.id, result.title, result.originalContent, result.newContent, result.toolCallId, type, chatId, messageId]);
 
   if (result.error) {
      return (
diff --git a/apps/saru/components/document/version-rail.tsx b/apps/saru/components/document/version-rail.tsx
@@ -69,13 +69,19 @@ export function VersionRail({ versions, currentIndex, onIndexChange, baseDocumen
     });
   },[versions]);
   React.useEffect(() => {
-    if (currentIndex === versions.length - 1) {
-      setSelectedIndex(currentIndex);
-    } else {
-      setSelectedIndex(currentIndex);
-    }
-  }, [currentIndex, versions.length]);
-  
+    setSelectedIndex(currentIndex);
+  }, [currentIndex]);
+
+  React.useEffect(() => {
+    return () => {
+      window.dispatchEvent(
+        new CustomEvent('cancel-document-update', {
+          detail: { documentId: baseDocumentId },
+        })
+      );
+    };
+  }, [baseDocumentId]);
+
   const isViewingHistory = selectedIndex !== null && selectedIndex < versions.length - 1;
   if (isLoading || versions.length <= 1) {
     return <div className="w-full border-b bg-background h-1 group-hover:h-12 transition-all duration-200" />;
@@ -191,24 +197,22 @@ export function VersionRail({ versions, currentIndex, onIndexChange, baseDocumen
 
   const handlePointerLeave = () => {
     setPressStart(null);
-    if (hoverIndex !== null) {
-      if (isViewingHistory && selectedIndex !== null && selectedIndex < versions.length - 1) {
-        const v = versions[selectedIndex];
-        window.dispatchEvent(
-          new CustomEvent('preview-document-update', {
-            detail: { documentId: baseDocumentId, newContent: v.content },
-          })
-        );
-      } else {
-        window.dispatchEvent(
-          new CustomEvent('cancel-document-update', {
-            detail: { documentId: baseDocumentId },
-          })
-        );
-      }
-      setHoverIndex(null);
-    }
+    setHoverIndex(null);
 
+    if (selectedIndex !== null && selectedIndex < versions.length - 1) {
+      const v = versions[selectedIndex];
+      window.dispatchEvent(
+        new CustomEvent('preview-document-update', {
+          detail: { documentId: baseDocumentId, newContent: v.content },
+        })
+      );
+    } else {
+      window.dispatchEvent(
+        new CustomEvent('cancel-document-update', {
+          detail: { documentId: baseDocumentId },
+        })
+      );
+    }
   };
 
   const Tooltip = ({ active, payload }: { active?: boolean; payload?: Array<{ payload: { ts: string; additions: number; deletions: number } }> }) => {
diff --git a/apps/saru/lib/db/queries.ts b/apps/saru/lib/db/queries.ts

Original file line number	Diff line number	Diff line change
`@@ -232,6 +232,8 @@ const PurePreviewMessage = ({`
`232`	`232`	`type={actionType}`
`233`	`233`	`result={{ ...result, toolCallId: p.toolCallId } as DocumentToolResultProps['result']}`
`234`	`234`	`isReadonly={isReadonly}`
	`235`	`+ chatId={chatId}`
	`236`	`+ messageId={message.id}`
`235`	`237`	`/>`
`236`	`238`	`);`
`237`	`239`	`}`