Update: Improved MCP setup and config, async handling for sessions, auto mcp install if present in settings config (on compose)
deci committed
May 13, 2025 at 22:30 UTC
cafcef00e1daeeea270931c4d389168427c49799
5 files changed
+420
-87
docker/run/Dockerfile.cuda
new
+174
@@ -0,0 +1,174 @@
1
+# Use the NVIDIA CUDA base image with Ubuntu
2
+FROM nvidia/cuda:12.8.1-base-ubuntu22.04
3
+
4
+# Set non-interactive installation and timezone
5
+ENV DEBIAN_FRONTEND=noninteractive
6
+ENV TZ=UTC
7
+
8
+# Check if the argument is provided, else throw an error
9
+ARG BRANCH
10
+RUN if [ -z "$BRANCH" ]; then echo "ERROR: BRANCH is not set!" >&2; exit 1; fi
11
+ENV BRANCH=$BRANCH
12
+
13
+# Set locale to en_US.UTF-8 and timezone to UTC (matching main Dockerfile)
14
+RUN apt-get update && apt-get install -y locales tzdata
15
+RUN sed -i -e 's/# \(en_US\.UTF-8 .*\)/\1/' /etc/locale.gen && \
16
+ dpkg-reconfigure --frontend=noninteractive locales && \
17
+ update-locale LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8
18
+RUN ln -sf /usr/share/zoneinfo/UTC /etc/localtime
19
+RUN echo "UTC" > /etc/timezone
20
+RUN dpkg-reconfigure -f noninteractive tzdata
21
+ENV LANG=en_US.UTF-8
22
+ENV LANGUAGE=en_US:en
23
+ENV LC_ALL=en_US.UTF-8
24
+ENV TZ=UTC
25
+
26
+# Copy contents of the project to root directory
27
+COPY ./fs/ /
28
+
29
+# Fix permissions for cron files from pre_install.sh
30
+RUN chmod 0644 /etc/cron.d/*
31
+
32
+# Install essential packages (from pre_install.sh but avoiding supervisor from apt)
33
+# Node.js and npm will be installed separately using NodeSource for a specific version
34
+RUN apt-get update && apt-get upgrade -y && apt-get -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" install -y \
35
+ python3 \
36
+ python3-venv \
37
+ python3-pip \
38
+ openssh-server \
39
+ sudo \
40
+ curl \
41
+ wget \
42
+ git \
43
+ ffmpeg \
44
+ nginx \
45
+ cron \
46
+ libmagic-dev \
47
+ poppler-utils \
48
+ tesseract-ocr \
49
+ qpdf \
50
+ libreoffice \
51
+ pandoc \
52
+ libgtk-3-0 \
53
+ libnss3 \
54
+ libatk1.0-0 \
55
+ libatk-bridge2.0-0 \
56
+ libcups2 \
57
+ libasound2 \
58
+ libasound2-data \
59
+ cargo
60
+
61
+# Install Node.js 20.x (LTS) and a compatible npm
62
+RUN apt-get update && apt-get install -y ca-certificates curl gnupg
63
+RUN mkdir -p /etc/apt/keyrings
64
+RUN curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg
65
+RUN echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" | tee /etc/apt/sources.list.d/nodesource.list
66
+RUN apt-get update && apt-get install -y nodejs || { echo "CRITICAL ERROR: Failed to install Node.js from NodeSource." ; exit 1; }
67
+
68
+# Prepare SSH daemon (from pre_install.sh)
69
+RUN bash /ins/setup_ssh.sh $BRANCH
70
+
71
+# Configure Python 3.12 (specific to Ubuntu-based CUDA image)
72
+RUN apt-get update && apt-get install -y \
73
+ software-properties-common && \
74
+ add-apt-repository -y ppa:deadsnakes/ppa && \
75
+ apt-get update && \
76
+ apt-get install -y python3.12 python3.12-venv python3.12-dev && \
77
+ update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 && \
78
+ update-alternatives --set python3 /usr/bin/python3.12
79
+
80
+# Bootstrap pip for Python 3.12 and install supervisor
81
+RUN python3 -m ensurepip --upgrade && \
82
+ python3 -m pip install --upgrade pip setuptools wheel && \
83
+ python3 -m pip install supervisor
84
+
85
+# Create supervisor directories
86
+RUN mkdir -p /var/log/supervisor /etc/supervisor/conf.d
87
+
88
+# Create root dotfiles with appropriate permissions
89
+RUN touch /root/.bashrc /root/.profile && \
90
+ chmod 644 /root/.bashrc /root/.profile
91
+
92
+# Create a symlink for supervisord to the expected location
93
+RUN ln -sf $(which supervisord) /usr/bin/supervisord
94
+
95
+# Install additional software
96
+RUN bash /ins/install_additional.sh $BRANCH
97
+
98
+# Install core CUDA dependencies (minimized to essentials)
99
+RUN apt-get update && apt-get install -y --no-install-recommends \
100
+ cuda-cudart-dev-12-8 \
101
+ libcublas-dev-12-8 \
102
+ libcudnn8 \
103
+ && apt-get clean \
104
+ && rm -rf /var/lib/apt/lists/*
105
+
106
+# Set CUDA environment variables
107
+ENV PATH=/usr/local/cuda/bin:${PATH}
108
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
109
+ENV CUDA_HOME=/usr/local/cuda
110
+ENV CUDA_VERSION=12.8.1
111
+
112
+# Install A0
113
+RUN bash /ins/install_A0.sh $BRANCH
114
+
115
+# Create and set up the shared instruments virtual environment
116
+ENV INSTRUMENTS_VENV_PATH=/opt/instruments_venv
117
+RUN python3 -m venv $INSTRUMENTS_VENV_PATH
118
+
119
+# Switch to bash for the next RUN step (required for set -o pipefail)
120
+SHELL ["/bin/bash", "-c"]
121
+
122
+# Install all heavy dependencies into the instruments venv, with explicit checks
123
+RUN set -euxo pipefail; \
124
+ . $INSTRUMENTS_VENV_PATH/bin/activate; \
125
+ echo "=== Upgrading pip, setuptools, wheel ==="; \
126
+ pip install --upgrade pip setuptools wheel; \
127
+ echo "=== Installing PyTorch with CUDA ==="; \
128
+ pip install --no-cache-dir torch==2.6.0+cu124 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124; \
129
+ echo "=== Installing other heavy dependencies ==="; \
130
+ pip install --no-cache-dir \
131
+ huggingface-hub==0.20.3 \
132
+ safetensors==0.4.1 \
133
+ accelerate==0.21.0 \
134
+ diffusers==0.25.0 \
135
+ transformers==4.38.2 \
136
+ scipy==1.15.2 \
137
+ xformers==0.0.29.post3; \
138
+ echo "=== Checking all critical imports ==="; \
139
+ python -c "import torch; print(f'PyTorch {torch.__version__} imported successfully. CUDA: {getattr(torch, 'cuda', None) and torch.cuda.is_available()}')" || (echo 'PyTorch import failed!' && exit 1); \
140
+ python -c "import torchvision; print(f'Torchvision {torchvision.__version__} imported successfully.')" || (echo 'Torchvision import failed!' && exit 1); \
141
+ python -c "import torchaudio; print(f'Torchaudio {torchaudio.__version__} imported successfully.')" || (echo 'Torchaudio import failed!' && exit 1); \
142
+ python -c "import diffusers; print(f'Diffusers {diffusers.__version__} imported successfully.')" || (echo 'Diffusers import failed!' && exit 1); \
143
+ python -c "import transformers; print(f'Transformers {transformers.__version__} imported successfully.')" || (echo 'Transformers import failed!' && exit 1); \
144
+ python -c "import xformers; print(f'Xformers {xformers.__version__} imported successfully.')" || (echo 'Xformers import failed!' && exit 1); \
145
+ python -c "import accelerate; print(f'Accelerate {accelerate.__version__} imported successfully.')" || (echo 'Accelerate import failed!' && exit 1); \
146
+ python -c "import safetensors; print(f'Safetensors {safetensors.__version__} imported successfully.')" || (echo 'Safetensors import failed!' && exit 1); \
147
+ python -c "import scipy; print(f'Scipy {scipy.__version__} imported successfully.')" || (echo 'Scipy import failed!' && exit 1); \
148
+ python -c "import huggingface_hub; print(f'Huggingface_hub {huggingface_hub.__version__} imported successfully.')" || (echo 'Huggingface_hub import failed!' && exit 1); \
149
+ echo "=== All critical imports succeeded ==="
150
+
151
+# Optionally revert to sh for subsequent steps if needed
152
+SHELL ["/bin/sh", "-c"]
153
+
154
+# The existing global PyTorch install for /opt/venv (A0's main venv) is below.
155
+# It might be used by other core A0 components, so we keep it.
156
+# Our instruments will use the dedicated $INSTRUMENTS_VENV_PATH.
157
+RUN . /opt/venv/bin/activate && \
158
+ pip uninstall -y torch torchvision && \
159
+ pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124
160
+
161
+# Cleanup repo and install A0 without caching, this speeds up builds
162
+ARG CACHE_DATE=none
163
+RUN echo "cache buster $CACHE_DATE" && bash /ins/install_A02.sh $BRANCH
164
+
165
+# Post installation steps
166
+RUN bash /ins/post_install.sh $BRANCH
167
+
168
+# Expose ports
169
+EXPOSE 22 80
170
+
171
+RUN chmod +x /exe/initialize.sh /exe/run_A0.sh /exe/run_searxng.sh
172
+
173
+# Initialize runtime
174
+CMD ["/exe/initialize.sh", "$BRANCH"]
docker/run/fs/ins/pre_install.sh
+12
-8
@@ -14,8 +14,6 @@ apt-get update && apt-get upgrade -y && apt-get -o Dpkg::Options::="--force-conf
14
python3 \
15
python3-venv \
16
python3-pip \
17
- nodejs \
18
- npm \
17
openssh-server \
18
sudo \
19
curl \
@@ -38,15 +36,21 @@ apt-get update && apt-get upgrade -y && apt-get -o Dpkg::Options::="--force-conf
36
libcups2 \
37
libasound2 \
38
libasound2-data \
41
- cargo
39
+ cargo \
40
+ ca-certificates \
41
+ gnupg
42
43
echo "=====MID UPDATE====="
44
45
-# for some reason npm crashes builds on amd64 in this version and has to be installed separately
46
-# A0 can install it when needed
47
-# The line below is now redundant as npm is included in the main install list above
48
-# apt-get install -y \
49
-# npm
45
+# Install Node.js 20.x (LTS) and a compatible npm using NodeSource
46
+echo "Setting up NodeSource repository for Node.js 20.x..."
47
+mkdir -p /etc/apt/keyrings
48
+curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg
49
+echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" | tee /etc/apt/sources.list.d/nodesource.list
50
+
51
+echo "Installing Node.js..."
52
+apt-get update # Update package list again after adding new source
53
+apt-get install -y nodejs || { echo "CRITICAL ERROR: Failed to install Node.js from NodeSource." ; exit 1; }
54
55
echo "=====AFTER UPDATE====="
56
initialize.py
+53
@@ -2,10 +2,63 @@ import asyncio
2
import models
3
from agent import AgentConfig, ModelConfig
4
from python.helpers import dotenv, files, rfc_exchange, runtime, settings, docker, log
5
+import subprocess
6
+import shutil
7
+from python.helpers.print_style import PrintStyle
8
9
10
def initialize():
11
12
+ PrintStyle(background_color="blue", font_color="white", padding=True).print(
13
+ "Attempting to ensure MCP server 'mcp-server-sequential-thinking' is available..."
14
+ )
15
+ # Check if npm is available first, as it's needed for the install.
16
+ if shutil.which("npm"):
17
+ if not shutil.which("mcp-server-sequential-thinking"):
18
+ PrintStyle(font_color="yellow", padding=True).print(
19
+ "'mcp-server-sequential-thinking' not found in PATH. Attempting global npm install..."
20
+ )
21
+ try:
22
+ # Attempt to install @modelcontextprotocol/server-sequential-thinking globally
23
+ npm_command = ["npm", "i", "-g", "@modelcontextprotocol/server-sequential-thinking", "--no-fund", "--no-audit"]
24
+ process = subprocess.run(npm_command, capture_output=True, text=True, check=False)
25
+ if process.returncode == 0:
26
+ PrintStyle(font_color="green", padding=True).print(
27
+ "Successfully installed @modelcontextprotocol/server-sequential-thinking globally via npm."
28
+ )
29
+ # Re-check if it's available in PATH now. This depends on npm's global bin location being in PATH.
30
+ if shutil.which("mcp-server-sequential-thinking"):
31
+ PrintStyle(font_color="green", padding=True).print(
32
+ "'mcp-server-sequential-thinking' is now available in PATH after install."
33
+ )
34
+ else:
35
+ PrintStyle(font_color="orange", padding=True).print(
36
+ "WARNING: npm install reported success, but 'mcp-server-sequential-thinking' still not found in PATH by shutil.which(). " +
37
+ "The 'npx' command in settings.json might still be necessary and hopefully works."
38
+ )
39
+ else:
40
+ PrintStyle(font_color="red", padding=True).print(
41
+ f"Failed to install @modelcontextprotocol/server-sequential-thinking globally via npm. Return code: {process.returncode}"
42
+ )
43
+ PrintStyle(font_color="red", padding=False).print(f"npm stdout: {process.stdout.strip()}")
44
+ PrintStyle(font_color="red", padding=False).print(f"npm stderr: {process.stderr.strip()}")
45
+ except FileNotFoundError:
46
+ PrintStyle(font_color="red", padding=True).print(
47
+ "ERROR: 'npm' command not found. Cannot attempt to install @modelcontextprotocol/server-sequential-thinking."
48
+ )
49
+ except Exception as e:
50
+ PrintStyle(font_color="red", padding=True).print(
51
+ f"Exception during npm install of @modelcontextprotocol/server-sequential-thinking: {e}"
52
+ )
53
+ else:
54
+ PrintStyle(font_color="green", padding=True).print(
55
+ "'mcp-server-sequential-thinking' already found in PATH."
56
+ )
57
+ else:
58
+ PrintStyle(font_color="red", padding=True).print(
59
+ "ERROR: 'npm' command not found. Cannot check for or install 'mcp-server-sequential-thinking'."
60
+ )
61
+
62
current_settings = settings.get_settings()
63
64
# chat model from user settings
python/helpers/mcp_handler.py
+174
-76
@@ -1,5 +1,5 @@
1
from abc import ABC, abstractmethod
2
-from typing import List, Dict, Optional, Any, Union, Literal, Annotated, ClassVar, cast
2
+from typing import List, Dict, Optional, Any, Union, Literal, Annotated, ClassVar, cast, Callable, Awaitable, TypeVar
3
import threading
4
import asyncio
5
from contextlib import AsyncExitStack
@@ -7,7 +7,7 @@ from shutil import which
7
from datetime import timedelta
8
9
import os
10
-print(f"DEBUG: Listing /opt/venv/lib/python3.11/site-packages/ before mcp import: {os.listdir('/opt/venv/lib/python3.11/site-packages/')}")
10
+# print(f"DEBUG: Listing /opt/venv/lib/python3.11/site-packages/ before mcp import: {os.listdir('/opt/venv/lib/python3.11/site-packages/')}") # This line caused FileNotFoundError, **FOR CUDA CHANGE TO '3.12'**
11
12
from mcp import ClientSession, StdioServerParameters
13
from mcp.client.stdio import stdio_client
@@ -202,21 +202,98 @@ class MCPConfig(BaseModel):
202
203
@classmethod
204
def update(cls, config_str: str) -> Any:
205
- """Parse the MCP config string into a MCPConfig object."""
205
with cls.__lock:
207
- try:
208
- servers = DirtyJson.parse_string(config_str)
209
- except Exception as e:
210
- raise ValueError(f"Failed to parse MCP config: {e}") from e
211
- cls.get_instance().__init__(servers_list=servers)
206
+ servers_data: List[Dict[str, Any]] = [] # Default to empty list
207
+
208
+ if config_str and config_str.strip(): # Only parse if non-empty and not just whitespace
209
+ try:
210
+ # Try with standard json.loads first, as it should handle escaped strings correctly
211
+ import json
212
+ parsed_value = json.loads(config_str)
213
+
214
+ if isinstance(parsed_value, list):
215
+ valid_servers = []
216
+ for item in parsed_value:
217
+ if isinstance(item, dict):
218
+ valid_servers.append(item)
219
+ else:
220
+ PrintStyle(background_color="yellow", font_color="black", padding=True).print(
221
+ f"Warning: MCP config item (from json.loads) was not a dictionary and was ignored: {item}"
222
+ )
223
+ servers_data = valid_servers
224
+ else:
225
+ PrintStyle(background_color="red", font_color="white", padding=True).print(
226
+ f"Error: Parsed MCP config (from json.loads) top-level structure is not a list. Config string was: '{config_str}'"
227
+ )
228
+ # servers_data remains empty
229
+ except Exception as e_json: # Catch json.JSONDecodeError specifically if possible, or general Exception
230
+ # Fallback to DirtyJson or log error if standard json.loads fails
231
+ PrintStyle(background_color="orange", font_color="black", padding=True).print(
232
+ f"Standard json.loads failed for MCP config: {e_json}. Attempting DirtyJson as fallback."
233
+ )
234
+ try:
235
+ parsed_value = DirtyJson.parse_string(config_str)
236
+ if isinstance(parsed_value, list):
237
+ valid_servers = []
238
+ for item in parsed_value:
239
+ if isinstance(item, dict):
240
+ valid_servers.append(item)
241
+ else:
242
+ PrintStyle(background_color="yellow", font_color="black", padding=True).print(
243
+ f"Warning: MCP config item (from DirtyJson) was not a dictionary and was ignored: {item}"
244
+ )
245
+ servers_data = valid_servers
246
+ else:
247
+ PrintStyle(background_color="red", font_color="white", padding=True).print(
248
+ f"Error: Parsed MCP config (from DirtyJson) top-level structure is not a list. Config string was: '{config_str}'"
249
+ )
250
+ # servers_data remains empty
251
+ except Exception as e_dirty:
252
+ PrintStyle(background_color="red", font_color="white", padding=True).print(
253
+ f"Error parsing MCP config string with DirtyJson as well: {e_dirty}. Config string was: '{config_str}'"
254
+ )
255
+ # servers_data remains empty, allowing graceful degradation
256
+
257
+ # Initialize/update the singleton instance with the (potentially empty) list of server data
258
+ instance = cls.get_instance()
259
+ # Directly update the servers attribute of the existing instance or re-initialize carefully
260
+ # For simplicity and to ensure __init__ logic runs if needed for setup:
261
+ new_instance_data = {'servers': servers_data} # Prepare data for re-initialization or update
262
+
263
+ # Option 1: Re-initialize the existing instance (if __init__ is idempotent for other fields)
264
+ instance.__init__(servers_list=servers_data)
265
+
266
+ # Option 2: Or, if __init__ has side effects we don't want to repeat,
267
+ # and 'servers' is the primary thing 'update' changes:
268
+ # instance.servers = [] # Clear existing servers first
269
+ # for server_item_data in servers_data:
270
+ # try:
271
+ # if server_item_data.get("url", None):
272
+ # instance.servers.append(MCPServerRemote(server_item_data))
273
+ # else:
274
+ # instance.servers.append(MCPServerLocal(server_item_data))
275
+ # except Exception as e_init:
276
+ # PrintStyle(background_color="grey", font_color="red", padding=True).print(
277
+ # f"MCPConfig.update: Failed to create MCPServer from item '{server_item_data.get('name', 'Unknown')}': {e_init}"
278
+ # )
279
+
280
cls.__initialized = True
213
- return cls.get_instance()
281
+ return instance
282
283
def __init__(self, servers_list: List[Dict[str, Any]]):
284
from collections.abc import Mapping, Iterable
285
218
- # This empties the servers list
219
- super().__init__()
286
+ # DEBUG: Print the received servers_list
287
+ PrintStyle(background_color="blue", font_color="white", padding=True).print(f"MCPConfig.__init__ received servers_list: {servers_list}")
288
+
289
+ # This empties the servers list if MCPConfig is a Pydantic model and servers is a field.
290
+ # If servers is a field like `servers: List[MCPServer] = Field(default_factory=list)`,
291
+ # then super().__init__() might try to initialize it.
292
+ # We are re-assigning self.servers later in this __init__.
293
+ super().__init__()
294
+
295
+ # Clear any servers potentially initialized by super().__init__() before we populate based on servers_list
296
+ self.servers = []
297
298
if not isinstance(servers_list, Iterable):
299
(
@@ -343,7 +420,7 @@ class MCPConfig(BaseModel):
420
def get_tool(self, agent: Any, tool_name: str) -> MCPTool | None:
421
if not self.has_tool(tool_name):
422
return None
346
- return MCPTool(agent, tool_name, {}, "", **{})
423
+ return MCPTool(agent=agent, name=tool_name, method=None, args={}, message="")
424
425
async def call_tool(self, tool_name: str, input_data: Dict[str, Any]) -> CallToolResult:
426
"""Call a tool with the given input data"""
@@ -357,68 +434,87 @@ class MCPConfig(BaseModel):
434
raise ValueError(f"Tool {tool_name} not found")
435
436
360
-class MCPClientBase(ABC):
361
- session: Optional[ClientSession] = None
362
- exit_stack: AsyncExitStack = AsyncExitStack()
363
- stdio: Optional[MemoryObjectReceiveStream[JSONRPCMessage | Exception]] = None
364
- write: Optional[MemoryObjectSendStream[JSONRPCMessage]] = None
437
+T = TypeVar('T')
438
366
- tools: List[dict[str, Any]] = []
367
- server: Optional[Union[MCPServerLocal, MCPServerRemote]] = None
439
+class MCPClientBase(ABC):
440
+ # server: Union[MCPServerLocal, MCPServerRemote] # Defined in __init__
441
+ # tools: List[dict[str, Any]] # Defined in __init__
442
+ # No self.session, self.exit_stack, self.stdio, self.write as persistent instance fields
443
444
__lock: ClassVar[threading.Lock] = threading.Lock()
445
446
def __init__(self, server: Union[MCPServerLocal, MCPServerRemote]):
447
self.server = server
448
+ self.tools: List[dict[str, Any]] = [] # Tools are cached on the client instance
449
450
# Protected method
451
@abstractmethod
376
- async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
377
- """Connect to an MCP server, init client and save stdio/write streams"""
452
+ async def _create_stdio_transport(self, current_exit_stack: AsyncExitStack) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
453
+ """Create stdio/write streams using the provided exit_stack."""
454
...
455
380
- async def __connect_to_server(self) -> Any:
381
- """Connect to an MCP server"""
382
- with self.__lock:
383
- self.stdio, self.write = await self._connect_client()
384
-
385
- self.session = (
386
- await self.exit_stack.enter_async_context(
456
+ async def _execute_with_session(self, coro_func: Callable[[ClientSession], Awaitable[T]]) -> T:
457
+ """
458
+ Manages the lifecycle of an MCP session for a single operation.
459
+ Creates a temporary session, executes coro_func with it, and ensures cleanup.
460
+ """
461
+ operation_name = coro_func.__name__ # For logging
462
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Creating new session for operation '{operation_name}'...")
463
+ async with AsyncExitStack() as temp_stack:
464
+ try:
465
+ stdio, write = await self._create_stdio_transport(temp_stack)
466
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...")
467
+ session = await temp_stack.enter_async_context(
468
ClientSession(
388
- self.stdio,
389
- self.write,
390
- read_timeout_seconds=timedelta(seconds=15)
469
+ stdio,
470
+ write,
471
+ read_timeout_seconds=timedelta(seconds=600)
472
)
473
)
393
- )
394
-
395
- # Initialize session
396
- await self.session.initialize()
397
- return self
398
-
399
- async def update_tools(self) -> Any:
400
- """List available tools from the server"""
401
- try:
402
- await self.__connect_to_server()
403
- with self.__lock:
404
- response: ListToolsResult = await self.session.list_tools()
405
- available_tools = [{
474
+ await session.initialize()
475
+ PrintStyle(font_color="green").print(f"MCPClientBase ({self.server.name} - {operation_name}): Session initialized.")
476
+
477
+ result = await coro_func(session)
478
+
479
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Operation successful.")
480
+ return result
481
+ except Exception as e:
482
+ PrintStyle(background_color="#AA4455", font_color="white", padding=False).print(
483
+ f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(e).__name__}: {e}"
484
+ )
485
+ raise # Re-raise the exception to be handled by the caller of update_tools/call_tool
486
+ finally:
487
+ PrintStyle(font_color="cyan").print(
488
+ f"MCPClientBase ({self.server.name} - {operation_name}): Session and transport will be closed by AsyncExitStack."
489
+ )
490
+ # temp_stack.aclose() is called automatically here.
491
+
492
+ async def update_tools(self) -> "MCPClientBase":
493
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Starting 'update_tools' operation...")
494
+
495
+ async def list_tools_op(current_session: ClientSession):
496
+ response: ListToolsResult = await current_session.list_tools()
497
+ with self.__lock:
498
+ self.tools = [{
499
"name": tool.name,
500
"description": tool.description,
501
"input_schema": tool.inputSchema
502
} for tool in response.tools]
503
+ PrintStyle(font_color="green").print(f"MCPClientBase ({self.server.name}): Tools updated. Found {len(self.tools)} tools.")
504
411
- self.tools = available_tools
412
- await self.exit_stack.aclose()
413
- return self
505
+ try:
506
+ await self._execute_with_session(list_tools_op)
507
except Exception as e:
415
- PrintStyle(
416
- background_color="#CC34C3", font_color="white", bold=True, padding=True
417
- ).print("MCPClientLocal::Failed to update tools:")
418
- PrintStyle(background_color="#AA4455", font_color="white", padding=False).print(str(e))
508
+ # Error already logged by _execute_with_session, this is for specific handling if needed
509
+ PrintStyle(background_color="#CC34C3", font_color="white", bold=True, padding=True).print(
510
+ f"MCPClientBase ({self.server.name}): 'update_tools' operation failed: {e}"
511
+ )
512
+ with self.__lock:
513
+ self.tools = [] # Ensure tools are cleared on failure
514
+ return self
515
516
def has_tool(self, tool_name: str) -> bool:
421
- """Check if a tool is available"""
517
+ """Check if a tool is available (uses cached tools)"""
518
with self.__lock:
519
for tool in self.tools:
520
if tool["name"] == tool_name:
@@ -426,42 +522,44 @@ class MCPClientBase(ABC):
522
return False
523
524
def get_tools(self) -> List[dict[str, Any]]:
429
- """Get all tools from the server"""
525
+ """Get all tools from the server (uses cached tools)"""
526
with self.__lock:
527
return self.tools
528
529
async def call_tool(self, tool_name: str, input_data: Dict[str, Any]) -> CallToolResult:
434
- """Call a tool with the given input data"""
530
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Preparing for 'call_tool' operation for tool '{tool_name}'.")
531
if not self.has_tool(tool_name):
436
- await self.update_tools()
437
-
438
- await self.__connect_to_server()
532
+ PrintStyle(font_color="orange").print(f"MCPClientBase ({self.server.name}): Tool '{tool_name}' not in cache for 'call_tool', refreshing tools...")
533
+ await self.update_tools() # This will use its own properly managed session
534
+ if not self.has_tool(tool_name):
535
+ PrintStyle(font_color="red").print(f"MCPClientBase ({self.server.name}): Tool '{tool_name}' not found after refresh. Raising ValueError.")
536
+ raise ValueError(f"Tool {tool_name} not found after refreshing tool list for server {self.server.name}.")
537
+ PrintStyle(font_color="green").print(f"MCPClientBase ({self.server.name}): Tool '{tool_name}' found after updating tools.")
538
+
539
+ async def call_tool_op(current_session: ClientSession):
540
+ PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Executing 'call_tool' for '{tool_name}' via MCP session...")
541
+ response: CallToolResult = await current_session.call_tool(tool_name, input_data)
542
+ PrintStyle(font_color="green").print(f"MCPClientBase ({self.server.name}): Tool '{tool_name}' call successful via session.")
543
+ return response
544
440
- with self.__lock:
441
- for tool in self.tools:
442
- if tool["name"] == tool_name:
443
- response: CallToolResult = await self.session.call_tool(tool_name, input_data)
444
- # after connect have to close the stack within this function
445
- await self.exit_stack.aclose()
446
- return response
447
- raise ValueError(f"Tool {tool_name} not found")
545
+ try:
546
+ return await self._execute_with_session(call_tool_op)
547
+ except Exception as e:
548
+ # Error logged by _execute_with_session. Re-raise a specific error for the caller.
549
+ PrintStyle(background_color="#AA4455", font_color="white", padding=True).print(
550
+ f"MCPClientBase ({self.server.name}): 'call_tool' operation for '{tool_name}' failed: {type(e).__name__}: {e}"
551
+ )
552
+ raise ConnectionError(f"MCPClientBase::Failed to call tool '{tool_name}' on server '{self.server.name}'. Original error: {type(e).__name__}: {e}")
553
554
555
class MCPClientLocal(MCPClientBase):
451
- async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
556
+ async def _create_stdio_transport(self, current_exit_stack: AsyncExitStack) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
557
"""Connect to an MCP server, init client and save stdio/write streams"""
558
server: MCPServerLocal = cast(MCPServerLocal, self.server)
559
560
if not which(server.command):
561
raise ValueError(f"Command {server.command} not found")
562
458
- # which_args = 0
459
- # for arg in server.args:
460
- # if which(arg):
461
- # which_args = which_args + 1
462
- # if which_args == 0:
463
- # raise ValueError(f"None of the arguments {server.args} is a file")
464
-
563
server_params = StdioServerParameters(
564
command=server.command,
565
args=server.args,
@@ -469,15 +567,15 @@ class MCPClientLocal(MCPClientBase):
567
encoding=server.encoding,
568
encoding_error_handler=server.encoding_error_handler
569
)
472
- stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
570
+ stdio_transport = await current_exit_stack.enter_async_context(stdio_client(server_params))
571
return stdio_transport
572
573
574
class MCPClientRemote(MCPClientBase):
477
- async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
575
+ async def _create_stdio_transport(self, current_exit_stack: AsyncExitStack) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
576
"""Connect to an MCP server, init client and save stdio/write streams"""
577
server: MCPServerRemote = cast(MCPServerRemote, self.server)
480
- stdio_transport = await self.exit_stack.enter_async_context(
578
+ stdio_transport = await current_exit_stack.enter_async_context(
579
sse_client(url=server.url, headers=server.headers, timeout=server.timeout, sse_read_timeout=server.sse_read_timeout)
580
)
581
return stdio_transport
python/helpers/settings.py
+7
-3
@@ -890,13 +890,16 @@ def _apply_settings(previous: Settings | None):
890
891
async def update_mcp_settings(mcp_servers: str):
892
PrintStyle(background_color="black", font_color="white", padding=True).print("Updating MCP config...")
893
- AgentContext.first().log.log(type="info", content="Updating MCP settings...", temp=True)
893
+ first_context = AgentContext.first()
894
+ if first_context:
895
+ first_context.log.log(type="info", content="Updating MCP settings...", temp=True)
896
897
mcp_config = MCPConfig.get_instance()
898
try:
899
MCPConfig.update(mcp_servers)
900
except Exception as e:
899
- AgentContext.first().log.log(type="warning", content=f"Failed to update MCP settings: {e}", temp=False)
901
+ if first_context:
902
+ first_context.log.log(type="warning", content=f"Failed to update MCP settings: {e}", temp=False)
903
(
904
PrintStyle(background_color="red", font_color="black", padding=True)
905
.print("Failed to update MCP settings")
@@ -913,7 +916,8 @@ def _apply_settings(previous: Settings | None):
916
PrintStyle(background_color="#334455", font_color="white", padding=False)
917
.print(mcp_config.model_dump_json())
918
)
916
- AgentContext.first().log.log(type="info", content="Finished updating MCP settings :)", temp=True)
919
+ if first_context:
920
+ first_context.log.log(type="info", content="Finished updating MCP settings :)", temp=True)
921
922
task2 = defer.DeferredTask().start_task(
923
update_mcp_settings, config.mcp_servers