Add support for edit credentials

j-sanaa · j-sanaa · commit db17014d1d0f · 2026-03-24T14:01:05.000-07:00
diff --git a/packages/agentflow/src/features/node-editor/AsyncInput.test.tsx b/packages/agentflow/src/features/node-editor/AsyncInput.test.tsx
@@ -15,15 +15,26 @@ jest.mock('reactflow', () => ({
 
 jest.mock('@tabler/icons-react', () => ({
     IconArrowsMaximize: () => <span data-testid='icon-expand' />,
+    IconEdit: () => <span data-testid='icon-edit' />,
     IconVariable: () => <span data-testid='icon-variable' />,
     IconRefresh: () => <span data-testid='icon-refresh' />
 }))
 
 jest.mock('./CreateCredentialDialog', () => ({
-    CreateCredentialDialog: ({ open, onCreated, onClose }: { open: boolean; onCreated: (id: string) => void; onClose: () => void }) =>
+    CreateCredentialDialog: ({
+        open,
+        onCreated,
+        onClose,
+        editCredentialId
+    }: {
+        open: boolean
+        onCreated: (id: string) => void
+        onClose: () => void
+        editCredentialId?: string
+    }) =>
         open ? (
-            <div data-testid='create-credential-dialog'>
-                <button onClick={() => onCreated('new-cred-id')}>Create</button>
+            <div data-testid={editCredentialId ? 'edit-credential-dialog' : 'create-credential-dialog'}>
+                <button onClick={() => onCreated(editCredentialId ?? 'new-cred-id')}>{editCredentialId ? 'Save' : 'Create'}</button>
                 <button onClick={onClose}>Close</button>
             </div>
         ) : null
@@ -521,3 +532,96 @@ describe('AsyncInput – Create New credential', () => {
         expect(mockUseAsyncOptions.mock.calls.length).toBeGreaterThan(initialCallCount)
     })
 })
+
+describe('AsyncInput – Edit credential', () => {
+    it('edit button does NOT appear when no credential is selected', () => {
+        mockUseAsyncOptions.mockReturnValue({
+            ...idleResult(),
+            options: [{ label: 'My API Key', name: 'cred-1' }]
+        })
+
+        render(
+            <AsyncInput
+                inputParam={makeParam({ type: 'asyncOptions', credentialNames: ['openAIApi'] })}
+                value=''
+                disabled={false}
+                onChange={jest.fn()}
+            />
+        )
+
+        expect(screen.queryByTitle('Edit Credential')).toBeNull()
+    })
+
+    it('edit button appears when a credential is selected', () => {
+        mockUseAsyncOptions.mockReturnValue({
+            ...idleResult(),
+            options: [{ label: 'My API Key', name: 'cred-1' }]
+        })
+
+        render(
+            <AsyncInput
+                inputParam={makeParam({ type: 'asyncOptions', credentialNames: ['openAIApi'] })}
+                value='cred-1'
+                disabled={false}
+                onChange={jest.fn()}
+            />
+        )
+
+        expect(screen.getByTitle('Edit Credential')).toBeTruthy()
+    })
+
+    it('edit button does NOT appear for non-credential async dropdowns', () => {
+        mockUseAsyncOptions.mockReturnValue({
+            ...idleResult(),
+            options: [{ label: 'GPT-4o', name: 'gpt-4o' }]
+        })
+
+        render(<AsyncInput inputParam={makeParam({ type: 'asyncOptions' })} value='gpt-4o' disabled={false} onChange={jest.fn()} />)
+
+        expect(screen.queryByTitle('Edit Credential')).toBeNull()
+    })
+
+    it('clicking edit button opens the edit dialog', () => {
+        mockUseAsyncOptions.mockReturnValue({
+            ...idleResult(),
+            options: [{ label: 'My API Key', name: 'cred-1' }]
+        })
+
+        render(
+            <AsyncInput
+                inputParam={makeParam({ type: 'asyncOptions', credentialNames: ['openAIApi'] })}
+                value='cred-1'
+                disabled={false}
+                onChange={jest.fn()}
+            />
+        )
+
+        fireEvent.click(screen.getByTitle('Edit Credential'))
+        expect(screen.getByTestId('edit-credential-dialog')).toBeTruthy()
+    })
+
+    it('after editing, onChange is called with credential ID and component remounts to refetch', () => {
+        const mockChange = jest.fn()
+        mockUseAsyncOptions.mockReturnValue({
+            ...idleResult(),
+            options: [{ label: 'My API Key', name: 'cred-1' }]
+        })
+
+        render(
+            <AsyncInput
+                inputParam={makeParam({ type: 'asyncOptions', credentialNames: ['openAIApi'] })}
+                value='cred-1'
+                disabled={false}
+                onChange={mockChange}
+            />
+        )
+
+        const initialCallCount = mockUseAsyncOptions.mock.calls.length
+
+        fireEvent.click(screen.getByTitle('Edit Credential'))
+        fireEvent.click(screen.getByText('Save'))
+
+        expect(mockChange).toHaveBeenCalledWith('cred-1')
+        expect(mockUseAsyncOptions.mock.calls.length).toBeGreaterThan(initialCallCount)
+    })
+})
diff --git a/packages/agentflow/src/features/node-editor/AsyncInput.tsx b/packages/agentflow/src/features/node-editor/AsyncInput.tsx
@@ -2,7 +2,7 @@ import { Fragment, useCallback, useState } from 'react'
 
 import { Box, CircularProgress, IconButton, TextField, Typography } from '@mui/material'
 import Autocomplete from '@mui/material/Autocomplete'
-import { IconRefresh } from '@tabler/icons-react'
+import { IconEdit, IconRefresh } from '@tabler/icons-react'
 
 import type { AsyncInputProps } from '@/atoms'
 import type { NodeOption } from '@/core/types'
@@ -31,8 +31,11 @@ function buildAsyncParams(
 function AsyncOptionsInput({ inputParam, value, disabled, onChange, nodeName, inputValues }: AsyncInputProps) {
     const isCredential = !!inputParam.credentialNames?.length
     const [createDialogOpen, setCreateDialogOpen] = useState(false)
+    const [editDialogOpen, setEditDialogOpen] = useState(false)
     const [reloadKey, setReloadKey] = useState(0)
 
+    const selectedCredentialId = isCredential && typeof value === 'string' && value ? value : null
+
     const handleCreated = useCallback(
         (newCredentialId: string) => {
             setCreateDialogOpen(false)
@@ -45,19 +48,35 @@ function AsyncOptionsInput({ inputParam, value, disabled, onChange, nodeName, in
         [onChange]
     )
 
+    const handleEdited = useCallback(
+        (credentialId: string) => {
+            setEditDialogOpen(false)
+            onChange(credentialId)
+            setReloadKey((k) => k + 1)
+        },
+        [onChange]
+    )
+
     return (
         <>
-            <AsyncOptionsDropdown
-                key={reloadKey}
-                inputParam={inputParam}
-                value={value}
-                disabled={disabled}
-                onChange={onChange}
-                nodeName={nodeName}
-                inputValues={inputValues}
-                isCredential={isCredential}
-                onCreateNew={() => setCreateDialogOpen(true)}
-            />
+            <Box sx={{ display: 'flex', flexDirection: 'row', alignItems: 'center', mt: 1 }}>
+                <AsyncOptionsDropdown
+                    key={reloadKey}
+                    inputParam={inputParam}
+                    value={value}
+                    disabled={disabled}
+                    onChange={onChange}
+                    nodeName={nodeName}
+                    inputValues={inputValues}
+                    isCredential={isCredential}
+                    onCreateNew={() => setCreateDialogOpen(true)}
+                />
+                {selectedCredentialId && (
+                    <IconButton title='Edit Credential' color='primary' size='small' onClick={() => setEditDialogOpen(true)}>
+                        <IconEdit size={18} />
+                    </IconButton>
+                )}
+            </Box>
             {isCredential && (
                 <CreateCredentialDialog
                     open={createDialogOpen}
@@ -66,6 +85,15 @@ function AsyncOptionsInput({ inputParam, value, disabled, onChange, nodeName, in
                     onCreated={handleCreated}
                 />
             )}
+            {isCredential && selectedCredentialId && (
+                <CreateCredentialDialog
+                    open={editDialogOpen}
+                    credentialNames={inputParam.credentialNames!}
+                    onClose={() => setEditDialogOpen(false)}
+                    onCreated={handleEdited}
+                    editCredentialId={selectedCredentialId}
+                />
+            )}
         </>
     )
 }
@@ -123,7 +151,7 @@ function AsyncOptionsDropdown({
             }}
             loading={loading}
             noOptionsText={loading ? 'Loading…' : 'No options available'}
-            sx={{ mt: 1 }}
+            sx={{ flexGrow: 1 }}
             renderOption={(props, option) => (
                 <Box component='li' {...props} sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
                     {option.name === CREATE_NEW_SENTINEL ? (
diff --git a/packages/agentflow/src/features/node-editor/CreateCredentialDialog.test.tsx b/packages/agentflow/src/features/node-editor/CreateCredentialDialog.test.tsx
@@ -54,11 +54,15 @@ jest.mock('@/atoms/Dropdown', () => ({
 
 const mockGetComponentCredentialSchema = jest.fn()
 const mockCreateCredential = jest.fn()
+const mockGetCredentialById = jest.fn()
+const mockUpdateCredential = jest.fn()
 
 const mockApiContext = {
     credentialsApi: {
         getComponentCredentialSchema: mockGetComponentCredentialSchema,
-        createCredential: mockCreateCredential
+        createCredential: mockCreateCredential,
+        getCredentialById: mockGetCredentialById,
+        updateCredential: mockUpdateCredential
     },
     apiBaseUrl: 'http://localhost:3000'
 }
diff --git a/packages/agentflow/src/features/node-editor/CreateCredentialDialog.tsx b/packages/agentflow/src/features/node-editor/CreateCredentialDialog.tsx
@@ -30,15 +30,18 @@ export interface CreateCredentialDialogProps {
     credentialNames: string[]
     onClose: () => void
     onCreated: (credentialId: string) => void
+    /** When set, the dialog opens in edit mode for the given credential ID. */
+    editCredentialId?: string
 }
 
 /**
- * Dialog for creating a new credential from within the node editor.
+ * Dialog for creating or editing a credential from within the node editor.
  * Fetches the credential schema from the backend and renders a dynamic form.
  */
-export function CreateCredentialDialog({ open, credentialNames, onClose, onCreated }: CreateCredentialDialogProps) {
+export function CreateCredentialDialog({ open, credentialNames, onClose, onCreated, editCredentialId }: CreateCredentialDialogProps) {
     const { credentialsApi, apiBaseUrl } = useApiContext()
     const theme = useTheme()
+    const isEditMode = !!editCredentialId
 
     const [schemas, setSchemas] = useState<ComponentCredentialSchema[]>([])
     const [selectedSchema, setSelectedSchema] = useState<ComponentCredentialSchema | null>(null)
@@ -65,7 +68,7 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
         setFormValues(defaults)
     }, [])
 
-    // Fetch credential schema(s) when dialog opens
+    // Fetch credential schema(s) when dialog opens, and load existing data in edit mode
     useEffect(() => {
         if (!open) return
 
@@ -95,6 +98,17 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
                         setSchemas(results)
                     }
                 }
+
+                // In edit mode, fetch existing credential and populate the form
+                if (editCredentialId && !cancelled) {
+                    const existing = await credentialsApi.getCredentialById(editCredentialId)
+                    if (!cancelled) {
+                        setCredentialName(existing.name)
+                        if (existing.plainDataObj) {
+                            setFormValues(existing.plainDataObj)
+                        }
+                    }
+                }
             } catch (err) {
                 if (!cancelled) {
                     setError(err instanceof Error ? err.message : 'Failed to load credential schema')
@@ -111,7 +125,7 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
         return () => {
             cancelled = true
         }
-    }, [open, credentialNamesKey, credentialsApi, selectSchema])
+    }, [open, credentialNamesKey, credentialsApi, selectSchema, editCredentialId])
 
     const handleFieldChange = useCallback((fieldName: string, value: unknown) => {
         setFormValues((prev) => ({ ...prev, [fieldName]: value }))
@@ -135,18 +149,21 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
                 }
             }
 
-            const result = await credentialsApi.createCredential({
+            const body = {
                 name: credentialName.trim(),
                 credentialName: selectedSchema.name,
                 plainDataObj
-            })
+            }
+            const result = isEditMode
+                ? await credentialsApi.updateCredential(editCredentialId!, body)
+                : await credentialsApi.createCredential(body)
             onCreated(result.id)
         } catch (err) {
-            setError(err instanceof Error ? err.message : 'Failed to create credential')
+            setError(err instanceof Error ? err.message : isEditMode ? 'Failed to update credential' : 'Failed to create credential')
         } finally {
             setSubmitting(false)
         }
-    }, [selectedSchema, credentialName, formValues, credentialsApi, onCreated])
+    }, [selectedSchema, credentialName, formValues, credentialsApi, onCreated, isEditMode, editCredentialId])
 
     const handleClose = useCallback(() => {
         if (!submitting) onClose()
@@ -224,13 +241,15 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
                                         display: 'flex',
                                         flexDirection: 'row',
                                         borderRadius: 10,
-                                        background: theme.palette.warningBanner.background,
+                                        background: theme.palette.warningBanner?.background ?? '#fefcbf',
                                         padding: 10,
                                         marginTop: 10,
                                         marginBottom: 10
                                     }}
                                 >
-                                    <span style={{ color: theme.palette.warningBanner.text }}>{parser(selectedSchema.description)}</span>
+                                    <span style={{ color: theme.palette.warningBanner?.text ?? '#744210' }}>
+                                        {parser(selectedSchema.description)}
+                                    </span>
                                 </div>
                             </Box>
                         )}
@@ -269,7 +288,7 @@ export function CreateCredentialDialog({ open, credentialNames, onClose, onCreat
                 </Button>
                 {selectedSchema && (
                     <Button variant='contained' onClick={handleSubmit} disabled={!credentialName.trim() || submitting}>
-                        {submitting ? 'Adding...' : 'Add'}
+                        {submitting ? (isEditMode ? 'Saving...' : 'Adding...') : isEditMode ? 'Save' : 'Add'}
                     </Button>
                 )}
             </DialogActions>
@@ -323,14 +342,14 @@ function CredentialField({ input, value, onChange, disabled = false }: Credentia
                         display: 'flex',
                         flexDirection: 'row',
                         borderRadius: 10,
-                        background: theme.palette.warningBanner.background,
+                        background: theme.palette.warningBanner?.background ?? '#fefcbf',
                         padding: 10,
                         marginTop: 10,
                         marginBottom: 10
                     }}
                 >
                     <IconAlertTriangle size={36} color={theme.palette.warning.main} />
-                    <span style={{ color: theme.palette.warningBanner.text, marginLeft: 10 }}>{input.warning}</span>
+                    <span style={{ color: theme.palette.warningBanner?.text ?? '#744210', marginLeft: 10 }}>{input.warning}</span>
                 </div>
             )}
 
diff --git a/packages/agentflow/src/infrastructure/api/credentials.test.ts b/packages/agentflow/src/infrastructure/api/credentials.test.ts
@@ -4,7 +4,8 @@ import { bindCredentialsApi } from './credentials'
 
 const mockClient = {
     get: jest.fn(),
-    post: jest.fn()
+    post: jest.fn(),
+    put: jest.fn()
 } as unknown as jest.Mocked<AxiosInstance>
 
 beforeEach(() => {
@@ -50,4 +51,23 @@ describe('bindCredentialsApi', () => {
         expect(mockClient.post).toHaveBeenCalledWith('/credentials', body)
         expect(result).toEqual(mockCreated)
     })
+
+    it('getCredentialById calls GET /credentials/:id', async () => {
+        const mockCredential = { id: 'cred-1', name: 'My Key', credentialName: 'openAIApi', plainDataObj: { apiKey: 'sk-test' } }
+        ;(mockClient.get as jest.Mock).mockResolvedValue({ data: mockCredential })
+
+        const result = await api.getCredentialById('cred-1')
+        expect(mockClient.get).toHaveBeenCalledWith('/credentials/cred-1')
+        expect(result).toEqual(mockCredential)
+    })
+
+    it('updateCredential calls PUT /credentials/:id with body', async () => {
+        const body = { name: 'Updated Key', credentialName: 'openAIApi', plainDataObj: { apiKey: 'sk-new' } }
+        const mockUpdated = { id: 'cred-1', ...body }
+        ;(mockClient.put as jest.Mock).mockResolvedValue({ data: mockUpdated })
+
+        const result = await api.updateCredential('cred-1', body)
+        expect(mockClient.put).toHaveBeenCalledWith('/credentials/cred-1', body)
+        expect(result).toEqual(mockUpdated)
+    })
 })
diff --git a/packages/agentflow/src/infrastructure/api/credentials.ts b/packages/agentflow/src/infrastructure/api/credentials.ts
diff --git a/packages/agentflow/src/infrastructure/api/loadMethodRegistry.test.ts b/packages/agentflow/src/infrastructure/api/loadMethodRegistry.test.ts
