UI: Add MCP server support

modules/chat.py

@@ -1264,14 +1264,23 @@ def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False):
 
     # Load tools if any are selected
     selected = state.get('selected_tools', [])
+    mcp_servers = state.get('mcp_servers', '')
     parse_tool_call = None
     _tool_parsers = None
-    if selected:
-        from modules.tool_use import load_tools, execute_tool
+    if selected or mcp_servers:
+        from modules.tool_use import load_tools, load_mcp_tools, execute_tool
         from modules.tool_parsing import parse_tool_call, get_tool_call_id, detect_tool_call_format
 
-    if selected:
         tool_defs, tool_executors = load_tools(selected)
+        if mcp_servers:
+            mcp_defs, mcp_executors = load_mcp_tools(mcp_servers)
+            for td in mcp_defs:
+                fn = td['function']['name']
+                if fn in tool_executors:
+                    logger.warning(f'MCP tool "{fn}" conflicts with a local tool. Skipping.')
+                    continue
+                tool_defs.append(td)
+                tool_executors[fn] = mcp_executors[fn]
         state['tools'] = tool_defs
         tool_func_names = [t['function']['name'] for t in tool_defs]
         _template_str = state.get('instruction_template_str', '') if state.get('mode') == 'instruct' else state.get('chat_template_str', '')

modules/shared.py

@@ -259,6 +259,7 @@ settings = {
     'enable_web_search': False,
     'web_search_pages': 3,
     'selected_tools': [],
+    'mcp_servers': '',
     'prompt-notebook': '',
     'preset': 'Top-P' if (user_data_dir / 'presets/Top-P.yaml').exists() else None,
     'max_new_tokens': 512,

modules/tool_use.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib.util
 import json
 
@@ -55,6 +56,119 @@ def load_tools(selected_names):
     return tool_defs, executors
 
 
+def _parse_mcp_servers(servers_str):
+    """Parse MCP servers textbox: one server per line, format 'url' or 'url,Header: value,Header2: value2'."""
+    servers = []
+    for line in servers_str.strip().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        parts = line.split(',')
+        url = parts[0].strip()
+        headers = {}
+        for part in parts[1:]:
+            part = part.strip()
+            if ':' in part:
+                key, val = part.split(':', 1)
+                headers[key.strip()] = val.strip()
+        servers.append((url, headers))
+    return servers
+
+
+def _mcp_tool_to_openai(tool):
+    """Convert an MCP Tool object to OpenAI-format tool dict."""
+    return {
+        "type": "function",
+        "function": {
+            "name": tool.name,
+            "description": tool.description or "",
+            "parameters": tool.inputSchema or {"type": "object", "properties": {}}
+        }
+    }
+
+
+async def _mcp_session(url, headers, callback):
+    """Open an MCP session and pass it to the callback."""
+    from mcp.client.streamable_http import streamablehttp_client
+    from mcp import ClientSession
+
+    async with streamablehttp_client(url, headers=headers or None) as (read_stream, write_stream, _):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+            return await callback(session)
+
+
+def _make_mcp_executor(name, url, headers):
+    def executor(arguments):
+        return asyncio.run(_call_mcp_tool(name, arguments, url, headers))
+    return executor
+
+
+async def _connect_mcp_server(url, headers):
+    """Connect to one MCP server and return (tool_defs, executors)."""
+
+    async def _discover(session):
+        result = await session.list_tools()
+        tool_defs = []
+        executors = {}
+        for tool in result.tools:
+            tool_defs.append(_mcp_tool_to_openai(tool))
+            executors[tool.name] = _make_mcp_executor(tool.name, url, headers)
+        return tool_defs, executors
+
+    return await _mcp_session(url, headers, _discover)
+
+
+async def _call_mcp_tool(name, arguments, url, headers):
+    """Connect to an MCP server and call a single tool."""
+
+    async def _invoke(session):
+        result = await session.call_tool(name, arguments)
+        parts = []
+        for content in result.content:
+            if hasattr(content, 'text'):
+                parts.append(content.text)
+            else:
+                parts.append(str(content))
+        return '\n'.join(parts) if parts else ''
+
+    return await _mcp_session(url, headers, _invoke)
+
+
+async def _connect_all_mcp_servers(servers):
+    """Connect to all MCP servers concurrently."""
+    results = await asyncio.gather(
+        *(_connect_mcp_server(url, headers) for url, headers in servers),
+        return_exceptions=True
+    )
+    all_defs = []
+    all_executors = {}
+    for (url, _), result in zip(servers, results):
+        if isinstance(result, Exception):
+            logger.exception(f'Failed to connect to MCP server "{url}"', exc_info=result)
+            continue
+        defs, execs = result
+        for td, (fn, ex) in zip(defs, execs.items()):
+            if fn in all_executors:
+                logger.warning(f'MCP tool "{fn}" from {url} conflicts with an already loaded tool. Skipping.')
+                continue
+            all_defs.append(td)
+            all_executors[fn] = ex
+    return all_defs, all_executors
+
+
+def load_mcp_tools(servers_str):
+    """
+    Parse MCP servers string and discover tools from each server.
+    Returns (tool_defs, executors) in the same format as load_tools.
+    """
+    servers = _parse_mcp_servers(servers_str)
+    if not servers:
+        return [], {}
+
+    return asyncio.run(_connect_all_mcp_servers(servers))
+
+
 def execute_tool(func_name, arguments, executors):
     """Execute a tool by function name. Returns result as a JSON string."""
     fn = executors.get(func_name)

modules/ui.py

@@ -209,6 +209,7 @@ def list_interface_input_elements():
         'textbox',
         'start_with',
         'selected_tools',
+        'mcp_servers',
         'mode',
         'chat_style',
         'chat-instruct_command',
@@ -434,6 +435,7 @@ def setup_auto_save():
         'custom_system_message',
         'chat_template_str',
         'selected_tools',
+        'mcp_servers',
 
         # Parameters tab (ui_parameters.py) - Generation parameters
         'preset_menu',

modules/ui_chat.py

@@ -105,6 +105,9 @@ def create_ui():
 
                 shared.gradio['selected_tools'].change(fn=sync_web_tools, inputs=[shared.gradio['selected_tools']], outputs=[shared.gradio['selected_tools']], show_progress=False)
 
+                with gr.Accordion('MCP servers', open=False):
+                    shared.gradio['mcp_servers'] = gr.Textbox(value=shared.settings.get('mcp_servers', ''), lines=3, max_lines=3, label='', info='One url per line. For headers, write url,Header: value,Header2: value2', elem_classes=['add_scrollbar'])
+
                 gr.HTML("<div class='sidebar-vertical-separator'></div>")
 
                 with gr.Row():