From b7057b1eda97417960c00ca4748e2eacff88285e Mon Sep 17 00:00:00 2001
From: Zane <75694352+zanesq@users.noreply.github.com>
Date: Thu, 17 Jul 2025 19:43:35 -0700
Subject: [PATCH] Fix tool call allow still showing initial state in chat after
 navigating back (#3498)

---
 ui/desktop/src/components/GooseMessage.tsx    |  2 +-
 .../src/components/ToolCallConfirmation.tsx   | 98 ++++++++++++++++---
 2 files changed, 84 insertions(+), 16 deletions(-)

diff --git a/ui/desktop/src/components/GooseMessage.tsx b/ui/desktop/src/components/GooseMessage.tsx
index 1fea8eee..62bbe953 100644
--- a/ui/desktop/src/components/GooseMessage.tsx
+++ b/ui/desktop/src/components/GooseMessage.tsx
@@ -220,7 +220,7 @@ export default function GooseMessage({
         {hasToolConfirmation && (
           <ToolCallConfirmation
             isCancelledMessage={messageIndex == messageHistoryIndex - 1}
-            isClicked={messageIndex < messageHistoryIndex - 1}
+            isClicked={messageIndex < messageHistoryIndex}
             toolConfirmationId={toolConfirmationContent.id}
             toolName={toolConfirmationContent.toolName}
           />
diff --git a/ui/desktop/src/components/ToolCallConfirmation.tsx b/ui/desktop/src/components/ToolCallConfirmation.tsx
index f2f9dfcc..764a4604 100644
--- a/ui/desktop/src/components/ToolCallConfirmation.tsx
+++ b/ui/desktop/src/components/ToolCallConfirmation.tsx
@@ -1,4 +1,4 @@
-import { useState } from 'react';
+import { useState, useEffect } from 'react';
 import { snakeToTitleCase } from '../utils';
 import PermissionModal from './settings/permission/PermissionModal';
 import { ChevronRight } from 'lucide-react';
@@ -9,6 +9,17 @@ const ALWAYS_ALLOW = 'always_allow';
 const ALLOW_ONCE = 'allow_once';
 const DENY = 'deny';
 
+// Global state to track tool confirmation decisions
+// This persists across navigation within the same session
+const toolConfirmationState = new Map<
+  string,
+  {
+    clicked: boolean;
+    status: string;
+    actionDisplay: string;
+  }
+>();
+
 interface ToolConfirmationProps {
   isCancelledMessage: boolean;
   isClicked: boolean;
@@ -22,30 +33,75 @@ export default function ToolConfirmation({
   toolConfirmationId,
   toolName,
 }: ToolConfirmationProps) {
-  const [clicked, setClicked] = useState(isClicked);
-  const [status, setStatus] = useState('unknown');
-  const [actionDisplay, setActionDisplay] = useState('');
+  // Check if we have a stored state for this tool confirmation
+  const storedState = toolConfirmationState.get(toolConfirmationId);
+
+  // Initialize state from stored state if available, otherwise use props/defaults
+  const [clicked, setClicked] = useState(storedState?.clicked ?? isClicked);
+  const [status, setStatus] = useState(storedState?.status ?? 'unknown');
+  const [actionDisplay, setActionDisplay] = useState(storedState?.actionDisplay ?? '');
   const [isModalOpen, setIsModalOpen] = useState(false);
 
-  const handleButtonClick = async (action: string) => {
-    setClicked(true);
-    setStatus(action);
-    if (action === ALWAYS_ALLOW) {
-      setActionDisplay('always allowed');
-    } else if (action === ALLOW_ONCE) {
-      setActionDisplay('allowed once');
-    } else {
-      setActionDisplay('denied');
+  // Sync internal state with stored state and props
+  useEffect(() => {
+    const currentStoredState = toolConfirmationState.get(toolConfirmationId);
+
+    // If we have stored state, use it
+    if (currentStoredState) {
+      setClicked(currentStoredState.clicked);
+      setStatus(currentStoredState.status);
+      setActionDisplay(currentStoredState.actionDisplay);
+    } else if (isClicked && !clicked) {
+      // Fallback to prop-based logic for historical confirmations
+      setClicked(isClicked);
+      if (status === 'unknown') {
+        setStatus('confirmed');
+        setActionDisplay('confirmed');
+
+        // Store this state for future renders
+        toolConfirmationState.set(toolConfirmationId, {
+          clicked: true,
+          status: 'confirmed',
+          actionDisplay: 'confirmed',
+        });
+      }
     }
+  }, [isClicked, clicked, status, toolName, toolConfirmationId]);
+
+  const handleButtonClick = async (action: string) => {
+    const newClicked = true;
+    const newStatus = action;
+    let newActionDisplay = '';
+
+    if (action === ALWAYS_ALLOW) {
+      newActionDisplay = 'always allowed';
+    } else if (action === ALLOW_ONCE) {
+      newActionDisplay = 'allowed once';
+    } else {
+      newActionDisplay = 'denied';
+    }
+
+    // Update local state
+    setClicked(newClicked);
+    setStatus(newStatus);
+    setActionDisplay(newActionDisplay);
+
+    // Store in global state for persistence across navigation
+    toolConfirmationState.set(toolConfirmationId, {
+      clicked: newClicked,
+      status: newStatus,
+      actionDisplay: newActionDisplay,
+    });
+
     try {
       const response = await confirmPermission({
         body: { id: toolConfirmationId, action, principal_type: 'Tool' },
       });
       if (response.error) {
-        console.error('Failed to confirm permission: ', response.error);
+        console.error('Failed to confirm permission:', response.error);
       }
     } catch (err) {
-      console.error('Error fetching tools:', err);
+      console.error('Error confirming permission:', err);
     }
   };
 
@@ -106,6 +162,18 @@ export default function ToolConfirmation({
                 <path strokeLinecap="round" strokeLinejoin="round" d="M6 18L18 6M6 6l12 12" />
               </svg>
             )}
+            {status === 'confirmed' && (
+              <svg
+                className="w-5 h-5 text-gray-500"
+                xmlns="http://www.w3.org/2000/svg"
+                fill="none"
+                viewBox="0 0 24 24"
+                stroke="currentColor"
+                strokeWidth={2}
+              >
+                <path strokeLinecap="round" strokeLinejoin="round" d="M5 13l4 4L19 7" />
+              </svg>
+            )}
             <span className="ml-2 text-textStandard">
               {isClicked
                 ? 'Tool confirmation is not available'