diff --git a/app/w/[id]/components/workflow-block/components/sub-block/components/short-input.tsx b/app/w/[id]/components/workflow-block/components/sub-block/components/short-input.tsx
index a210378065..7fd47f6368 100644
--- a/app/w/[id]/components/workflow-block/components/sub-block/components/short-input.tsx
+++ b/app/w/[id]/components/workflow-block/components/sub-block/components/short-input.tsx
@@ -4,7 +4,10 @@ import { useSubBlockValue } from '../hooks/use-sub-block-value'
 import { cn } from '@/lib/utils'
 import { SubBlockConfig } from '@/blocks/types'
 import { formatDisplayText } from '@/components/ui/formatted-text'
-import { EnvVarDropdown, checkEnvVarTrigger } from '@/components/ui/env-var-dropdown'
+import {
+  EnvVarDropdown,
+  checkEnvVarTrigger,
+} from '@/components/ui/env-var-dropdown'
 import { TagDropdown, checkTagTrigger } from '@/components/ui/tag-dropdown'
 
 interface ShortInputProps {
@@ -14,6 +17,8 @@ interface ShortInputProps {
   subBlockId: string
   isConnecting: boolean
   config: SubBlockConfig
+  value?: string
+  onChange?: (value: string) => void
 }
 
 export function ShortInput({
@@ -23,28 +28,39 @@ export function ShortInput({
   password,
   isConnecting,
   config,
+  value: propValue,
+  onChange,
 }: ShortInputProps) {
   const [isFocused, setIsFocused] = useState(false)
   const [showEnvVars, setShowEnvVars] = useState(false)
   const [showTags, setShowTags] = useState(false)
-  const [value, setValue] = useSubBlockValue(blockId, subBlockId)
+  const [storeValue, setStoreValue] = useSubBlockValue(blockId, subBlockId)
   const [searchTerm, setSearchTerm] = useState('')
   const [cursorPosition, setCursorPosition] = useState(0)
   const inputRef = useRef<HTMLInputElement>(null)
   const overlayRef = useRef<HTMLDivElement>(null)
 
+  // Use either controlled or uncontrolled value
+  const value = propValue !== undefined ? propValue : storeValue
+
   // Handle input changes
   const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
     const newValue = e.target.value
     const newCursorPosition = e.target.selectionStart ?? 0
-    setValue(newValue)
+
+    if (onChange) {
+      onChange(newValue)
+    } else {
+      setStoreValue(newValue)
+    }
+
     setCursorPosition(newCursorPosition)
-    
+
     // Check for environment variables trigger
     const envVarTrigger = checkEnvVarTrigger(newValue, newCursorPosition)
     setShowEnvVars(envVarTrigger.show)
     setSearchTerm(envVarTrigger.show ? envVarTrigger.searchTerm : '')
-    
+
     // Check for tag trigger
     const tagTrigger = checkTagTrigger(newValue, newCursorPosition)
     setShowTags(tagTrigger.show)
@@ -85,18 +101,22 @@ export function ShortInput({
       if (data.type !== 'connectionBlock') return
 
       // Get current cursor position or append to end
-      const dropPosition = inputRef.current?.selectionStart ?? value?.toString().length ?? 0
-      
+      const dropPosition =
+        inputRef.current?.selectionStart ?? value?.toString().length ?? 0
+
       // Insert '<' at drop position to trigger the dropdown
       const currentValue = value?.toString() ?? ''
-      const newValue = currentValue.slice(0, dropPosition) + '<' + currentValue.slice(dropPosition)
-      
+      const newValue =
+        currentValue.slice(0, dropPosition) +
+        '<' +
+        currentValue.slice(dropPosition)
+
       // Focus the input first
       inputRef.current?.focus()
 
       // Update all state in a single batch
       Promise.resolve().then(() => {
-        setValue(newValue)
+        setStoreValue(newValue)
         setCursorPosition(dropPosition + 1)
         setShowTags(true)
 
@@ -127,6 +147,15 @@ export function ShortInput({
       ? '•'.repeat(value?.toString().length ?? 0)
       : value?.toString() ?? ''
 
+  // Modify the EnvVarDropdown to use the correct setter
+  const handleEnvVarSelect = (newValue: string) => {
+    if (onChange) {
+      onChange(newValue)
+    } else {
+      setStoreValue(newValue)
+    }
+  }
+
   return (
     <div className="relative w-full">
       <Input
@@ -164,7 +193,7 @@ export function ShortInput({
       </div>
       <EnvVarDropdown
         visible={showEnvVars}
-        onSelect={setValue}
+        onSelect={handleEnvVarSelect}
         searchTerm={searchTerm}
         inputValue={value?.toString() ?? ''}
         cursorPosition={cursorPosition}
@@ -175,7 +204,7 @@ export function ShortInput({
       />
       <TagDropdown
         visible={showTags}
-        onSelect={setValue}
+        onSelect={handleEnvVarSelect}
         blockId={blockId}
         inputValue={value?.toString() ?? ''}
         cursorPosition={cursorPosition}
diff --git a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input.tsx b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input.tsx
index 733487ad37..bf6d83b54c 100644
--- a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input.tsx
+++ b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input.tsx
@@ -118,18 +118,17 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
     paramValue: string
   ) => {
     setValue(
-      selectedTools.map((tool) => {
-        if (tool.type === toolType) {
-          return {
-            ...tool,
-            params: {
-              ...tool.params,
-              [paramId]: paramValue,
-            },
-          }
-        }
-        return tool
-      })
+      selectedTools.map((tool) =>
+        tool.type === toolType
+          ? {
+              ...tool,
+              params: {
+                ...tool.params,
+                [paramId]: paramValue,
+              },
+            }
+          : tool
+      )
     )
   }
 
@@ -194,7 +193,6 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
       ) : (
         <div className="flex flex-wrap gap-2 min-h-[2.5rem] w-full rounded-md border border-input bg-transparent px-3 py-2 text-sm ring-offset-background">
           {selectedTools.map((tool) => {
-            // Get UI properties from toolBlocks for display
             const toolBlock = toolBlocks.find(
               (block) => block.type === tool.type
             )
@@ -233,20 +231,21 @@ export function ToolInput({ blockId, subBlockId }: ToolInputProps) {
                             {param.id}
                           </div>
                           <div className="relative">
-                            <input
-                              type={
-                                param.type === 'password' ? 'password' : 'text'
-                              }
-                              value={tool.params?.[param.id] || ''}
-                              onChange={(e) =>
-                                handleParamChange(
-                                  tool.type,
-                                  param.id,
-                                  e.target.value
-                                )
-                              }
+                            <ShortInput
+                              blockId={blockId}
+                              subBlockId={`${subBlockId}-param`}
                               placeholder={param.description}
-                              className="flex h-9 w-full rounded-md border border-input bg-transparent px-3 py-1 text-sm shadow-sm transition-colors file:border-0 file:bg-transparent file:text-sm file:font-medium placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:cursor-not-allowed disabled:opacity-50"
+                              password={param.type === 'password'}
+                              isConnecting={false}
+                              config={{
+                                id: `${subBlockId}-param`,
+                                type: 'short-input',
+                                title: param.id,
+                              }}
+                              value={tool.params[param.id] || ''}
+                              onChange={(value) =>
+                                handleParamChange(tool.type, param.id, value)
+                              }
                             />
                           </div>
                         </div>