From a345b2ffab815749d78cf37fae00b53a73e352ae Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Sat, 1 Mar 2025 17:03:03 -0800
Subject: [PATCH] feat(code): added complex code execution for agents' custom
 tools, added envvar resolution/styling/dropdown for agent custom tool code
 modal

---
 app/api/execute/route.ts                      |  47 ++++-
 .../tool-input/components/code-editor.tsx     |  39 +++-
 .../components/custom-tool-modal.tsx          | 149 +++++++++++++-
 components/ui/env-var-dropdown.tsx            |  58 +++---
 components/ui/tag-dropdown.tsx                |  58 +++---
 executor/handlers.ts                          | 191 ++++++++++++------
 tools/index.ts                                | 113 +++++++++++
 tools/types.ts                                |   3 +
 8 files changed, 526 insertions(+), 132 deletions(-)

diff --git a/app/api/execute/route.ts b/app/api/execute/route.ts
index e33593f52f..cddfe453b6 100644
--- a/app/api/execute/route.ts
+++ b/app/api/execute/route.ts
@@ -5,6 +5,40 @@ import { Script, createContext } from 'vm'
 export const dynamic = 'force-dynamic' // Disable static optimization
 export const runtime = 'nodejs' // Use Node.js runtime
 
+/**
+ * Resolves environment variables and tags in code
+ * @param code - Code with variables
+ * @param params - Parameters that may contain variable values
+ * @param envVars - Environment variables from the workflow
+ * @returns Resolved code
+ */
+function resolveCodeVariables(
+  code: string,
+  params: Record<string, any>,
+  envVars: Record<string, string> = {}
+): string {
+  let resolvedCode = code
+
+  // Resolve environment variables with {{var_name}} syntax
+  const envVarMatches = resolvedCode.match(/\{\{([^}]+)\}\}/g) || []
+  for (const match of envVarMatches) {
+    const varName = match.slice(2, -2).trim()
+    // Priority: 1. Environment variables from workflow, 2. Params, 3. process.env
+    const varValue = envVars[varName] || params[varName] || process.env[varName] || ''
+    resolvedCode = resolvedCode.replace(match, varValue)
+  }
+
+  // Resolve tags with <tag_name> syntax
+  const tagMatches = resolvedCode.match(/<([^>]+)>/g) || []
+  for (const match of tagMatches) {
+    const tagName = match.slice(1, -1).trim()
+    const tagValue = params[tagName] || ''
+    resolvedCode = resolvedCode.replace(match, tagValue)
+  }
+
+  return resolvedCode
+}
+
 export async function POST(req: NextRequest) {
   const startTime = Date.now()
   let stdout = ''
@@ -12,18 +46,15 @@ export async function POST(req: NextRequest) {
   try {
     const body = await req.json()
 
-    const { code, params = {}, timeout = 3000 } = body
+    const { code, params = {}, timeout = 3000, envVars = {} } = body
 
-    // Check if code contains unresolved template variables
-    if (code.includes('<') && code.includes('>')) {
-      throw new Error(
-        'Code contains unresolved template variables. Please ensure all variables are resolved before execution.'
-      )
-    }
+    // Resolve variables in the code with workflow environment variables
+    const resolvedCode = resolveCodeVariables(code, params, envVars)
 
     // Create a secure context with console logging
     const context = createContext({
       params,
+      environmentVariables: envVars, // Make environment variables available in the context
       console: {
         log: (...args: any[]) => {
           const logMessage =
@@ -46,7 +77,7 @@ export async function POST(req: NextRequest) {
     const script = new Script(`
       (async () => {
         try {
-          ${code}
+          ${resolvedCode}
         } catch (error) {
           console.error(error);
           throw error;
diff --git a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor.tsx b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor.tsx
index d003f44c0d..4dc6c92654 100644
--- a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor.tsx
+++ b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor.tsx
@@ -14,6 +14,8 @@ interface CodeEditorProps {
   placeholder?: string
   className?: string
   minHeight?: string
+  highlightVariables?: boolean
+  onKeyDown?: (e: React.KeyboardEvent) => void
 }
 
 export function CodeEditor({
@@ -23,6 +25,8 @@ export function CodeEditor({
   placeholder = '',
   className = '',
   minHeight = '360px',
+  highlightVariables = true,
+  onKeyDown,
 }: CodeEditorProps) {
   const [code, setCode] = useState(value)
   const [visualLineHeights, setVisualLineHeights] = useState<number[]>([])
@@ -103,6 +107,38 @@ export function CodeEditor({
     return numbers
   }
 
+  // Custom highlighter that highlights environment variables and tags
+  const customHighlight = (code: string) => {
+    if (!highlightVariables || language !== 'javascript') {
+      // Use default Prism highlighting for non-JS or when variable highlighting is off
+      return highlight(code, languages[language], language)
+    }
+
+    // First, get the default Prism highlighting
+    let highlighted = highlight(code, languages[language], language)
+
+    // Then, highlight environment variables with {{var_name}} syntax in blue
+    if (highlighted.includes('{{')) {
+      highlighted = highlighted.replace(
+        /\{\{([^}]+)\}\}/g,
+        '<span class="text-blue-500">{{$1}}</span>'
+      )
+    }
+
+    // Also highlight tags with <tag_name> syntax in blue
+    if (highlighted.includes('<') && !language.includes('html')) {
+      highlighted = highlighted.replace(/<([^>\s/]+)>/g, (match, group) => {
+        // Avoid replacing HTML tags in comments
+        if (match.startsWith('<!--') || match.includes('</')) {
+          return match
+        }
+        return `<span class="text-blue-500">&lt;${group}&gt;</span>`
+      })
+    }
+
+    return highlighted
+  }
+
   return (
     <div
       className={cn(
@@ -159,7 +195,8 @@ export function CodeEditor({
               onChange(newCode)
             }
           }}
-          highlight={(code) => highlight(code, languages[language], language)}
+          onKeyDown={onKeyDown}
+          highlight={(code) => customHighlight(code)}
           padding={12}
           style={{
             fontFamily: 'inherit',
diff --git a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal.tsx b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal.tsx
index 2d1bc4734e..7cef756df7 100644
--- a/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal.tsx
+++ b/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useState } from 'react'
+import { useEffect, useMemo, useRef, useState } from 'react'
 import { Code, FileJson, X } from 'lucide-react'
 import { Button } from '@/components/ui/button'
 import {
@@ -9,7 +9,9 @@ import {
   DialogHeader,
   DialogTitle,
 } from '@/components/ui/dialog'
+import { EnvVarDropdown, checkEnvVarTrigger } from '@/components/ui/env-var-dropdown'
 import { Label } from '@/components/ui/label'
+import { TagDropdown, checkTagTrigger } from '@/components/ui/tag-dropdown'
 import { cn } from '@/lib/utils'
 import { useCustomToolsStore } from '@/stores/custom-tools/store'
 import { CodeEditor } from './code-editor'
@@ -52,6 +54,16 @@ export function CustomToolModal({
   const [isEditing, setIsEditing] = useState(false)
   const [toolId, setToolId] = useState<string | undefined>(undefined)
 
+  // Environment variables and tags dropdown state
+  const [showEnvVars, setShowEnvVars] = useState(false)
+  const [showTags, setShowTags] = useState(false)
+  const [searchTerm, setSearchTerm] = useState('')
+  const [cursorPosition, setCursorPosition] = useState(0)
+  const codeEditorRef = useRef<HTMLDivElement>(null)
+  const [activeSourceBlockId, setActiveSourceBlockId] = useState<string | null>(null)
+  // Add state for dropdown positioning
+  const [dropdownPosition, setDropdownPosition] = useState({ top: 0, left: 0 })
+
   const addTool = useCustomToolsStore((state) => state.addTool)
   const updateTool = useCustomToolsStore((state) => state.updateTool)
 
@@ -240,6 +252,78 @@ export function CustomToolModal({
     if (codeError) {
       setCodeError(null)
     }
+
+    // Check for environment variables and tags
+    const textarea = codeEditorRef.current?.querySelector('textarea')
+    if (textarea) {
+      const pos = textarea.selectionStart
+      setCursorPosition(pos)
+
+      // Calculate cursor position for dropdowns
+      const textBeforeCursor = value.substring(0, pos)
+      const lines = textBeforeCursor.split('\n')
+      const currentLine = lines.length
+      const currentCol = lines[lines.length - 1].length
+
+      // Find position of cursor in the editor
+      try {
+        if (codeEditorRef.current) {
+          const editorRect = codeEditorRef.current.getBoundingClientRect()
+          const lineHeight = 21 // Same as in CodeEditor
+
+          // Calculate approximate position
+          const top = currentLine * lineHeight + 5
+          const left = Math.min(currentCol * 8, editorRect.width - 260) // Prevent dropdown from going off-screen
+
+          setDropdownPosition({ top, left })
+        }
+      } catch (error) {
+        console.error('Error calculating cursor position:', error)
+      }
+
+      // Check if we should show the environment variables dropdown
+      const envVarTrigger = checkEnvVarTrigger(value, pos)
+      setShowEnvVars(envVarTrigger.show)
+      setSearchTerm(envVarTrigger.show ? envVarTrigger.searchTerm : '')
+
+      // Check if we should show the tags dropdown
+      const tagTrigger = checkTagTrigger(value, pos)
+      setShowTags(tagTrigger.show)
+      if (!tagTrigger.show) {
+        setActiveSourceBlockId(null)
+      }
+    }
+  }
+
+  // Handle environment variable selection
+  const handleEnvVarSelect = (newValue: string) => {
+    setFunctionCode(newValue)
+    setShowEnvVars(false)
+  }
+
+  // Handle tag selection
+  const handleTagSelect = (newValue: string) => {
+    setFunctionCode(newValue)
+    setShowTags(false)
+    setActiveSourceBlockId(null)
+  }
+
+  // Handle key press events
+  const handleKeyDown = (e: React.KeyboardEvent) => {
+    // Only handle Escape directly if dropdowns aren't visible
+    // Otherwise, let the dropdowns handle their own keyboard events
+    if (e.key === 'Escape' && !showEnvVars && !showTags) {
+      setShowEnvVars(false)
+      setShowTags(false)
+    }
+
+    // Don't handle other keys when dropdowns are visible
+    if (showEnvVars || showTags) {
+      if (['ArrowDown', 'ArrowUp', 'Enter'].includes(e.key)) {
+        e.preventDefault()
+        e.stopPropagation()
+      }
+    }
   }
 
   const navigationItems = [
@@ -358,14 +442,61 @@ export function CustomToolModal({
                   <span className="text-sm text-red-600 ml-4 flex-shrink-0">{codeError}</span>
                 )}
               </div>
-              <CodeEditor
-                value={functionCode}
-                onChange={handleFunctionCodeChange}
-                language="javascript"
-                placeholder={`// This code will be executed when the tool is called`}
-                minHeight="340px"
-                className={cn(codeError ? 'border-red-500' : '')}
-              />
+              <div ref={codeEditorRef} className="relative">
+                <CodeEditor
+                  value={functionCode}
+                  onChange={handleFunctionCodeChange}
+                  language="javascript"
+                  placeholder={`// This code will be executed when the tool is called. You can use environment variables with {{VARIABLE_NAME}}.`}
+                  minHeight="340px"
+                  className={cn(codeError ? 'border-red-500' : '')}
+                  highlightVariables={true}
+                  onKeyDown={handleKeyDown}
+                />
+
+                {/* Environment variables dropdown */}
+                {showEnvVars && (
+                  <EnvVarDropdown
+                    visible={showEnvVars}
+                    onSelect={handleEnvVarSelect}
+                    searchTerm={searchTerm}
+                    inputValue={functionCode}
+                    cursorPosition={cursorPosition}
+                    onClose={() => {
+                      setShowEnvVars(false)
+                      setSearchTerm('')
+                    }}
+                    className="w-64"
+                    style={{
+                      position: 'absolute',
+                      top: `${dropdownPosition.top}px`,
+                      left: `${dropdownPosition.left}px`,
+                    }}
+                  />
+                )}
+
+                {/* Tags dropdown */}
+                {showTags && (
+                  <TagDropdown
+                    visible={showTags}
+                    onSelect={handleTagSelect}
+                    blockId=""
+                    activeSourceBlockId={activeSourceBlockId}
+                    inputValue={functionCode}
+                    cursorPosition={cursorPosition}
+                    onClose={() => {
+                      setShowTags(false)
+                      setActiveSourceBlockId(null)
+                    }}
+                    className="w-64"
+                    style={{
+                      position: 'absolute',
+                      top: `${dropdownPosition.top}px`,
+                      left: `${dropdownPosition.left}px`,
+                    }}
+                  />
+                )}
+              </div>
               <div className="h-6"></div>
             </div>
           </div>
diff --git a/components/ui/env-var-dropdown.tsx b/components/ui/env-var-dropdown.tsx
index a250cc68e5..8d3dbc3976 100644
--- a/components/ui/env-var-dropdown.tsx
+++ b/components/ui/env-var-dropdown.tsx
@@ -10,6 +10,7 @@ interface EnvVarDropdownProps {
   inputValue: string
   cursorPosition: number
   onClose?: () => void
+  style?: React.CSSProperties
 }
 
 export const EnvVarDropdown: React.FC<EnvVarDropdownProps> = ({
@@ -20,6 +21,7 @@ export const EnvVarDropdown: React.FC<EnvVarDropdownProps> = ({
   inputValue,
   cursorPosition,
   onClose,
+  style,
 }) => {
   const envVars = useEnvironmentStore((state) => Object.keys(state.variables))
   const [selectedIndex, setSelectedIndex] = useState(0)
@@ -57,35 +59,38 @@ export const EnvVarDropdown: React.FC<EnvVarDropdownProps> = ({
     onClose?.()
   }
 
-  // Handle keyboard navigation
-  const handleKeyDown = (e: KeyboardEvent) => {
-    if (!visible || filteredEnvVars.length === 0) return
-
-    switch (e.key) {
-      case 'ArrowDown':
-        e.preventDefault()
-        setSelectedIndex((prev) => (prev < filteredEnvVars.length - 1 ? prev + 1 : prev))
-        break
-      case 'ArrowUp':
-        e.preventDefault()
-        setSelectedIndex((prev) => (prev > 0 ? prev - 1 : prev))
-        break
-      case 'Enter':
-        e.preventDefault()
-        handleEnvVarSelect(filteredEnvVars[selectedIndex])
-        break
-      case 'Escape':
-        e.preventDefault()
-        onClose?.()
-        break
-    }
-  }
-
   // Add and remove keyboard event listener
   useEffect(() => {
     if (visible) {
-      window.addEventListener('keydown', handleKeyDown)
-      return () => window.removeEventListener('keydown', handleKeyDown)
+      const handleKeyboardEvent = (e: KeyboardEvent) => {
+        if (!filteredEnvVars.length) return
+
+        switch (e.key) {
+          case 'ArrowDown':
+            e.preventDefault()
+            e.stopPropagation()
+            setSelectedIndex((prev) => (prev < filteredEnvVars.length - 1 ? prev + 1 : prev))
+            break
+          case 'ArrowUp':
+            e.preventDefault()
+            e.stopPropagation()
+            setSelectedIndex((prev) => (prev > 0 ? prev - 1 : prev))
+            break
+          case 'Enter':
+            e.preventDefault()
+            e.stopPropagation()
+            handleEnvVarSelect(filteredEnvVars[selectedIndex])
+            break
+          case 'Escape':
+            e.preventDefault()
+            e.stopPropagation()
+            onClose?.()
+            break
+        }
+      }
+
+      window.addEventListener('keydown', handleKeyboardEvent, true)
+      return () => window.removeEventListener('keydown', handleKeyboardEvent, true)
     }
   }, [visible, selectedIndex, filteredEnvVars])
 
@@ -97,6 +102,7 @@ export const EnvVarDropdown: React.FC<EnvVarDropdownProps> = ({
         'absolute z-[9999] w-full mt-1 overflow-hidden bg-popover rounded-md border shadow-md',
         className
       )}
+      style={style}
     >
       {filteredEnvVars.length === 0 ? (
         <div className="px-3 py-2 text-sm text-muted-foreground">
diff --git a/components/ui/tag-dropdown.tsx b/components/ui/tag-dropdown.tsx
index 905231b913..d53760578d 100644
--- a/components/ui/tag-dropdown.tsx
+++ b/components/ui/tag-dropdown.tsx
@@ -27,6 +27,7 @@ interface TagDropdownProps {
   inputValue: string
   cursorPosition: number
   onClose?: () => void
+  style?: React.CSSProperties
 }
 
 export const TagDropdown: React.FC<TagDropdownProps> = ({
@@ -38,6 +39,7 @@ export const TagDropdown: React.FC<TagDropdownProps> = ({
   inputValue,
   cursorPosition,
   onClose,
+  style,
 }) => {
   const [selectedIndex, setSelectedIndex] = useState(0)
 
@@ -198,35 +200,38 @@ export const TagDropdown: React.FC<TagDropdownProps> = ({
     onClose?.()
   }
 
-  // Handle keyboard navigation
-  const handleKeyDown = (e: KeyboardEvent) => {
-    if (!visible || filteredTags.length === 0) return
-
-    switch (e.key) {
-      case 'ArrowDown':
-        e.preventDefault()
-        setSelectedIndex((prev) => (prev < filteredTags.length - 1 ? prev + 1 : prev))
-        break
-      case 'ArrowUp':
-        e.preventDefault()
-        setSelectedIndex((prev) => (prev > 0 ? prev - 1 : prev))
-        break
-      case 'Enter':
-        e.preventDefault()
-        handleTagSelect(filteredTags[selectedIndex])
-        break
-      case 'Escape':
-        e.preventDefault()
-        onClose?.()
-        break
-    }
-  }
-
   // Add and remove keyboard event listener
   useEffect(() => {
     if (visible) {
-      window.addEventListener('keydown', handleKeyDown)
-      return () => window.removeEventListener('keydown', handleKeyDown)
+      const handleKeyboardEvent = (e: KeyboardEvent) => {
+        if (!filteredTags.length) return
+
+        switch (e.key) {
+          case 'ArrowDown':
+            e.preventDefault()
+            e.stopPropagation()
+            setSelectedIndex((prev) => (prev < filteredTags.length - 1 ? prev + 1 : prev))
+            break
+          case 'ArrowUp':
+            e.preventDefault()
+            e.stopPropagation()
+            setSelectedIndex((prev) => (prev > 0 ? prev - 1 : prev))
+            break
+          case 'Enter':
+            e.preventDefault()
+            e.stopPropagation()
+            handleTagSelect(filteredTags[selectedIndex])
+            break
+          case 'Escape':
+            e.preventDefault()
+            e.stopPropagation()
+            onClose?.()
+            break
+        }
+      }
+
+      window.addEventListener('keydown', handleKeyboardEvent, true)
+      return () => window.removeEventListener('keydown', handleKeyboardEvent, true)
     }
   }, [visible, selectedIndex, filteredTags])
 
@@ -239,6 +244,7 @@ export const TagDropdown: React.FC<TagDropdownProps> = ({
         'absolute z-[9999] w-full mt-1 overflow-hidden bg-popover rounded-md border shadow-md',
         className
       )}
+      style={style}
     >
       <div className="py-1">
         {filteredTags.length === 0 ? (
diff --git a/executor/handlers.ts b/executor/handlers.ts
index f3b03e1f3e..4b06b32c85 100644
--- a/executor/handlers.ts
+++ b/executor/handlers.ts
@@ -37,6 +37,62 @@ export interface BlockHandler {
   ): Promise<BlockOutput>
 }
 
+/**
+ * Shared helper for executing code with WebContainer and VM fallback
+ * @param code - The code to execute
+ * @param params - Parameters to pass to the code
+ * @param timeout - Execution timeout in milliseconds
+ * @returns Execution result
+ */
+async function executeCodeWithFallback(
+  code: string,
+  params: Record<string, any> = {},
+  timeout: number = 5000
+): Promise<{ success: boolean; output: any; error?: string }> {
+  // Only try WebContainer in browser environment with direct execution
+  const isBrowser = typeof window !== 'undefined'
+  if (isBrowser && window.crossOriginIsolated) {
+    try {
+      // Dynamically import WebContainer to prevent server-side import
+      const { executeCode } = await import('@/lib/webcontainer')
+
+      // Execute directly in the browser
+      const result = await executeCode(code, params, timeout)
+
+      if (!result.success) {
+        console.warn(`WebContainer API execution failed: ${result.error}`)
+        throw new Error(result.error || `WebContainer execution failed with no error message`)
+      }
+
+      return { success: true, output: result.output }
+    } catch (error: any) {
+      console.warn('WebContainer execution failed, falling back to VM:', error)
+      console.error('WebContainer error details:', {
+        name: error.name,
+        message: error.message,
+        stack: error.stack,
+      })
+    }
+  }
+
+  // Fall back to VM execution if WebContainer fails or not available
+  try {
+    const vmResult = await executeTool('function_execute', { code, ...params }, true)
+
+    if (!vmResult.success) {
+      throw new Error(vmResult.error || `Function execution failed with no error message`)
+    }
+
+    return { success: true, output: vmResult.output }
+  } catch (vmError: any) {
+    return {
+      success: false,
+      output: null,
+      error: `Function execution failed: ${vmError.message}`,
+    }
+  }
+}
+
 /**
  * Handler for Agent blocks that process LLM requests with optional tools.
  */
@@ -68,31 +124,72 @@ export class AgentBlockHandler implements BlockHandler {
 
     // Format tools for provider API
     const formattedTools = Array.isArray(inputs.tools)
-      ? inputs.tools
-          .map((tool: any) => {
-            // Handle custom tools
-            if (tool.type === 'custom-tool' && tool.schema) {
-              return {
-                id: `custom_${tool.title}`,
-                name: tool.schema.function.name,
-                description: tool.schema.function.description || '',
-                params: tool.params || {},
-                parameters: {
-                  type: tool.schema.function.parameters.type,
-                  properties: tool.schema.function.parameters.properties,
-                  required: tool.schema.function.parameters.required || [],
-                },
-              }
-            }
+      ? (
+          await Promise.all(
+            inputs.tools.map(async (tool: any) => {
+              // Handle custom tools
+              if (tool.type === 'custom-tool' && tool.schema) {
+                // Add function execution capability to custom tools with code
+                if (tool.code) {
+                  // Store the tool's code and make it available for execution
+                  const toolName = tool.schema.function.name
+                  const params = tool.params || {}
 
-            // Handle regular block tools with operation selection
-            return transformBlockTool(tool, {
-              selectedOperation: tool.operation,
-              getAllBlocks,
-              getTool,
+                  // Create a tool that can execute the code
+                  return {
+                    id: `custom_${tool.title}`,
+                    name: toolName,
+                    description: tool.schema.function.description || '',
+                    params: params,
+                    parameters: {
+                      type: tool.schema.function.parameters.type,
+                      properties: tool.schema.function.parameters.properties,
+                      required: tool.schema.function.parameters.required || [],
+                    },
+                    executeFunction: async (callParams: Record<string, any>) => {
+                      try {
+                        // Execute the code with WebContainer fallback
+                        const result = await executeCodeWithFallback(
+                          tool.code,
+                          { ...params, ...callParams },
+                          tool.timeout || 5000
+                        )
+
+                        if (!result.success) {
+                          throw new Error(result.error || 'Function execution failed')
+                        }
+
+                        return result.output
+                      } catch (error: any) {
+                        console.error(`Error executing custom tool ${toolName}:`, error)
+                        throw new Error(`Error in ${toolName}: ${error.message}`)
+                      }
+                    },
+                  }
+                }
+
+                return {
+                  id: `custom_${tool.title}`,
+                  name: tool.schema.function.name,
+                  description: tool.schema.function.description || '',
+                  params: tool.params || {},
+                  parameters: {
+                    type: tool.schema.function.parameters.type,
+                    properties: tool.schema.function.parameters.properties,
+                    required: tool.schema.function.parameters.required || [],
+                  },
+                }
+              }
+
+              // Handle regular block tools with operation selection
+              return transformBlockTool(tool, {
+                selectedOperation: tool.operation,
+                getAllBlocks,
+                getTool,
+              })
             })
-          })
-          .filter((t): t is NonNullable<typeof t> => t !== null)
+          )
+        ).filter((t: any): t is NonNullable<typeof t> => t !== null)
       : []
 
     // Ensure context is properly formatted for the provider
@@ -460,49 +557,19 @@ export class FunctionBlockHandler implements BlockHandler {
     inputs: Record<string, any>,
     context: ExecutionContext
   ): Promise<BlockOutput> {
-    // Only try WebContainer in browser environment with direct execution
-    const isBrowser = typeof window !== 'undefined'
-    if (isBrowser && window.crossOriginIsolated) {
-      try {
-        // Dynamically import WebContainer to prevent server-side import
-        const { executeCode } = await import('@/lib/webcontainer')
+    // Prepare code for execution
+    const codeContent = Array.isArray(inputs.code)
+      ? inputs.code.map((c: { content: string }) => c.content).join('\n')
+      : inputs.code
 
-        // Prepare code for execution
-        const codeContent = Array.isArray(inputs.code)
-          ? inputs.code.map((c: { content: string }) => c.content).join('\n')
-          : inputs.code
+    // Use the shared helper function
+    const result = await executeCodeWithFallback(codeContent, inputs, inputs.timeout || 5000)
 
-        // Execute directly in the browser
-        const result = await executeCode(codeContent, inputs, inputs.timeout || 5000)
-
-        if (!result.success) {
-          console.warn(`WebContainer API execution failed: ${result.error}`)
-          throw new Error(result.error || `WebContainer execution failed with no error message`)
-        }
-
-        return { response: result.output }
-      } catch (error: any) {
-        console.warn('WebContainer execution failed, falling back to VM:', error)
-        console.error('WebContainer error details:', {
-          name: error.name,
-          message: error.message,
-          stack: error.stack,
-        })
-      }
+    if (!result.success) {
+      throw new Error(result.error || 'Function execution failed')
     }
 
-    // Fall back to VM execution if WebContainer fails or not available
-    try {
-      const vmResult = await executeTool('function_execute', inputs, true)
-
-      if (!vmResult.success) {
-        throw new Error(vmResult.error || `Function execution failed with no error message`)
-      }
-
-      return { response: vmResult.output }
-    } catch (vmError: any) {
-      throw new Error(`Function execution failed: ${vmError.message}`)
-    }
+    return { response: result.output }
   }
 }
 
diff --git a/tools/index.ts b/tools/index.ts
index 83ac9a6135..86a33e5cd4 100644
--- a/tools/index.ts
+++ b/tools/index.ts
@@ -1,4 +1,5 @@
 import { useCustomToolsStore } from '@/stores/custom-tools/store'
+import { useEnvironmentStore } from '@/stores/settings/environment/store'
 import { visionTool as crewAIVision } from './crewai/vision'
 import { scrapeTool } from './firecrawl/scrape'
 import { functionExecuteTool as functionExecute } from './function/execute'
@@ -82,6 +83,16 @@ export function getTool(toolId: string): ToolConfig | undefined {
   return undefined
 }
 
+// Check if we're running in the browser
+function isBrowser(): boolean {
+  return typeof window !== 'undefined'
+}
+
+// Check if WebContainer is available
+function isWebContainerAvailable(): boolean {
+  return isBrowser() && !!window.crossOriginIsolated
+}
+
 // Create a tool config from a custom tool definition
 function getCustomTool(customToolId: string): ToolConfig | undefined {
   // Extract the identifier part (could be UUID or title)
@@ -131,16 +142,109 @@ function getCustomTool(customToolId: string): ToolConfig | undefined {
       method: 'POST',
       headers: () => ({ 'Content-Type': 'application/json' }),
       body: (params: Record<string, any>) => {
+        // Get environment variables from the store
+        const envStore = useEnvironmentStore.getState()
+        const allEnvVars = envStore.getAllVariables()
+
+        // Convert environment variables to a simple key-value object
+        const envVars = Object.entries(allEnvVars).reduce(
+          (acc, [key, variable]) => {
+            acc[key] = variable.value
+            return acc
+          },
+          {} as Record<string, string>
+        )
+
         // Include everything needed for execution
         return {
           code: customTool.code,
           params: params, // These will be available in the VM context
           schema: customTool.schema.function.parameters, // For validation on the client side
+          envVars: envVars, // Pass environment variables for server-side resolution
         }
       },
       isInternalRoute: true,
     },
 
+    // Direct execution support for browser environment with WebContainer
+    directExecution: async (params: Record<string, any>) => {
+      // If there's no code, we can't execute directly
+      if (!customTool.code) {
+        return {
+          success: false,
+          output: {},
+          error: 'No code provided for tool execution',
+        }
+      }
+
+      // If we're in a browser with WebContainer available, use it
+      if (isWebContainerAvailable()) {
+        try {
+          // Get environment variables from the store
+          const envStore = useEnvironmentStore.getState()
+          const envVars = envStore.getAllVariables()
+
+          // Create a merged params object that includes environment variables
+          const mergedParams = { ...params }
+
+          // Add environment variables to the params
+          Object.entries(envVars).forEach(([key, variable]) => {
+            if (variable.value && !mergedParams[key]) {
+              mergedParams[key] = variable.value
+            }
+          })
+
+          // Resolve environment variables and tags in the code
+          let resolvedCode = customTool.code
+
+          // Resolve environment variables with {{var_name}} syntax
+          const envVarMatches = resolvedCode.match(/\{\{([^}]+)\}\}/g) || []
+          for (const match of envVarMatches) {
+            const varName = match.slice(2, -2).trim()
+            // Look for the variable in our environment store first, then in params
+            const envVar = envVars[varName]
+            const varValue = envVar ? envVar.value : mergedParams[varName] || ''
+            resolvedCode = resolvedCode.replace(match, varValue)
+          }
+
+          // Resolve tags with <tag_name> syntax
+          const tagMatches = resolvedCode.match(/<([^>]+)>/g) || []
+          for (const match of tagMatches) {
+            const tagName = match.slice(1, -1).trim()
+            const tagValue = mergedParams[tagName] || ''
+            resolvedCode = resolvedCode.replace(match, tagValue)
+          }
+
+          // Dynamically import the executeCode function
+          const { executeCode } = await import('@/lib/webcontainer')
+
+          // Execute the code with resolved variables
+          const result = await executeCode(
+            resolvedCode,
+            mergedParams, // Use the merged params that include env vars
+            5000 // Default timeout
+          )
+
+          if (!result.success) {
+            throw new Error(result.error || 'WebContainer execution failed')
+          }
+
+          return {
+            success: true,
+            output: result.output.result || result.output,
+            error: undefined,
+          }
+        } catch (error: any) {
+          console.warn('WebContainer execution failed, falling back to API:', error.message)
+          // Fall back to API route if WebContainer fails
+          return undefined
+        }
+      }
+
+      // No WebContainer or not in browser, return undefined to use regular API route
+      return undefined
+    },
+
     // Response handling
     transformResponse: async (response: Response) => {
       const data = await response.json()
@@ -177,6 +281,15 @@ export async function executeTool(
       throw new Error(`Tool not found: ${toolId}`)
     }
 
+    // For custom tools, try direct execution in browser first if available
+    if (toolId.startsWith('custom_') && tool.directExecution) {
+      const directResult = await tool.directExecution(params)
+      if (directResult) {
+        return directResult
+      }
+      // If directExecution returns undefined, fall back to API route
+    }
+
     // For internal routes or when skipProxy is true, call the API directly
     if (tool.request.isInternalRoute || skipProxy) {
       const result = await handleInternalRequest(toolId, tool, params)
diff --git a/tools/types.ts b/tools/types.ts
index be6e6a5cfb..6bf2d7a593 100644
--- a/tools/types.ts
+++ b/tools/types.ts
@@ -34,6 +34,9 @@ export interface ToolConfig<P = any, R extends ToolResponse = ToolResponse> {
     isInternalRoute?: boolean // Whether this is an internal API route
   }
 
+  // Direct execution in browser (optional) - bypasses HTTP request
+  directExecution?: (params: P) => Promise<R | undefined>
+
   // Response handling
   transformResponse: (response: Response) => Promise<R>
   transformError: (error: any) => string