Improve AutoGen Studio: deprecate FunctionTool, harden MCP WebSocket endpoint

- Deprecate FunctionTool creation in the UI; show deprecation warning for existing configs
- Skip FunctionTool instantiation during validation to avoid exec() on user code
- Store MCP server params server-side instead of passing via WebSocket query string
- Remove FunctionTool template and direct users to MCP Workbenches

Author: victordibia

.gitignore

@@ -203,3 +203,6 @@ registry.json
 # files created by the gitty agent in python/samples/gitty
 .gitty/
 .aider*
+
+# Claude Code
+.claude/

python/packages/autogen-studio/autogenstudio/validation/validation_service.py

@@ -115,6 +115,14 @@ def validate_instantiation(component: ComponentModel) -> Optional[ValidationErro
         """Validate that the component can be instantiated"""
         try:
             model = component.model_copy(deep=True)
+
+            # SECURITY: Skip instantiation for FunctionTool to prevent arbitrary code execution.
+            # FunctionTool._from_config() uses exec() on user-provided source_code, which is an RCE vector.
+            # Schema validation is sufficient for FunctionTool - we validate the config structure without
+            # actually executing the code. This blocks drive-by attacks via the /api/validate/ endpoint.
+            if "FunctionTool" in model.provider:
+                return None
+
             # Attempt to load the component
             module_path, class_name = model.provider.rsplit(".", maxsplit=1)
             module = importlib.import_module(module_path)

python/packages/autogen-studio/autogenstudio/web/routes/mcp.py

@@ -1,8 +1,6 @@
-import base64
-import json
 import uuid
 from datetime import datetime, timezone
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from autogen_ext.tools.mcp._config import (
     McpServerParams,
@@ -32,6 +30,11 @@
 # Global session tracking for status endpoint
 active_sessions: Dict[str, Dict[str, Any]] = {}
 
+# Server-side storage for pending MCP session parameters.
+# Params are registered via POST /ws/connect and consumed (popped) when the WebSocket connects.
+# This prevents attackers from injecting arbitrary server_params via the WebSocket query string.
+pending_session_params: Dict[str, Union[StdioServerParams, SseServerParams, StreamableHttpServerParams]] = {}
+
 
 class CreateWebSocketConnectionRequest(BaseModel):
     server_params: McpServerParams
@@ -129,35 +132,19 @@ async def create_mcp_session(bridge: MCPWebSocketBridge, server_params: McpServe
 
 @router.websocket("/ws/{session_id}")
 async def mcp_websocket(websocket: WebSocket, session_id: str):
-    """Main WebSocket endpoint - now a thin layer"""
+    """Main WebSocket endpoint - looks up server params from server-side storage"""
+    # Look up pre-registered server params (one-time use)
+    server_params = pending_session_params.pop(session_id, None)
+    if server_params is None:
+        await websocket.close(code=4004, reason="Unknown or expired session")
+        return
+
     await websocket.accept()
     logger.info(f"MCP WebSocket connection established for session {session_id}")
 
     bridge = None
 
     try:
-        # Parse server parameters
-        query_params = dict(websocket.query_params)
-        server_params_encoded = query_params.get("server_params")
-
-        if not server_params_encoded:
-            await websocket.close(code=4000, reason="Missing server_params")
-            return
-
-        decoded_params = base64.b64decode(server_params_encoded).decode("utf-8")
-        server_params_dict = json.loads(decoded_params)
-
-        # Create appropriate server params object
-        if server_params_dict.get("type") == "StdioServerParams":
-            server_params = StdioServerParams(**server_params_dict)
-        elif server_params_dict.get("type") == "SseServerParams":
-            server_params = SseServerParams(**server_params_dict)
-        elif server_params_dict.get("type") == "StreamableHttpServerParams":
-            server_params = StreamableHttpServerParams(**server_params_dict)
-        else:
-            await websocket.close(code=4000, reason="Invalid server parameters")
-            return
-
         # Create bridge and run MCP session
         bridge = MCPWebSocketBridge(websocket, session_id)
         await create_mcp_session(bridge, server_params, session_id)
@@ -197,18 +184,18 @@ async def mcp_websocket(websocket: WebSocket, session_id: str):
 
 @router.post("/ws/connect")
 async def create_mcp_websocket_connection(request: CreateWebSocketConnectionRequest):
-    """Create WebSocket connection URL"""
+    """Register server params and return a WebSocket URL with session_id only"""
     try:
         session_id = str(uuid.uuid4())
 
-        server_params_json = json.dumps(serialize_for_json(request.server_params.model_dump()))
-        server_params_encoded = base64.b64encode(server_params_json.encode("utf-8")).decode("utf-8")
+        # Store params server-side — WebSocket handler will pop them on connect
+        pending_session_params[session_id] = request.server_params
 
         return {
             "status": True,
             "message": "WebSocket connection URL created",
             "session_id": session_id,
-            "websocket_url": f"/api/mcp/ws/{session_id}?server_params={server_params_encoded}",
+            "websocket_url": f"/api/mcp/ws/{session_id}",
             "timestamp": datetime.now(timezone.utc).isoformat(),
         }
 

python/packages/autogen-studio/frontend/src/components/types/component-templates.ts

@@ -229,35 +229,9 @@ export const MODEL_TEMPLATES: ComponentTemplate<ModelConfig>[] = [
 ];
 
 // Tool Templates
+// NOTE: FunctionTool has been removed due to security concerns (arbitrary code execution via exec()).
+// Use MCP Workbenches instead for custom tool functionality.
 export const TOOL_TEMPLATES: ComponentTemplate<ToolConfig>[] = [
-  {
-    id: "function-tool",
-    label: "Function Tool",
-    description: "A custom Python function that can be called by agents",
-    provider: PROVIDERS.FUNCTION_TOOL,
-    component_type: "tool",
-    version: 1,
-    component_version: 1,
-    config: {
-      name: "my_function",
-      description: "A custom function that performs a specific task",
-      source_code: `def my_function(input_text: str) -> str:
-    """
-    A template function that processes input text.
-    
-    Args:
-        input_text: The text to process
-    
-    Returns:
-        Processed text result
-    """
-    # Replace this with your custom function logic
-    result = f"Processed: {input_text}"
-    return result`,
-      global_imports: [],
-      has_cancellation_support: false,
-    } as FunctionToolConfig,
-  },
   {
     id: "code-execution-tool",
     label: "Code Execution Tool",

python/packages/autogen-studio/frontend/src/components/views/teambuilder/builder/component-editor/fields/agent-fields.tsx

@@ -121,94 +121,8 @@ export const AgentFields: React.FC<AgentFieldsProps> = ({
     [component, handleComponentUpdate]
   );
 
-  const handleAddTool = useCallback(() => {
-    if (!isAssistantAgent(component)) return;
-
-    const blankTool: Component<FunctionToolConfig> = {
-      provider: "autogen_core.tools.FunctionTool",
-      component_type: "tool",
-      version: 1,
-      component_version: 1,
-      description: "Create custom tools by wrapping standard Python functions.",
-      label: "New Tool",
-      config: {
-        source_code: "def new_function():\n    pass",
-        name: "new_function",
-        description: "Description of the new function",
-        global_imports: [],
-        has_cancellation_support: false,
-      },
-    };
-
-    // Get or create workbenches array
-    let workbenches = normalizeWorkbenches(component.config.workbench);
-
-    // Find existing StaticWorkbench or create one
-    let workbenchIndex = workbenches.findIndex((wb) => isStaticWorkbench(wb));
-
-    let workbench: Component<StaticWorkbenchConfig>;
-    if (workbenchIndex === -1) {
-      // Create a new StaticWorkbench
-      workbench = {
-        provider: "autogen_core.tools.StaticWorkbench",
-        component_type: "workbench",
-        config: {
-          tools: [],
-        },
-        label: "Static Workbench",
-      } as Component<StaticWorkbenchConfig>;
-      workbenches = [...workbenches, workbench];
-      workbenchIndex = workbenches.length - 1;
-    } else {
-      workbench = workbenches[
-        workbenchIndex
-      ] as Component<StaticWorkbenchConfig>;
-    }
-
-    const staticConfig = workbench.config as StaticWorkbenchConfig;
-    const currentTools = staticConfig.tools || [];
-    const updatedTools = [...currentTools, blankTool];
-
-    // Update workbench config
-    const updatedWorkbench = {
-      ...workbench,
-      config: {
-        ...staticConfig,
-        tools: updatedTools,
-      },
-    };
-
-    // Update the workbenches array
-    const updatedWorkbenches = [...workbenches];
-    updatedWorkbenches[workbenchIndex] = updatedWorkbench;
-
-    handleConfigUpdate("workbench", updatedWorkbenches);
-
-    // If working copy functionality is available, update that too
-    if (
-      workingCopy &&
-      setWorkingCopy &&
-      updateComponentAtPath &&
-      getCurrentComponent &&
-      editPath
-    ) {
-      const updatedCopy = updateComponentAtPath(workingCopy, editPath, {
-        config: {
-          ...getCurrentComponent(workingCopy)?.config,
-          workbench: updatedWorkbenches,
-        },
-      });
-      setWorkingCopy(updatedCopy);
-    }
-  }, [
-    component,
-    handleConfigUpdate,
-    workingCopy,
-    setWorkingCopy,
-    updateComponentAtPath,
-    getCurrentComponent,
-    editPath,
-  ]);
+  // NOTE: handleAddTool removed - FunctionTool creation is deprecated due to security concerns
+  // (arbitrary code execution via exec()). Users should use MCP Workbenches instead.
 
   // Helper functions to add different types of workbenches
   const addStaticWorkbench = useCallback(() => {