Use test connection button

Author: ajhai

llmstack/client/src/components/connections/AddConnectionModal.jsx

@@ -34,6 +34,7 @@ function AddConnectionModal({ open, onCancelCb, onSaveCb, connection }) {
   );
   const [localConnection, setLocalConnection] = useState(connection || {});
   const [validationErrors, setValidationErrors] = useState([]);
+  const [connectionActive, setConnectionActive] = useState(false);
 
   const validateForm = () => {
     let errors = [];
@@ -105,6 +106,7 @@ function AddConnectionModal({ open, onCancelCb, onSaveCb, connection }) {
               enqueueSnackbar("Connection test successful", {
                 variant: "success",
               });
+              setConnectionActive(true);
               ws.close();
               onCancelCb();
             }
@@ -268,32 +270,37 @@ function AddConnectionModal({ open, onCancelCb, onSaveCb, connection }) {
         />
       </DialogContent>
       <DialogActions>
-        <Button
-          onClick={testConnection({
-            ...localConnection,
-            ...{
-              provider_slug: connectionType?.provider_slug,
-              connection_type_slug: connectionType?.slug,
-            },
-          })}
-        >
-          Test Connection
-        </Button>
         <Button onClick={handleCloseCb}>Cancel</Button>
-        <Button
-          onClick={() =>
-            handleSaveCb({
+        {!connectionActive && (
+          <Button
+            onClick={testConnection({
               ...localConnection,
               ...{
                 provider_slug: connectionType?.provider_slug,
                 connection_type_slug: connectionType?.slug,
               },
-            })
-          }
-          variant="contained"
-        >
-          Save
-        </Button>
+            })}
+            variant="contained"
+          >
+            Test Connection
+          </Button>
+        )}
+        {connectionActive && (
+          <Button
+            onClick={() =>
+              handleSaveCb({
+                ...localConnection,
+                ...{
+                  provider_slug: connectionType?.provider_slug,
+                  connection_type_slug: connectionType?.slug,
+                },
+              })
+            }
+            variant="contained"
+          >
+            Save
+          </Button>
+        )}
       </DialogActions>
     </Dialog>
   );

llmstack/connections/handlers/linkedin_login.py

@@ -55,7 +55,13 @@ async def activate(self, connection) -> Iterator[Union[Connection, dict]]:
                 return
 
             # Wait for login to complete and redirect to /feed/
-            await page.wait_for_url('https://www.linkedin.com/feed/')
+            await page.wait_for_url('https://www.linkedin.com/feed/', timeout=5000)
+
+            if page.url != 'https://www.linkedin.com/feed/':
+                connection.status = ConnectionStatus.FAILED
+                await browser.close()
+                yield {'error': 'Login failed', 'connection': connection}
+                return
 
             # Get storage state
             storage_state = await context.storage_state()