Skip to content

Commit 07e7604

Browse files
j-sanaaj-sanaa
authored
Feat/agentflow Support for Additional Load Methods (#6003)
* Generate itemKeys synchronously * Support for listActions and listTables * Have generic fallback method for loadMethods which are not registered * Fix comments * Fix text error
1 parent 787a7bb commit 07e7604

9 files changed

Lines changed: 339 additions & 93 deletions

File tree

‎packages/agentflow/src/atoms/ArrayInput.tsx‎

Lines changed: 5 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { type ComponentType, useCallback, useMemo, useRef, useState } from 'react'
1+
import { type ComponentType, useCallback, useMemo } from 'react'
22

33
import { Box, Button, Chip, IconButton } from '@mui/material'
44
import { useTheme } from '@mui/material/styles'
@@ -7,6 +7,7 @@ import { IconPlus, IconTrash } from '@tabler/icons-react'
77
import type { InputParam, NodeData } from '@/core/types'
88

99
import { type AsyncInputProps, type ConfigInputComponentProps, NodeInputHandler } from './NodeInputHandler'
10+
import { useStableKeys } from './useStableKeys'
1011

1112
export interface ArrayInputProps {
1213
inputParam: InputParam
@@ -42,33 +43,7 @@ export function ArrayInput({
4243
[data.inputValues, inputParam.name]
4344
)
4445

45-
// Stable keys for array items — avoids using index as React key.
46-
// Keys are held in state so they persist across renders. A local variable
47-
// (`effectiveKeys`) is used for the current render pass so that newly
48-
// generated keys are available immediately (setState alone would only
49-
// take effect on the *next* render, leaving keys undefined in this one).
50-
const idCounterRef = useRef(0)
51-
const [itemKeys, setItemKeys] = useState<string[]>(() => {
52-
const initial: string[] = []
53-
while (initial.length < arrayItems.length) {
54-
initial.push(`item-${idCounterRef.current++}`)
55-
}
56-
return initial
57-
})
58-
59-
let effectiveKeys = itemKeys
60-
if (effectiveKeys.length < arrayItems.length) {
61-
const nextKeys = [...effectiveKeys]
62-
while (nextKeys.length < arrayItems.length) {
63-
nextKeys.push(`item-${idCounterRef.current++}`)
64-
}
65-
setItemKeys(nextKeys)
66-
effectiveKeys = nextKeys
67-
} else if (effectiveKeys.length > arrayItems.length) {
68-
const trimmed = effectiveKeys.slice(0, arrayItems.length)
69-
setItemKeys(trimmed)
70-
effectiveKeys = trimmed
71-
}
46+
const { keys: effectiveKeys, removeKey } = useStableKeys(arrayItems.length, 'item')
7247

7348
// Use pre-computed itemParameters
7449
// Falls back to raw field definitions for nested arrays without show/hide conditions.
@@ -130,12 +105,12 @@ export function ArrayInput({
130105
const handleDeleteItem = useCallback(
131106
(indexToDelete: number) => {
132107
const updatedArrayItems = arrayItems.filter((_, i) => i !== indexToDelete)
133-
setItemKeys((prev) => prev.filter((_, i) => i !== indexToDelete))
108+
removeKey(indexToDelete)
134109

135110
// Notify parent of change (parent will update props, causing re-render)
136111
onDataChange?.({ inputParam, newValue: updatedArrayItems })
137112
},
138-
[arrayItems, inputParam, onDataChange]
113+
[arrayItems, inputParam, onDataChange, removeKey]
139114
)
140115

141116
// Pre-compute stable per-item onDataChange handlers to avoid new closures on every render

‎packages/agentflow/src/atoms/ConditionBuilder.tsx‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useCallback, useEffect, useMemo, useRef } from 'react'
1+
import { useCallback, useMemo } from 'react'
22

33
import { Box, Button, Chip, IconButton, Typography } from '@mui/material'
44
import { useTheme } from '@mui/material/styles'
@@ -7,6 +7,7 @@ import { IconPlus, IconTrash } from '@tabler/icons-react'
77
import type { InputParam, NodeData } from '@/core/types'
88

99
import { NodeInputHandler } from './NodeInputHandler'
10+
import { useStableKeys } from './useStableKeys'
1011

1112
export interface ConditionBuilderProps {
1213
inputParam: InputParam
@@ -29,20 +30,13 @@ export function ConditionBuilder({
2930
itemParameters: itemParametersProp
3031
}: ConditionBuilderProps) {
3132
const theme = useTheme()
32-
const idCounterRef = useRef(0)
33-
const itemKeysRef = useRef<string[]>([])
3433

3534
const arrayItems = useMemo(
3635
() => (Array.isArray(data.inputValues?.[inputParam.name]) ? (data.inputValues[inputParam.name] as Record<string, unknown>[]) : []),
3736
[data.inputValues, inputParam.name]
3837
)
3938

40-
// Grow keys array when new items appear (e.g. on mount or external data changes)
41-
useEffect(() => {
42-
while (itemKeysRef.current.length < arrayItems.length) {
43-
itemKeysRef.current.push(`condition-${idCounterRef.current++}`)
44-
}
45-
}, [arrayItems.length])
39+
const { keys: effectiveKeys, removeKey } = useStableKeys(arrayItems.length, 'condition')
4640

4741
const itemParameters = useMemo<InputParam[][]>(
4842
() => itemParametersProp ?? arrayItems.map(() => inputParam.array || []),
@@ -85,10 +79,10 @@ export function ConditionBuilder({
8579

8680
const handleDeleteItem = useCallback(
8781
(indexToDelete: number) => {
88-
itemKeysRef.current.splice(indexToDelete, 1)
82+
removeKey(indexToDelete)
8983
onDataChange?.({ inputParam, newValue: arrayItems.filter((_, i) => i !== indexToDelete) })
9084
},
91-
[arrayItems, inputParam, onDataChange]
85+
[arrayItems, inputParam, onDataChange, removeKey]
9286
)
9387

9488
const itemHandlers = useMemo(
@@ -111,7 +105,7 @@ export function ConditionBuilder({
111105

112106
return (
113107
<Box
114-
key={itemKeysRef.current[index]}
108+
key={effectiveKeys[index]}
115109
sx={{
116110
p: 2,
117111
mt: 2,

‎packages/agentflow/src/atoms/MessagesInput.tsx‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import type { InputParam, NodeData } from '@/core/types'
88

99
import { ExpandTextDialog } from './ExpandTextDialog'
1010
import { RichTextEditor } from './RichTextEditor.lazy'
11+
import { useStableKeys } from './useStableKeys'
1112

1213
const MESSAGE_ROLES = [
1314
{ label: 'System', value: 'system' },
@@ -37,18 +38,13 @@ export interface MessagesInputProps {
3738
*/
3839
export function MessagesInput({ inputParam, data, disabled = false, onDataChange }: MessagesInputProps) {
3940
const theme = useTheme()
40-
const idCounterRef = useRef(0)
41-
const itemKeysRef = useRef<string[]>([])
4241

4342
const messages = useMemo(
4443
() => (Array.isArray(data.inputValues?.[inputParam.name]) ? (data.inputValues[inputParam.name] as MessageEntry[]) : []),
4544
[data.inputValues, inputParam.name]
4645
)
4746

48-
// Grow keys array synchronously so keys are available on the first render
49-
while (itemKeysRef.current.length < messages.length) {
50-
itemKeysRef.current.push(`message-${idCounterRef.current++}`)
51-
}
47+
const { keys: effectiveKeys, removeKey } = useStableKeys(messages.length, 'message')
5248

5349
const handleRoleChange = useCallback(
5450
(index: number, role: string) => {
@@ -61,17 +57,17 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
6157

6258
// Track latest inline content locally so the expand dialog always has fresh values,
6359
// even if the parent hasn't round-tripped onDataChange back into data yet.
64-
// Keyed by itemKeysRef values so deletes are a simple Map.delete() with no index rebasing.
60+
// Keyed by item key values so deletes are a simple Map.delete() with no index rebasing.
6561
const latestContentRef = useRef<Map<string, string>>(new Map())
6662

6763
const handleContentChange = useCallback(
6864
(index: number, content: string) => {
69-
latestContentRef.current.set(itemKeysRef.current[index], content)
65+
latestContentRef.current.set(effectiveKeys[index], content)
7066
const updated = [...messages]
7167
updated[index] = { ...updated[index], content }
7268
onDataChange?.({ inputParam, newValue: updated })
7369
},
74-
[messages, inputParam, onDataChange]
70+
[effectiveKeys, messages, inputParam, onDataChange]
7571
)
7672

7773
const handleAddMessage = useCallback(() => {
@@ -81,11 +77,11 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
8177

8278
const handleDeleteMessage = useCallback(
8379
(indexToDelete: number) => {
84-
latestContentRef.current.delete(itemKeysRef.current[indexToDelete])
85-
itemKeysRef.current.splice(indexToDelete, 1)
80+
latestContentRef.current.delete(effectiveKeys[indexToDelete])
81+
removeKey(indexToDelete)
8682
onDataChange?.({ inputParam, newValue: messages.filter((_, i) => i !== indexToDelete) })
8783
},
88-
[messages, inputParam, onDataChange]
84+
[effectiveKeys, messages, inputParam, onDataChange, removeKey]
8985
)
9086

9187
// Expand dialog state
@@ -98,12 +94,12 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
9894
const handleExpandConfirm = useCallback(
9995
(value: string) => {
10096
if (expandIndex !== null) {
101-
latestContentRef.current.set(itemKeysRef.current[expandIndex], value)
97+
latestContentRef.current.set(effectiveKeys[expandIndex], value)
10298
handleContentChange(expandIndex, value)
10399
}
104100
setExpandIndex(null)
105101
},
106-
[expandIndex, handleContentChange]
102+
[effectiveKeys, expandIndex, handleContentChange]
107103
)
108104

109105
const handleExpandCancel = useCallback(() => {
@@ -122,7 +118,7 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
122118

123119
{messages.map((message, index) => (
124120
<Box
125-
key={itemKeysRef.current[index]}
121+
key={effectiveKeys[index]}
126122
sx={{
127123
p: 2,
128124
mt: 2,
@@ -232,7 +228,7 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
232228
{expandIndex !== null && (
233229
<ExpandTextDialog
234230
open={true}
235-
value={latestContentRef.current.get(itemKeysRef.current[expandIndex]) ?? messages[expandIndex]?.content ?? ''}
231+
value={latestContentRef.current.get(effectiveKeys[expandIndex]) ?? messages[expandIndex]?.content ?? ''}
236232
title='Content'
237233
placeholder='Message content (supports {{ variable }} syntax)'
238234
disabled={disabled}

‎packages/agentflow/src/atoms/StructuredOutputBuilder.tsx‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
1+
import { useCallback, useMemo, useState } from 'react'
22

33
import { Box, Button, Chip, IconButton, MenuItem, Select, TextField, Tooltip, Typography } from '@mui/material'
44
import { useTheme } from '@mui/material/styles'
@@ -7,6 +7,8 @@ import { IconArrowsMaximize, IconInfoCircle, IconPlus, IconTrash } from '@tabler
77
import { ExpandTextDialog } from '@/atoms'
88
import type { InputParam, NodeData } from '@/core/types'
99

10+
import { useStableKeys } from './useStableKeys'
11+
1012
const OUTPUT_TYPES = [
1113
{ label: 'String', value: 'string' },
1214
{ label: 'String Array', value: 'stringArray' },
@@ -40,20 +42,13 @@ export interface StructuredOutputBuilderProps {
4042
*/
4143
export function StructuredOutputBuilder({ inputParam, data, disabled = false, onDataChange }: StructuredOutputBuilderProps) {
4244
const theme = useTheme()
43-
const idCounterRef = useRef(0)
44-
const itemKeysRef = useRef<string[]>([])
4545

4646
const entries = useMemo(
4747
() => (Array.isArray(data.inputValues?.[inputParam.name]) ? (data.inputValues[inputParam.name] as StructuredOutputEntry[]) : []),
4848
[data.inputValues, inputParam.name]
4949
)
5050

51-
// Grow keys array when new items appear (e.g. on mount or external data changes)
52-
useEffect(() => {
53-
while (itemKeysRef.current.length < entries.length) {
54-
itemKeysRef.current.push(`output-${idCounterRef.current++}`)
55-
}
56-
}, [entries.length])
51+
const { keys: effectiveKeys, removeKey } = useStableKeys(entries.length, 'output')
5752

5853
const handleFieldChange = useCallback(
5954
(index: number, field: string, value: string) => {
@@ -79,10 +74,10 @@ export function StructuredOutputBuilder({ inputParam, data, disabled = false, on
7974

8075
const handleDeleteEntry = useCallback(
8176
(indexToDelete: number) => {
82-
itemKeysRef.current.splice(indexToDelete, 1)
77+
removeKey(indexToDelete)
8378
onDataChange?.({ inputParam, newValue: entries.filter((_, i) => i !== indexToDelete) })
8479
},
85-
[entries, inputParam, onDataChange]
80+
[entries, inputParam, onDataChange, removeKey]
8681
)
8782

8883
const isDeleteVisible = !inputParam.minItems || entries.length > inputParam.minItems
@@ -100,7 +95,7 @@ export function StructuredOutputBuilder({ inputParam, data, disabled = false, on
10095

10196
{entries.map((entry, index) => (
10297
<Box
103-
key={itemKeysRef.current[index]}
98+
key={effectiveKeys[index]}
10499
sx={{
105100
p: 2,
106101
mt: 2,

0 commit comments

Comments
 (0)