Skip to content

feat(autofill): consolidated tool-params & sub-block store for one unified store that handles sub-block values #237

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
waleedlatif1 merged 1 commit into from
Apr 10, 2025
Merged

feat(autofill): consolidated tool-params & sub-block store for one unified store that handles sub-block values #237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions .../[id]/components/workflow-block/components/sub-block/components/tool-input/tool-input.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import { OAuthProvider } from '@/lib/oauth'
import { cn } from '@/lib/utils'
import { useCustomToolsStore } from '@/stores/custom-tools/store'
import { useGeneralStore } from '@/stores/settings/general/store'
import { useToolParamsStore } from '@/stores/tool-params/store'
import { useSubBlockStore } from '@/stores/workflows/subblock/store'
import { useWorkflowStore } from '@/stores/workflows/workflow/store'
import { getAllBlocks } from '@/blocks'
import { getTool } from '@/tools'
Expand Down Expand Up @@ -122,8 +122,12 @@ const getOperationOptions = (blockType: string): { label: string; id: string }[]
const initializeToolParams = (
toolId: string,
params: ToolParam[],
toolParamsStore: {
resolveParamValue: (toolId: string, paramId: string, instanceId?: string) => string | undefined
subBlockStore: {
resolveToolParamValue: (
toolId: string,
paramId: string,
instanceId?: string
) => string | undefined
},
isAutoFillEnabled: boolean,
instanceId?: string
Expand All @@ -134,7 +138,7 @@ const initializeToolParams = (
if (isAutoFillEnabled) {
// For each parameter, check if we have a stored/resolved value
params.forEach((param) => {
const resolvedValue = toolParamsStore.resolveParamValue(toolId, param.id, instanceId)
const resolvedValue = subBlockStore.resolveToolParamValue(toolId, param.id, instanceId)
if (resolvedValue) {
initialParams[param.id] = resolvedValue
}
Expand All @@ -152,7 +156,7 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
const [searchQuery, setSearchQuery] = useState('')
const isWide = useWorkflowStore((state) => state.blocks[blockId]?.isWide)
const customTools = useCustomToolsStore((state) => state.getAllTools())
const toolParamsStore = useToolParamsStore()
const subBlockStore = useSubBlockStore()
const isAutoFillEnvVarsEnabled = useGeneralStore((state) => state.isAutoFillEnvVarsEnabled)

const toolBlocks = getAllBlocks().filter((block) => block.category === 'tools')
Expand Down Expand Up @@ -194,7 +198,7 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
const initialParams = initializeToolParams(
toolId,
requiredParams,
toolParamsStore,
subBlockStore,
isAutoFillEnvVarsEnabled,
blockId
)
Expand Down Expand Up @@ -247,7 +251,7 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
const initialParams = initializeToolParams(
toolId,
toolParams,
toolParamsStore,
subBlockStore,
isAutoFillEnvVarsEnabled,
blockId
)
Expand Down Expand Up @@ -327,7 +331,7 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {

// Only store non-empty values
if (paramValue.trim()) {
toolParamsStore.setParam(toolId, paramId, paramValue)
subBlockStore.setToolParam(toolId, paramId, paramValue)
}

// Update the value in the workflow
Expand Down
67 changes: 45 additions & 22 deletions sim/app/w/[id]/components/workflow-block/components/sub-block/hooks/use-sub-block-value.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import { useCallback, useEffect, useRef } from 'react'
import { isEqual } from 'lodash'
import { useGeneralStore } from '@/stores/settings/general/store'
import { useToolParamsStore } from '@/stores/tool-params/store'
import { useSubBlockStore } from '@/stores/workflows/subblock/store'
import { useWorkflowStore } from '@/stores/workflows/workflow/store'
import { getProviderFromModel } from '@/providers/utils'
Expand All @@ -24,20 +23,19 @@ function handleAgentBlockApiKey(
// Skip if we couldn't determine a provider
if (!provider || provider === 'ollama') return

const toolParamsStore = useToolParamsStore.getState()
const subBlockStore = useSubBlockStore.getState()

// Try to get a saved API key for this provider
const savedValue = toolParamsStore.resolveParamValue(provider, 'apiKey', blockId)
const savedValue = subBlockStore.resolveToolParamValue(provider, 'apiKey', blockId)

// If we have a valid API key, use it
if (savedValue && savedValue !== '') {
// Only update if different from current value
if (savedValue !== storeValue) {
subBlockStore.setValue(blockId, subBlockId, savedValue)
}
// Always update the value when switching models, even if it appears the same
// This handles cases where the field shows masked values but needs to update
subBlockStore.setValue(blockId, subBlockId, savedValue)
} else {
// No API key found for this provider - ALWAYS clear the field
// Always clear the field when switching to a model with no API key
// Don't wait for user interaction to clear it
subBlockStore.setValue(blockId, subBlockId, '')
}
}
Expand All @@ -53,16 +51,16 @@ function handleStandardBlockApiKey(
) {
if (!blockType) return

const toolParamsStore = useToolParamsStore.getState()
const subBlockStore = useSubBlockStore.getState()

// Only auto-fill if the field is empty
if (!storeValue || storeValue === '') {
// Pass the blockId as instanceId to check if this specific instance has been cleared
const savedValue = toolParamsStore.resolveParamValue(blockType, 'apiKey', blockId)
const savedValue = subBlockStore.resolveToolParamValue(blockType, 'apiKey', blockId)

if (savedValue && savedValue !== '' && savedValue !== storeValue) {
// Auto-fill the API key from the param store
useSubBlockStore.getState().setValue(blockId, subBlockId, savedValue)
subBlockStore.setValue(blockId, subBlockId, savedValue)
}
}
// Handle environment variable references
Expand All @@ -73,13 +71,13 @@ function handleStandardBlockApiKey(
storeValue.endsWith('}}')
) {
// Pass the blockId as instanceId
const currentValue = toolParamsStore.resolveParamValue(blockType, 'apiKey', blockId)
const currentValue = subBlockStore.resolveToolParamValue(blockType, 'apiKey', blockId)

if (currentValue !== storeValue) {
// If we got a replacement or null, update the field
if (currentValue) {
// Replacement found - update to new reference
useSubBlockStore.getState().setValue(blockId, subBlockId, currentValue)
subBlockStore.setValue(blockId, subBlockId, currentValue)
}
}
}
Expand All @@ -97,32 +95,39 @@ function storeApiKeyValue(
) {
if (!blockType) return

const toolParamsStore = useToolParamsStore.getState()
const subBlockStore = useSubBlockStore.getState()

// Check if this is an empty value for an API key field that previously had a value
// This indicates the user has deliberately cleared the field
// Check if this is user explicitly clearing a field that had a value
// We only want to mark it as cleared if it's a user action, not an automatic
// clearing from model switching
if (
storeValue &&
storeValue !== '' &&
(newValue === null || newValue === '' || String(newValue).trim() === '')
) {
// Mark this specific instance as cleared so we don't auto-fill it
toolParamsStore.markParamAsCleared(blockId, 'apiKey')
subBlockStore.markParamAsCleared(blockId, 'apiKey')
return
}

// Only store non-empty values
if (!newValue || String(newValue).trim() === '') return

// If user enters a value, we should clear any "cleared" flag
// to ensure auto-fill will work in the future
if (subBlockStore.isParamCleared(blockId, 'apiKey')) {
subBlockStore.unmarkParamAsCleared(blockId, 'apiKey')
}

// For agent blocks, store the API key under the provider name
if (blockType === 'agent' && modelValue) {
const provider = getProviderFromModel(modelValue)
if (provider && provider !== 'ollama') {
toolParamsStore.setParam(provider, 'apiKey', String(newValue))
subBlockStore.setToolParam(provider, 'apiKey', String(newValue))
}
} else {
// For other blocks, store under the block type
toolParamsStore.setParam(blockType, 'apiKey', String(newValue))
subBlockStore.setToolParam(blockType, 'apiKey', String(newValue))
}
}

Expand Down Expand Up @@ -184,9 +189,27 @@ export function useSubBlockValue<T = any>(
// Update the previous model reference
prevModelRef.current = modelValue

// Handle API key autofill for the new model
if (isAutoFillEnvVarsEnabled && modelValue) {
handleAgentBlockApiKey(blockId, subBlockId, modelValue, storeValue)
// For agent blocks, always clear the field if needed
// But only fill with saved values if auto-fill is enabled
if (modelValue) {
const provider = getProviderFromModel(modelValue)

// Skip if we couldn't determine a provider
if (!provider || provider === 'ollama') return

const subBlockStore = useSubBlockStore.getState()

// Check if there's a saved value for this provider
const savedValue = subBlockStore.resolveToolParamValue(provider, 'apiKey', blockId)

if (savedValue && savedValue !== '' && isAutoFillEnvVarsEnabled) {
// Only auto-fill if the feature is enabled
subBlockStore.setValue(blockId, subBlockId, savedValue)
} else {
// Always clear immediately when switching to a model with no saved key
// or when auto-fill is disabled
subBlockStore.setValue(blockId, subBlockId, '')
}
}
}
}, [blockId, subBlockId, blockType, isApiKey, modelValue, isAutoFillEnvVarsEnabled, storeValue])
Expand Down
6 changes: 1 addition & 5 deletions sim/stores/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ import { useConsoleStore } from './panel/console/store'
import { useVariablesStore } from './panel/variables/store'
import { useEnvironmentStore } from './settings/environment/store'
import { getSyncManagers, initializeSyncManagers, resetSyncManagers } from './sync-registry'
import { useToolParamsStore } from './tool-params/store'
import {
loadRegistry,
loadSubblockValues,
Expand All @@ -18,7 +17,6 @@ import {
} from './workflows/persistence'
import { useWorkflowRegistry } from './workflows/registry/store'
import { useSubBlockStore } from './workflows/subblock/store'
import { workflowSync } from './workflows/sync'
import { useWorkflowStore } from './workflows/workflow/store'

const logger = createLogger('Stores')
Expand Down Expand Up @@ -232,7 +230,6 @@ export {
useChatStore,
useCustomToolsStore,
useVariablesStore,
useToolParamsStore,
}

// Helper function to reset all stores
Expand All @@ -246,6 +243,7 @@ export const resetAllStores = () => {
})
useWorkflowStore.getState().clear()
useSubBlockStore.getState().clear()
useSubBlockStore.getState().clearToolParams()
useNotificationStore.setState({ notifications: [] })
useEnvironmentStore.setState({
variables: {},
Expand All @@ -257,7 +255,6 @@ export const resetAllStores = () => {
useChatStore.setState({ messages: [], isProcessing: false, error: null })
useCustomToolsStore.setState({ tools: {} })
useVariablesStore.getState().resetLoaded() // Reset variables store tracking
useToolParamsStore.getState().clear()
}

// Helper function to log all store states
Expand All @@ -273,7 +270,6 @@ export const logAllStores = () => {
customTools: useCustomToolsStore.getState(),
subBlock: useSubBlockStore.getState(),
variables: useVariablesStore.getState(),
toolParams: useToolParamsStore.getState(),
}

return state
Expand Down
Loading