From 80342f60534edf82dd0d3432416a56846617f41f Mon Sep 17 00:00:00 2001
From: Max Techera <maxi.techerag@gmail.com>
Date: Thu, 19 Feb 2026 19:11:52 -0300
Subject: [PATCH] fix: enable drive picker in docstores

---
 .../ui-component/drive/GoogleDrivePicker.jsx  | 59 ++++++++++++---
 .../views/docstore/DocStoreInputHandler.jsx   | 72 ++++++++++---------
 2 files changed, 88 insertions(+), 43 deletions(-)

diff --git a/packages/ui/src/ui-component/drive/GoogleDrivePicker.jsx b/packages/ui/src/ui-component/drive/GoogleDrivePicker.jsx
index 83ad31577d7..e58b08b5c1e 100644
--- a/packages/ui/src/ui-component/drive/GoogleDrivePicker.jsx
+++ b/packages/ui/src/ui-component/drive/GoogleDrivePicker.jsx
@@ -130,7 +130,7 @@ const useGooglePicker = (accessToken, onFilesSelected) => {
     const createPicker = useCallback(async () => {
         if (!accessToken || !window.google?.picker) {
             console.error('Google Picker not available or no access token')
-            return
+            return false
         }
 
         try {
@@ -149,8 +149,10 @@ const useGooglePicker = (accessToken, onFilesSelected) => {
 
             picker.setVisible(true)
             setPickerInstance(picker)
+            return true
         } catch (error) {
             console.error('Error creating picker:', error)
+            return false
         }
     }, [accessToken, createPickerView, pickerCallback])
 
@@ -185,6 +187,17 @@ export const GoogleDrivePicker = ({ onChange, value, disabled, credentialId, cre
     const getCredentialDataApi = useApi(credentialsApi.getSpecificCredential)
     const scriptsLoaded = useGoogleAPILoader(accessToken)
 
+    const getTokenStateFromCredential = useCallback((plainDataObj = {}) => {
+        const token = plainDataObj.access_token || plainDataObj.googleAccessToken || ''
+        const expiresAt = plainDataObj.expires_at || plainDataObj.expiresAt
+        const isExpired = expiresAt ? new Date(expiresAt) < new Date() : false
+
+        return {
+            token,
+            isExpired
+        }
+    }, [])
+
     // Handle new files selection from picker
     const handleFilesSelected = useCallback(
         (newFiles) => {
@@ -219,23 +232,21 @@ export const GoogleDrivePicker = ({ onChange, value, disabled, credentialId, cre
     // Handle credential data from prop
     useEffect(() => {
         if (credentialData?.plainDataObj) {
-            const { access_token, expires_at } = credentialData.plainDataObj
-            const isExpired = expires_at ? new Date(expires_at) < new Date() : false
+            const { token, isExpired } = getTokenStateFromCredential(credentialData.plainDataObj)
             setIsTokenExpired(isExpired)
-            setAccessToken(access_token ?? '')
+            setAccessToken(token)
         }
-    }, [credentialData])
+    }, [credentialData, getTokenStateFromCredential])
 
     // Handle credential data from API
     useEffect(() => {
         if (getCredentialDataApi.data) {
-            const { access_token, expires_at } = getCredentialDataApi.data?.plainDataObj
-            const isExpired = expires_at ? new Date(expires_at) < new Date() : false
+            const { token, isExpired } = getTokenStateFromCredential(getCredentialDataApi.data?.plainDataObj)
             setIsTokenExpired(isExpired)
-            setAccessToken(access_token ?? '')
+            setAccessToken(token)
             handleCredentialDataChange?.(getCredentialDataApi.data)
         }
-    }, [getCredentialDataApi.data, handleCredentialDataChange])
+    }, [getCredentialDataApi.data, getTokenStateFromCredential, handleCredentialDataChange])
 
     // Handle outside clicks and keyboard events for picker
     useEffect(() => {
@@ -320,6 +331,32 @@ export const GoogleDrivePicker = ({ onChange, value, disabled, credentialId, cre
         }
     }, [credentialId, getCredentialDataApi, showSnackbar])
 
+    const handleOpenPicker = useCallback(async () => {
+        const opened = await createPicker()
+        if (opened) return
+
+        if (!credentialId) {
+            showSnackbar('Select a credential before opening Google Drive.', 'warning')
+            return
+        }
+
+        if (!accessToken) {
+            showSnackbar('No access token available. Refresh or re-authenticate to load files.', 'error')
+            return
+        }
+
+        if (!scriptsLoaded) {
+            showSnackbar('Google picker is still loading. Try again in a moment.', 'info')
+            return
+        }
+
+        if (isTokenExpired) {
+            showSnackbar('Access token expired. Refresh or re-authenticate before selecting files.', 'warning')
+        } else {
+            showSnackbar('Unable to open Google Drive picker. Please try again.', 'error')
+        }
+    }, [accessToken, createPicker, credentialId, isTokenExpired, scriptsLoaded, showSnackbar])
+
     const handleReauthenticate = useCallback(async () => {
         if (!credentialId) return
 
@@ -358,7 +395,7 @@ export const GoogleDrivePicker = ({ onChange, value, disabled, credentialId, cre
         }
     }, [credentialId, getCredentialDataApi, showSnackbar])
 
-    const isPickerReady = scriptsLoaded && !disabled && !isTokenExpired
+    const isPickerReady = !disabled
     const hasSelectedFiles = selectedFiles.length > 0
 
     return (
@@ -366,7 +403,7 @@ export const GoogleDrivePicker = ({ onChange, value, disabled, credentialId, cre
             <Stack direction='column' spacing={2} sx={{ mb: 2 }}>
                 <Button
                     variant='outlined'
-                    onClick={createPicker}
+                    onClick={handleOpenPicker}
                     disabled={!isPickerReady}
                     sx={{
                         '&.Mui-disabled': {
diff --git a/packages/ui/src/views/docstore/DocStoreInputHandler.jsx b/packages/ui/src/views/docstore/DocStoreInputHandler.jsx
index 1cc2ae02538..e3d999221e5 100644
--- a/packages/ui/src/views/docstore/DocStoreInputHandler.jsx
+++ b/packages/ui/src/views/docstore/DocStoreInputHandler.jsx
@@ -1,5 +1,5 @@
 import PropTypes from 'prop-types'
-import { useState, useContext, useCallback } from 'react'
+import { useState, useContext, useCallback, useEffect } from 'react'
 import { useSelector } from 'react-redux'
 
 // material-ui
@@ -42,9 +42,16 @@ const DocStoreInputHandler = ({ inputParam, data, disabled = false, onNodeDataCh
     const [showManageScrapedLinksDialog, setShowManageScrapedLinksDialog] = useState(false)
     const [manageScrapedLinksDialogProps, setManageScrapedLinksDialogProps] = useState({})
     const [reloadTimestamp, setReloadTimestamp] = useState(Date.now().toString())
-    const [selectedCredential, setSelectedCredential] = useState(data.credential || null)
+    const [selectedCredential, setSelectedCredential] = useState(
+        data.credential || data.inputs?.credential || data.inputs?.[FLOWISE_CREDENTIAL_ID] || null
+    )
     const [selectedCredentialData, setSelectedCredentialData] = useState(null)
 
+    useEffect(() => {
+        const credentialFromData = data.credential || data.inputs?.credential || data.inputs?.[FLOWISE_CREDENTIAL_ID] || null
+        setSelectedCredential((prevCredential) => (prevCredential === credentialFromData ? prevCredential : credentialFromData))
+    }, [data.credential, data.inputs?.credential, data.inputs?.[FLOWISE_CREDENTIAL_ID]])
+
     const handleCredentialDataChange = useCallback((credData) => {
         setSelectedCredentialData(credData)
     }, [])
@@ -275,37 +282,38 @@ const DocStoreInputHandler = ({ inputParam, data, disabled = false, onNodeDataCh
                                 value={data.inputs[inputParam.name] ?? inputParam.default ?? 'choose an option'}
                             />
                         )}
-                        {(inputParam.type === 'asyncOptions' || inputParam.type === 'asyncMultiOptions') && (
-                            <>
-                                {data.inputParams?.length === 1 && <div style={{ marginTop: 10 }} />}
-                                <div style={{ display: 'flex', flexDirection: 'row' }}>
-                                    <div key={reloadTimestamp} style={{ flex: 1 }}>
-                                        <AsyncDropdown
-                                            key={JSON.stringify(inputParam)}
-                                            disabled={disabled}
-                                            name={inputParam.name}
-                                            nodeData={data}
-                                            freeSolo={inputParam.freeSolo}
-                                            multiple={inputParam.type === 'asyncMultiOptions'}
-                                            value={data.inputs[inputParam.name] ?? inputParam.default ?? 'choose an option'}
-                                            onSelect={(newValue) => handleDataChange({ inputParam, newValue })}
-                                            onCreateNew={() => addAsyncOption(inputParam.name)}
-                                            fullWidth={true}
-                                        />
+                        {(inputParam.type === 'asyncOptions' || inputParam.type === 'asyncMultiOptions') &&
+                            !(data.name === 'googleDrive' && inputParam.name === 'selectedFiles') && (
+                                <>
+                                    {data.inputParams?.length === 1 && <div style={{ marginTop: 10 }} />}
+                                    <div style={{ display: 'flex', flexDirection: 'row' }}>
+                                        <div key={reloadTimestamp} style={{ flex: 1 }}>
+                                            <AsyncDropdown
+                                                key={JSON.stringify(inputParam)}
+                                                disabled={disabled}
+                                                name={inputParam.name}
+                                                nodeData={data}
+                                                freeSolo={inputParam.freeSolo}
+                                                multiple={inputParam.type === 'asyncMultiOptions'}
+                                                value={data.inputs[inputParam.name] ?? inputParam.default ?? 'choose an option'}
+                                                onSelect={(newValue) => handleDataChange({ inputParam, newValue })}
+                                                onCreateNew={() => addAsyncOption(inputParam.name)}
+                                                fullWidth={true}
+                                            />
+                                        </div>
+                                        {inputParam.refresh && (
+                                            <IconButton
+                                                title='Refresh'
+                                                color='primary'
+                                                size='small'
+                                                onClick={() => setReloadTimestamp(Date.now().toString())}
+                                            >
+                                                <IconRefresh />
+                                            </IconButton>
+                                        )}
                                     </div>
-                                    {inputParam.refresh && (
-                                        <IconButton
-                                            title='Refresh'
-                                            color='primary'
-                                            size='small'
-                                            onClick={() => setReloadTimestamp(Date.now().toString())}
-                                        >
-                                            <IconRefresh />
-                                        </IconButton>
-                                    )}
-                                </div>
-                            </>
-                        )}
+                                </>
+                            )}
                         {inputParam.type === 'array' && (
                             <ArrayRenderer inputParam={inputParam} data={data} disabled={disabled} isDocStore={true} />
                         )}