fix: Improve AutoGen Studio: deprecate FunctionTool, harden MCP WebSocket endpoint (#7362)

Author: victordibia

.github/workflows/docs.yml

@@ -237,7 +237,7 @@ jobs:
       - run: |
           uv venv --python=3.11
           source .venv/bin/activate
-          uv sync --locked --all-extras
+          uv sync --locked
           poe --directory ${{ matrix.version.poe-dir }} docs-build
           mkdir -p docs-staging/${{ matrix.version.dest-dir }}/
           mv ${{ matrix.version.poe-dir }}/docs/build/* docs-staging/${{ matrix.version.dest-dir }}/
@@ -363,7 +363,7 @@ jobs:
         uses: actions/setup-dotnet@v4
         with:
           global-json-file: dotnet/global.json
-      - run: dotnet tool update -g docfx
+      - run: dotnet tool update -g docfx --version 2.67.5
       - run: |
           docfx docs/dotnet/docfx.json
           mkdir -p build/dotnet/

.gitignore

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

docs/dotnet/docfx.json

@@ -21,7 +21,7 @@
       "noRestore": false,
       "namespaceLayout": "flattened",
       "memberLayout": "samePage",
-      "allowCompilationErrors": false
+      "allowCompilationErrors": true
     }
   ],
   "build": {

python/packages/autogen-ext/pyproject.toml

@@ -67,7 +67,7 @@ video-surfer = [
     "autogen-agentchat==0.7.5",
     "opencv-python>=4.5",
     "ffmpeg-python",
-    "openai-whisper",
+    "openai-whisper>=20250625",
 ]
 diskcache = [
     "diskcache>=5.6.3"

python/packages/autogen-ext/src/autogen_ext/memory/redis/_redis_memory.py

@@ -288,6 +288,10 @@ async def query(
         top_k = kwargs.pop("top_k", self.config.top_k)
         distance_threshold = kwargs.pop("distance_threshold", self.config.distance_threshold)
 
+        # return empty results for empty/whitespace queries
+        if isinstance(query, str) and not query.strip():
+            return MemoryQueryResult(results=[])
+
         # if sequential memory is requested skip prompt creation
         sequential = bool(kwargs.pop("sequential", self.config.sequential))
         if self.config.sequential and not sequential:

python/packages/autogen-ext/tests/memory/test_redis_memory.py

@@ -71,6 +71,26 @@ async def test_redis_memory_close_with_mock() -> None:
         mock_history.delete.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_redis_memory_query_empty_string_with_mock() -> None:
+    with patch("autogen_ext.memory.redis._redis_memory.SemanticMessageHistory") as MockHistory:
+        mock_history = MagicMock()
+        MockHistory.return_value = mock_history
+
+        config = RedisMemoryConfig()
+        memory = RedisMemory(config=config)
+
+        # Empty string should return empty results without calling the vectorizer
+        result = await memory.query("")
+        assert result.results == []
+        mock_history.get_relevant.assert_not_called()
+
+        # Whitespace-only string should also return empty results
+        result = await memory.query("   ")
+        assert result.results == []
+        mock_history.get_relevant.assert_not_called()
+
+
 def redis_available() -> bool:
     try:
         client = Redis.from_url("redis://localhost:6379")  # type: ignore[reportUnkownMemberType]
@@ -465,8 +485,9 @@ async def test_markdown_memory_type(semantic_memory: RedisMemory) -> None:
     results = await semantic_memory.query("how can I make itemized lists, or italicize text with asterisks?")
     assert results.results[0].content == markdown_data
 
-    # test we can query with markdown interpreted as a text string also
+    # empty query should return empty results without error
     results = await semantic_memory.query("")
+    assert results.results == []
 
     # we can also if the markdown is within a MemoryContent container
     results = await semantic_memory.query(

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(() => {