From 999ffd1bdb072cbc554549e284b880c4959cfc55 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 14 Mar 2025 17:27:30 -0700
Subject: [PATCH] workign run test connection many times

---
 .../add_model/ConnectionErrorDisplay.tsx      | 11 +++++-
 .../components/add_model/add_model_tab.tsx    | 37 ++++++++++++-------
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/add_model/ConnectionErrorDisplay.tsx b/ui/litellm-dashboard/src/components/add_model/ConnectionErrorDisplay.tsx
index 7eb7e2a190..4a3fc30377 100644
--- a/ui/litellm-dashboard/src/components/add_model/ConnectionErrorDisplay.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ConnectionErrorDisplay.tsx
@@ -12,6 +12,7 @@ interface ConnectionErrorDisplayProps {
   testMode: string;
   modelName?: string;
   onClose?: () => void;
+  onTestComplete?: () => void;
 }
 
 const ConnectionErrorDisplay: React.FC<ConnectionErrorDisplayProps> = ({ 
@@ -19,7 +20,8 @@ const ConnectionErrorDisplay: React.FC<ConnectionErrorDisplayProps> = ({
   accessToken, 
   testMode, 
   modelName = "this model", 
-  onClose 
+  onClose,
+  onTestComplete
 }) => {
   const [error, setError] = React.useState<Error | string | null>(null);
   const [rawRequest, setRawRequest] = React.useState<any>(null);
@@ -30,6 +32,12 @@ const ConnectionErrorDisplay: React.FC<ConnectionErrorDisplayProps> = ({
 
   const testModelConnection = async () => {
     setIsLoading(true);
+    setShowDetails(false);
+    setError(null);
+    setRawRequest(null);
+    setRawResponse(null);
+    setIsSuccess(false);
+    
     try {
       const result = await prepareModelAddRequest(formValues, accessToken, null);
       if (!result) throw new Error("Failed to prepare model data");
@@ -55,6 +63,7 @@ const ConnectionErrorDisplay: React.FC<ConnectionErrorDisplayProps> = ({
       setIsSuccess(false);
     } finally {
       setIsLoading(false);
+      if (onTestComplete) onTestComplete();
     }
   };
 
diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
index e67dba781e..1ae625d590 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
@@ -45,16 +45,21 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
   credentials,
   accessToken,
 }) => {
-  // Add state for test mode and connection error
+  // State for test mode and connection testing
   const [testMode, setTestMode] = useState<string>("chat");
   const [isResultModalVisible, setIsResultModalVisible] = useState<boolean>(false);
   const [isTestingConnection, setIsTestingConnection] = useState<boolean>(false);
+  // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test
+  const [connectionTestId, setConnectionTestId] = useState<string>("");
 
-  // Test connection directly when button is clicked
+  // Test connection when button is clicked
   const handleTestConnection = async () => {
     setIsTestingConnection(true);
+    // Generate a new test ID (using timestamp for uniqueness)
+    // This forces React to create a new instance of ConnectionErrorDisplay
+    setConnectionTestId(`test-${Date.now()}`);
+    // Show the modal with the fresh test
     setIsResultModalVisible(true);
-    // The actual testing is handled in ConnectionErrorDisplay component
   };
 
   return (
@@ -241,16 +246,22 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
         ]}
         width={700}
       >
-        <ConnectionErrorDisplay 
-          formValues={form.getFieldsValue()}
-          accessToken={accessToken}
-          testMode={testMode}
-          modelName={form.getFieldValue('model_name') || form.getFieldValue('model')}
-          onClose={() => {
-            setIsResultModalVisible(false);
-            setIsTestingConnection(false);
-          }}
-        />
+        {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */}
+        {isResultModalVisible && (
+          <ConnectionErrorDisplay 
+            // The key prop tells React to create a fresh component instance when it changes
+            key={connectionTestId}
+            formValues={form.getFieldsValue()}
+            accessToken={accessToken}
+            testMode={testMode}
+            modelName={form.getFieldValue('model_name') || form.getFieldValue('model')}
+            onClose={() => {
+              setIsResultModalVisible(false);
+              setIsTestingConnection(false);
+            }}
+            onTestComplete={() => setIsTestingConnection(false)}
+          />
+        )}
       </Modal>
     </>
   );