Skip to content

Commit fedc6eb

Browse files
jocelynlin-wdclaude
jocelynlin-wd
and
claude
authored
feat(agentflow): enhance state management with StateKeyValueInput and variable support (#6119)
* feat(agentflow): enhance state management with StateKeyValueInput and variable support - Added StateKeyValueInput component for managing key-value pairs in flow state updates. - Integrated VariableInput for value fields to support variable syntax. - Updated MessagesInput and ExpandTextDialog to utilize variable items for enhanced user experience. - Introduced new types for state updates in core types. - Enhanced Jest tests to cover new functionality and ensure reliability. This update improves the flexibility and usability of state management within the agentflow package. * fix(agentflow): address PR #6119 review feedback - Add autoFocus prop to VariableInput, apply in ExpandTextDialog for consistent focus behavior with RichTextEditor fallback - Extract duplicated VariableItem→SuggestionItem mapping into shared toSuggestionItems utility, used by MessagesInput and NodeInputHandler - Remove redundant default export from StateKeyValueInput Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * clear selected key and update dropdown when a state key is removed * cleanup unused type --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 49a2259 commit fedc6eb

29 files changed

Lines changed: 868 additions & 117 deletions

‎packages/agentflow/jest.config.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ module.exports = {
4949
'./src/atoms/ExpandTextDialog.tsx': { branches: 80, functions: 80, lines: 80, statements: 80 },
5050
'./src/atoms/MessagesInput.tsx': { branches: 80, functions: 80, lines: 80, statements: 80 },
5151
'./src/atoms/ScenariosInput.tsx': { branches: 80, functions: 80, lines: 80, statements: 80 },
52+
'./src/atoms/StateKeyValueInput.tsx': { branches: 80, functions: 80, lines: 80, statements: 80 },
5253
// Tier 3 UI atom — only the onChange/disabled/sync logic is tested, not styled internals
5354
'./src/atoms/RichTextEditor.tsx': { branches: 30, functions: 50, lines: 50, statements: 50 },
5455
'./src/core/': { branches: 80, functions: 80, lines: 80, statements: 80 },

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type { InputParam, NodeData } from '@/core/types'
99

1010
import { type AsyncInputProps, type ConfigInputComponentProps, NodeInputHandler } from './NodeInputHandler'
1111
import { useStableKeys } from './useStableKeys'
12+
import type { VariableItem } from './VariablePicker'
1213

1314
export interface ArrayInputProps {
1415
inputParam: InputParam
@@ -23,6 +24,8 @@ export interface ArrayInputProps {
2324
configValues: Record<string, unknown>,
2425
arrayContext?: { parentParamName: string; arrayIndex: number }
2526
) => void
27+
/** Variable items passed through to sub-field NodeInputHandlers for {{ autocomplete. */
28+
variableItems?: VariableItem[]
2629
}
2730

2831
export function ArrayInput({
@@ -33,7 +36,8 @@ export function ArrayInput({
3336
itemParameters: itemParametersProp,
3437
AsyncInputComponent,
3538
ConfigInputComponent,
36-
onConfigChange
39+
onConfigChange,
40+
variableItems
3741
}: ArrayInputProps) {
3842
const theme = useTheme()
3943

@@ -174,6 +178,7 @@ export function ArrayInput({
174178
onConfigChange={onConfigChange}
175179
arrayIndex={index}
176180
parentArrayParam={inputParam}
181+
variableItems={param.acceptVariable ? variableItems : undefined}
177182
/>
178183
))}
179184
</Box>

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

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import { Box, Button, Dialog, DialogActions, DialogContent, TextField, Typograph
44

55
import { CodeInput } from './CodeInput'
66
import { RichTextEditor } from './RichTextEditor.lazy'
7+
import type { SuggestionItem } from './SuggestionDropdown'
8+
import { VariableInput } from './VariableInput'
79

810
export interface ExpandTextDialogProps {
911
open: boolean
@@ -15,6 +17,8 @@ export interface ExpandTextDialogProps {
1517
inputType?: string
1618
/** Language hint for 'code' mode (e.g. 'javascript', 'python', 'json'). */
1719
language?: string
20+
/** Variable suggestion items for `{{ }}` autocomplete in string mode. When provided, uses VariableInput instead of RichTextEditor. */
21+
suggestionItems?: SuggestionItem[]
1822
onConfirm: (value: string) => void
1923
onCancel: () => void
2024
}
@@ -31,6 +35,7 @@ export function ExpandTextDialog({
3135
disabled = false,
3236
inputType = 'string',
3337
language,
38+
suggestionItems,
3439
onConfirm,
3540
onCancel
3641
}: ExpandTextDialogProps) {
@@ -76,14 +81,26 @@ export function ExpandTextDialog({
7681
overflowX: 'hidden'
7782
}}
7883
>
79-
<RichTextEditor
80-
value={localValue}
81-
onChange={setLocalValue}
82-
placeholder={placeholder}
83-
disabled={disabled}
84-
rows={15}
85-
autoFocus
86-
/>
84+
{suggestionItems && suggestionItems.length > 0 ? (
85+
<VariableInput
86+
value={localValue}
87+
onChange={setLocalValue}
88+
placeholder={placeholder}
89+
disabled={disabled}
90+
rows={15}
91+
suggestionItems={suggestionItems}
92+
autoFocus
93+
/>
94+
) : (
95+
<RichTextEditor
96+
value={localValue}
97+
onChange={setLocalValue}
98+
placeholder={placeholder}
99+
disabled={disabled}
100+
rows={15}
101+
autoFocus
102+
/>
103+
)}
87104
</Box>
88105
) : (
89106
<TextField

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,30 @@ jest.mock('./RichTextEditor', () => ({
4040
)
4141
}))
4242

43+
jest.mock('./VariableInput', () => ({
44+
VariableInput: ({
45+
value,
46+
onChange,
47+
disabled,
48+
placeholder
49+
}: {
50+
value: string
51+
onChange: (html: string) => void
52+
disabled?: boolean
53+
placeholder?: string
54+
}) => (
55+
<div data-testid='variable-input'>
56+
<textarea
57+
data-testid='variable-input-content'
58+
value={value}
59+
disabled={disabled}
60+
placeholder={placeholder}
61+
onChange={(e) => onChange(e.target.value)}
62+
/>
63+
</div>
64+
)
65+
}))
66+
4367
describe('MessagesInput', () => {
4468
const mockInputParam: InputParam = {
4569
id: 'messages',
@@ -503,4 +527,72 @@ describe('MessagesInput', () => {
503527
const optionValues = options.map((opt) => opt.getAttribute('data-value'))
504528
expect(optionValues).toEqual(['system', 'assistant', 'developer', 'user'])
505529
})
530+
531+
// --- Variable support ---
532+
533+
const mockVariableItems = [
534+
{ label: 'question', description: "User's question", category: 'Chat Context', value: '{{question}}' },
535+
{ label: '$flow.sessionId', description: 'Session ID', category: 'Flow Variables', value: '{{$flow.sessionId}}' }
536+
]
537+
538+
it('should render VariableInput instead of RichTextEditor when variableItems are provided', async () => {
539+
const dataWithMessages: NodeData = {
540+
...mockNodeData,
541+
inputs: {
542+
agentMessages: [{ role: 'user', content: 'Hello' }]
543+
}
544+
} as NodeData
545+
546+
render(
547+
<MessagesInput
548+
inputParam={mockInputParam}
549+
data={dataWithMessages}
550+
onDataChange={mockOnDataChange}
551+
variableItems={mockVariableItems}
552+
/>
553+
)
554+
555+
expect(await screen.findByTestId('variable-input')).toBeInTheDocument()
556+
expect(screen.queryByTestId('rich-text-editor')).not.toBeInTheDocument()
557+
})
558+
559+
it('should fall back to RichTextEditor when variableItems is empty', async () => {
560+
const dataWithMessages: NodeData = {
561+
...mockNodeData,
562+
inputs: {
563+
agentMessages: [{ role: 'user', content: 'Hello' }]
564+
}
565+
} as NodeData
566+
567+
render(<MessagesInput inputParam={mockInputParam} data={dataWithMessages} onDataChange={mockOnDataChange} variableItems={[]} />)
568+
569+
expect(await screen.findByTestId('rich-text-editor')).toBeInTheDocument()
570+
expect(screen.queryByTestId('variable-input')).not.toBeInTheDocument()
571+
})
572+
573+
it('should update content via VariableInput when variableItems are provided', async () => {
574+
const dataWithMessages: NodeData = {
575+
...mockNodeData,
576+
inputs: {
577+
agentMessages: [{ role: 'user', content: '' }]
578+
}
579+
} as NodeData
580+
581+
render(
582+
<MessagesInput
583+
inputParam={mockInputParam}
584+
data={dataWithMessages}
585+
onDataChange={mockOnDataChange}
586+
variableItems={mockVariableItems}
587+
/>
588+
)
589+
590+
const textarea = await screen.findByTestId('variable-input-content')
591+
fireEvent.change(textarea, { target: { value: '{{ question }}' } })
592+
593+
expect(mockOnDataChange).toHaveBeenCalledWith({
594+
inputParam: mockInputParam,
595+
newValue: [{ role: 'user', content: '{{ question }}' }]
596+
})
597+
})
506598
})

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

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

99
import { ExpandTextDialog } from './ExpandTextDialog'
1010
import { RichTextEditor } from './RichTextEditor.lazy'
11+
import { toSuggestionItems } from './toSuggestionItems'
1112
import { useStableKeys } from './useStableKeys'
13+
import { VariableInput } from './VariableInput'
14+
import type { VariableItem } from './VariablePicker'
1215

1316
const MESSAGE_ROLES = [
1417
{ label: 'System', value: 'system' },
@@ -29,14 +32,16 @@ export interface MessagesInputProps {
2932
data: NodeData
3033
disabled?: boolean
3134
onDataChange?: (params: { inputParam: InputParam; newValue: unknown }) => void
35+
/** Variable items for {{ autocomplete in message content fields. */
36+
variableItems?: VariableItem[]
3237
}
3338

3439
/**
3540
* Specialized array input for message entries (Agent + LLM nodes).
3641
* Each entry has a role dropdown (system/assistant/developer/user)
3742
* and a multiline content textarea with variable support ({{ variable }} syntax).
3843
*/
39-
export function MessagesInput({ inputParam, data, disabled = false, onDataChange }: MessagesInputProps) {
44+
export function MessagesInput({ inputParam, data, disabled = false, onDataChange, variableItems }: MessagesInputProps) {
4045
const theme = useTheme()
4146

4247
const messages = useMemo(
@@ -46,6 +51,8 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
4651

4752
const { keys: effectiveKeys, removeKey } = useStableKeys(messages.length, 'message')
4853

54+
const suggestionItems = useMemo(() => toSuggestionItems(variableItems), [variableItems])
55+
4956
const handleRoleChange = useCallback(
5057
(index: number, role: string) => {
5158
const updated = [...messages]
@@ -200,13 +207,24 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
200207
<IconArrowsMaximize />
201208
</IconButton>
202209
</div>
203-
<RichTextEditor
204-
value={message.content}
205-
onChange={(html) => handleContentChange(index, html)}
206-
placeholder='Message content (supports {{ variable }} syntax)'
207-
disabled={disabled}
208-
rows={4}
209-
/>
210+
{suggestionItems && suggestionItems.length > 0 ? (
211+
<VariableInput
212+
value={message.content}
213+
onChange={(html) => handleContentChange(index, html)}
214+
placeholder='Message content (supports {{ variable }} syntax)'
215+
disabled={disabled}
216+
rows={4}
217+
suggestionItems={suggestionItems}
218+
/>
219+
) : (
220+
<RichTextEditor
221+
value={message.content}
222+
onChange={(html) => handleContentChange(index, html)}
223+
placeholder='Message content (supports {{ variable }} syntax)'
224+
disabled={disabled}
225+
rows={4}
226+
/>
227+
)}
210228
</Box>
211229
</Box>
212230
))}
@@ -233,6 +251,7 @@ export function MessagesInput({ inputParam, data, disabled = false, onDataChange
233251
placeholder='Message content (supports {{ variable }} syntax)'
234252
disabled={disabled}
235253
inputType='string'
254+
suggestionItems={suggestionItems}
236255
onConfirm={handleExpandConfirm}
237256
onCancel={handleExpandCancel}
238257
/>

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

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,18 @@ import { styled, useTheme } from '@mui/material/styles'
1919
import { tooltipClasses } from '@mui/material/Tooltip'
2020
import { IconArrowsMaximize, IconVariable } from '@tabler/icons-react'
2121

22-
import type { InputAnchor, InputParam, NodeData } from '@/core/types'
22+
import type { InputAnchor, InputParam, NodeData, StateUpdate } from '@/core/types'
2323

2424
import ArrayInput from './ArrayInput'
2525
import { CodeInput } from './CodeInput'
2626
import { Dropdown } from './Dropdown'
2727
import { ExpandTextDialog } from './ExpandTextDialog'
2828
import { JsonInput } from './JsonInput'
2929
import { RichTextEditor } from './RichTextEditor.lazy'
30-
import { SuggestionItem } from './SuggestionDropdown'
30+
import { StateKeyValueInput } from './StateKeyValueInput'
3131
import { SwitchInput } from './SwitchInput'
3232
import { TooltipWithParser } from './TooltipWithParser'
33+
import { toSuggestionItems } from './toSuggestionItems'
3334
import { VariableInput } from './VariableInput'
3435
import type { VariableItem } from './VariablePicker'
3536
import { VariablePicker } from './VariablePicker'
@@ -180,24 +181,10 @@ export function NodeInputHandler({
180181
['string', 'password', 'code'].includes(inputParam?.type ?? '')
181182
)
182183

183-
// Map VariableItem[] to SuggestionItem[] for the inline autocomplete.
184-
// ids must be unique for correct findIndex lookups and React keys — append
185-
// a counter suffix when the same base id appears more than once.
186-
const suggestionItems: SuggestionItem[] | undefined = useMemo(() => {
187-
if (!inputParam?.acceptVariable || !variableItems || variableItems.length === 0) return undefined
188-
const idCount = new Map<string, number>()
189-
return variableItems.map((v) => {
190-
const baseId = v.value.replace(/{{|}}/g, '')
191-
const count = idCount.get(baseId) ?? 0
192-
idCount.set(baseId, count + 1)
193-
return {
194-
id: count === 0 ? baseId : `${baseId}__${count}`,
195-
label: v.label,
196-
description: v.description,
197-
category: v.category
198-
}
199-
})
200-
}, [inputParam?.acceptVariable, variableItems])
184+
const suggestionItems = useMemo(
185+
() => (inputParam?.acceptVariable ? toSuggestionItems(variableItems) : undefined),
186+
[inputParam?.acceptVariable, variableItems]
187+
)
201188

202189
const renderInput = () => {
203190
if (!inputParam) return null
@@ -355,6 +342,16 @@ export function NodeInputHandler({
355342
</>
356343
)
357344

345+
case 'updateFlowState':
346+
return (
347+
<StateKeyValueInput
348+
value={Array.isArray(value) ? (value as StateUpdate[]) : []}
349+
onChange={(v) => handleDataChange(v)}
350+
disabled={disabled}
351+
suggestionItems={suggestionItems}
352+
/>
353+
)
354+
358355
case 'array':
359356
return (
360357
<ArrayInput
@@ -366,6 +363,7 @@ export function NodeInputHandler({
366363
AsyncInputComponent={AsyncInputComponent}
367364
ConfigInputComponent={ConfigInputComponent}
368365
onConfigChange={onConfigChange}
366+
variableItems={variableItems}
369367
/>
370368
)
371369

@@ -524,6 +522,7 @@ export function NodeInputHandler({
524522
disabled={disabled}
525523
inputType={inputParam?.type}
526524
language={inputParam?.type === 'code' ? inputParam.codeLanguage : undefined}
525+
suggestionItems={suggestionItems}
527526
onConfirm={handleExpandConfirm}
528527
onCancel={() => setExpandOpen(false)}
529528
/>

0 commit comments

Comments
 (0)