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

This commit is contained in:
Victor Dibia
2026-03-11 12:42:46 -07:00
committed by GitHub
parent 13e144e547
commit b0477309d2
13 changed files with 137 additions and 497 deletions

View File

@@ -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/

3
.gitignore vendored
View File

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

View File

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

View File

@@ -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"

View File

@@ -288,6 +288,10 @@ class RedisMemory(Memory, Component[RedisMemoryConfig]):
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:

View File

@@ -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(

View File

@@ -115,6 +115,14 @@ class ValidationService:
"""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)

View File

@@ -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 @@ router = APIRouter()
# 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(),
}

View File

@@ -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",

View File

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

View File

@@ -1,319 +1,82 @@
import React, { useCallback, useRef, useState } from "react";
import { Input, Switch, Select, Button, Space, Collapse } from "antd";
import { PlusCircle, MinusCircle, User, Settings, Code } from "lucide-react";
import React from "react";
import { Alert } from "antd";
import { AlertTriangle } from "lucide-react";
import {
Component,
ComponentConfig,
Import,
} from "../../../../../types/datamodel";
import { isFunctionTool } from "../../../../../types/guards";
import { MonacoEditor } from "../../../../monaco";
const { TextArea } = Input;
const { Option } = Select;
interface ToolFieldsProps {
component: Component<ComponentConfig>;
onChange: (updates: Partial<Component<ComponentConfig>>) => void;
}
interface ImportState {
module: string;
imports: string;
}
/**
* ToolFields component - displays a deprecation warning for FunctionTool.
*
* FunctionTool has been deprecated due to security concerns:
* - FunctionTool uses exec() to execute user-provided Python code
* - This creates a Remote Code Execution (RCE) vulnerability
* - The /api/validate/ endpoint could be exploited via drive-by attacks
*
* Users should migrate to MCP Workbenches for custom tool functionality.
* MCP provides better security through process isolation.
*/
export const ToolFields: React.FC<ToolFieldsProps> = ({
component,
onChange,
}) => {
if (!isFunctionTool(component)) return null;
const editorRef = useRef(null);
const [showAddImport, setShowAddImport] = useState(false);
const [importType, setImportType] = useState<"direct" | "fromModule">(
"direct"
);
const [directImport, setDirectImport] = useState("");
const [moduleImport, setModuleImport] = useState<ImportState>({
module: "",
imports: "",
});
const handleComponentUpdate = useCallback(
(updates: Partial<Component<ComponentConfig>>) => {
onChange({
...component,
...updates,
config: {
...component.config,
...(updates.config || {}),
},
});
},
[component, onChange]
);
const formatImport = (imp: Import): string => {
if (!imp) return "";
if (typeof imp === "string") {
return imp;
}
return `from ${imp.module} import ${imp.imports.join(", ")}`;
};
const handleAddImport = () => {
const currentImports = [...(component.config.global_imports || [])];
if (importType === "direct" && directImport) {
currentImports.push(directImport);
setDirectImport("");
} else if (
importType === "fromModule" &&
moduleImport.module &&
moduleImport.imports
) {
currentImports.push({
module: moduleImport.module,
imports: moduleImport.imports
.split(",")
.map((i) => i.trim())
.filter((i) => i),
});
setModuleImport({ module: "", imports: "" });
}
handleComponentUpdate({
config: {
...component.config,
global_imports: currentImports,
},
});
setShowAddImport(false);
};
const handleRemoveImport = (index: number) => {
const newImports = [...(component.config.global_imports || [])];
newImports.splice(index, 1);
handleComponentUpdate({
config: {
...component.config,
global_imports: newImports,
},
});
};
return (
<Collapse
defaultActiveKey={["details", "configuration"]}
className="border-0"
expandIconPosition="end"
items={[
{
key: "details",
label: (
<div className="flex items-center gap-2">
<User className="w-4 h-4 text-blue-500" />
<span className="font-medium">Component Details</span>
<div className="space-y-4 p-4">
<Alert
message="FunctionTool Deprecated"
description={
<div className="space-y-2">
<p>
<strong>FunctionTool has been deprecated due to security concerns.</strong>
</p>
<p>
FunctionTool executes arbitrary Python code, which creates security
vulnerabilities. This component type is no longer supported for new
configurations.
</p>
<p>
<strong>Recommended alternative:</strong> Use an{" "}
<span className="font-mono bg-gray-100 px-1 rounded">MCP Workbench</span>{" "}
instead. MCP (Model Context Protocol) servers provide the same tool
functionality with better security through process isolation.
</p>
<p className="text-sm text-gray-600 mt-2">
Existing FunctionTool configurations in the gallery will continue to
work, but creating or editing FunctionTool source code is no longer
available in the UI.
</p>
</div>
}
type="warning"
showIcon
icon={<AlertTriangle className="w-5 h-5" />}
/>
{/* Show read-only info about the existing tool */}
{component.config.name && (
<div className="bg-gray-50 rounded p-3 space-y-2">
<div>
<span className="text-sm font-medium text-gray-500">Tool Name:</span>
<span className="ml-2 text-sm">{component.config.name}</span>
</div>
{component.config.description && (
<div>
<span className="text-sm font-medium text-gray-500">Description:</span>
<span className="ml-2 text-sm">{component.config.description}</span>
</div>
),
children: (
<div className="space-y-4">
<label className="block">
<span className="text-sm font-medium text-gray-700">Name</span>
<Input
value={component.label || ""}
onChange={(e) =>
handleComponentUpdate({ label: e.target.value })
}
placeholder="Tool name"
className="mt-1"
/>
</label>
<label className="block">
<span className="text-sm font-medium text-gray-700">
Description
</span>
<TextArea
value={component.description || ""}
onChange={(e) =>
handleComponentUpdate({ description: e.target.value })
}
placeholder="Tool description"
rows={4}
className="mt-1"
/>
</label>
</div>
),
},
{
key: "configuration",
label: (
<div className="flex items-center gap-2">
<Settings className="w-4 h-4 text-green-500" />
<span className="font-medium">Tool Configuration</span>
</div>
),
children: (
<div className="space-y-4">
<label className="block">
<span className="text-sm font-medium text-gray-700">
Function Name
</span>
<Input
value={component.config.name || ""}
onChange={(e) =>
handleComponentUpdate({
config: { ...component.config, name: e.target.value },
})
}
placeholder="Function name"
className="mt-1"
/>
</label>
<div className="space-y-2">
<span className="text-sm font-medium text-gray-700">
Global Imports
</span>
<div className="flex flex-wrap gap-2 mt-2">
{(component.config.global_imports || []).map((imp, index) => (
<div
key={index}
className="flex items-center gap-2 bg-tertiary rounded px-2 py-1"
>
<span className="text-sm">{formatImport(imp)}</span>
<Button
type="text"
size="small"
className="flex items-center justify-center h-6 w-6 p-0"
onClick={() => handleRemoveImport(index)}
icon={<MinusCircle className="h-4 w-4" />}
/>
</div>
))}
</div>
{showAddImport ? (
<div className="border rounded p-3 space-y-3">
<Select
value={importType}
onChange={setImportType}
style={{ width: 200 }}
>
<Option value="direct">Direct Import</Option>
<Option value="fromModule">From Module Import</Option>
</Select>
{importType === "direct" ? (
<Space>
<Input
placeholder="Package name (e.g., os)"
className="w-64"
value={directImport}
onChange={(e) => setDirectImport(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter" && directImport) {
handleAddImport();
}
}}
/>
<Button
onClick={handleAddImport}
disabled={!directImport}
>
Add
</Button>
</Space>
) : (
<Space direction="vertical" className="w-full">
<Input
placeholder="Module name (e.g., typing)"
className="w-64"
value={moduleImport.module}
onChange={(e) =>
setModuleImport((prev) => ({
...prev,
module: e.target.value,
}))
}
/>
<Space className="w-full">
<Input
placeholder="Import names (comma-separated)"
className="w-64"
value={moduleImport.imports}
onChange={(e) =>
setModuleImport((prev) => ({
...prev,
imports: e.target.value,
}))
}
/>
<Button
onClick={handleAddImport}
disabled={
!moduleImport.module || !moduleImport.imports
}
>
Add
</Button>
</Space>
</Space>
)}
</div>
) : (
<Button
type="dashed"
onClick={() => setShowAddImport(true)}
className="w-full"
>
<PlusCircle className="h-4 w-4 mr-2" />
Add Import
</Button>
)}
</div>
<label className="block">
<span className="text-sm font-medium text-gray-700">
Source Code
</span>
<div className="mt-1 h-96">
<MonacoEditor
value={component.config.source_code || ""}
editorRef={editorRef}
language="python"
onChange={(value) =>
handleComponentUpdate({
config: { ...component.config, source_code: value },
})
}
/>
</div>
</label>
<div className="flex items-center gap-2">
<span className="text-sm font-medium text-gray-700">
Has Cancellation Support
</span>
<Switch
checked={component.config.has_cancellation_support || false}
onChange={(checked) =>
handleComponentUpdate({
config: {
...component.config,
has_cancellation_support: checked,
},
})
}
/>
</div>
</div>
),
},
]}
/>
)}
</div>
)}
</div>
);
};

View File

@@ -143,32 +143,8 @@ export const WorkbenchFields: React.FC<WorkbenchFieldsProps> = ({
if (isStaticWorkbench(component)) {
const staticConfig = component.config as StaticWorkbenchConfig;
const handleAddTool = () => {
const newTool: Component<FunctionToolConfig> = {
provider: "autogen_core.tools.FunctionTool",
component_type: "tool",
version: 1,
component_version: 1,
label: "New Tool",
description: "A new tool",
config: {
source_code:
'def new_tool():\n """A new tool function"""\n return "Hello from new tool"',
name: "new_tool",
description: "A new tool",
global_imports: [],
has_cancellation_support: false,
},
};
const updatedTools = [...(staticConfig.tools || []), newTool];
handleComponentUpdate({
config: {
...staticConfig,
tools: updatedTools,
},
});
};
// NOTE: handleAddTool removed - FunctionTool creation is deprecated due to security concerns
// (arbitrary code execution via exec()). Users should use MCP Workbenches instead.
const handleUpdateTool = (
index: number,
@@ -321,21 +297,11 @@ export const WorkbenchFields: React.FC<WorkbenchFieldsProps> = ({
),
children: (
<div className="space-y-4">
<div className="flex items-center justify-between">
<div className="flex items-center gap-2 text-sm text-secondary">
<Package className="w-4 h-4" />
<span>
Tools: {staticConfig.tools?.length || 0} configured
</span>
</div>
<Button
type="primary"
size="small"
onClick={handleAddTool}
icon={<PlusCircle className="h-4 w-4" />}
>
Add Tool
</Button>
<div className="flex items-center gap-2 text-sm text-secondary">
<Package className="w-4 h-4" />
<span>
Tools: {staticConfig.tools?.length || 0} configured
</span>
</div>
{staticConfig.tools && staticConfig.tools.length > 0 ? (
@@ -387,10 +353,10 @@ export const WorkbenchFields: React.FC<WorkbenchFieldsProps> = ({
) : (
<div className="text-center text-gray-500 py-8 border-2 border-dashed border-gray-200 rounded-lg">
<Package className="w-12 h-12 mx-auto mb-4 text-gray-400" />
<p className="mb-4">No tools configured</p>
<Button type="dashed" onClick={handleAddTool}>
Add Your First Tool
</Button>
<p className="mb-2">No tools configured</p>
<p className="text-xs text-gray-400">
Use MCP Workbenches for custom tool functionality
</p>
</div>
)}

8
python/uv.lock generated
View File

@@ -1,5 +1,5 @@
version = 1
revision = 3
revision = 2
requires-python = ">=3.10, <3.13"
resolution-markers = [
"python_full_version >= '3.12.4' and sys_platform == 'darwin'",
@@ -778,7 +778,7 @@ requires-dist = [
{ name = "neo4j", marker = "extra == 'mem0-local'", specifier = ">=5.25.0" },
{ name = "ollama", marker = "extra == 'ollama'", specifier = ">=0.4.7" },
{ name = "openai", marker = "extra == 'openai'", specifier = ">=1.93" },
{ name = "openai-whisper", marker = "extra == 'video-surfer'" },
{ name = "openai-whisper", marker = "extra == 'video-surfer'", specifier = ">=20250625" },
{ name = "opencv-python", marker = "extra == 'video-surfer'", specifier = ">=4.5" },
{ name = "pillow", marker = "extra == 'magentic-one'", specifier = ">=11.0.0" },
{ name = "pillow", marker = "extra == 'web-surfer'", specifier = ">=11.0.0" },
@@ -5239,7 +5239,7 @@ wheels = [
[[package]]
name = "openai-whisper"
version = "20240930"
version = "20250625"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "more-itertools" },
@@ -5250,7 +5250,7 @@ dependencies = [
{ name = "tqdm" },
{ name = "triton", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'linux2'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/f5/77/952ca71515f81919bd8a6a4a3f89a27b09e73880cebf90957eda8f2f8545/openai-whisper-20240930.tar.gz", hash = "sha256:b7178e9c1615576807a300024f4daa6353f7e1a815dac5e38c33f1ef055dd2d2", size = 800544, upload-time = "2024-09-30T18:21:22.596Z" }
sdist = { url = "https://files.pythonhosted.org/packages/35/8e/d36f8880bcf18ec026a55807d02fe4c7357da9f25aebd92f85178000c0dc/openai_whisper-20250625.tar.gz", hash = "sha256:37a91a3921809d9f44748ffc73c0a55c9f366c85a3ef5c2ae0cc09540432eb96", size = 803191, upload-time = "2025-06-26T01:06:13.34Z" }
[[package]]
name = "openapi-core"