| 1 | from abc import ABC, abstractmethod |
| 2 | import re |
| 3 | from typing import ( |
| 4 | List, |
| 5 | Dict, |
| 6 | Optional, |
| 7 | Any, |
| 8 | TextIO, |
| 9 | Union, |
| 10 | Literal, |
| 11 | Annotated, |
| 12 | ClassVar, |
| 13 | cast, |
| 14 | Callable, |
| 15 | Awaitable, |
| 16 | TypeVar, |
| 17 | ) |
| 18 | import threading |
| 19 | import asyncio |
| 20 | from contextlib import AsyncExitStack |
| 21 | from shutil import which |
| 22 | from datetime import timedelta |
| 23 | import json |
| 24 | import shlex |
| 25 | import uuid |
| 26 | from helpers import errors |
| 27 | from helpers import settings |
| 28 | from helpers.log import LogItem |
| 29 | |
| 30 | import httpx |
| 31 | |
| 32 | from mcp import ClientSession, StdioServerParameters |
| 33 | from mcp.client.stdio import stdio_client |
| 34 | from mcp.client.sse import sse_client |
| 35 | from mcp.client.streamable_http import streamablehttp_client |
| 36 | from mcp.shared.message import SessionMessage |
| 37 | from mcp.types import CallToolResult, ListToolsResult |
| 38 | from anyio.streams.memory import ( |
| 39 | MemoryObjectReceiveStream, |
| 40 | MemoryObjectSendStream, |
| 41 | ) |
| 42 | |
| 43 | from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr |
| 44 | from helpers import dirty_json, media_artifacts |
| 45 | from helpers.print_style import PrintStyle |
| 46 | from helpers.tool import Tool, Response |
| 47 | from helpers.defer import DeferredTask |
| 48 | from helpers.responses_tools import original_tool_name |
| 49 | |
| 50 | |
| 51 | MCP_MEDIA_TOKENS_ESTIMATE = 1500 |
| 52 | MAX_MCP_RESOURCE_TEXT_CHARS = 12_000 |
| 53 | MCP_SESSION_CLEANUP_TIMEOUT_SECONDS = 5.0 |
| 54 | MCP_OPERATION_TIMEOUT_GRACE_SECONDS = MCP_SESSION_CLEANUP_TIMEOUT_SECONDS + 2.0 |
| 55 | DEFAULT_MCP_SERVERS_CONFIG = '{\n "mcpServers": {}\n}' |
| 56 | |
| 57 | |
| 58 | def _mcp_get(item: Any, key: str, default: Any = None) -> Any: |
| 59 | if isinstance(item, dict): |
| 60 | return item.get(key, default) |
| 61 | return getattr(item, key, default) |
| 62 | |
| 63 | |
| 64 | def normalize_name(name: str) -> str: |
| 65 | # Lowercase and strip whitespace |
| 66 | name = name.strip().lower() |
| 67 | # Replace all non-alphanumeric (unicode) chars with underscore |
| 68 | # \W matches non-alphanumeric, but also matches underscore, so use [^\w] with re.UNICODE |
| 69 | # To also replace underscores from non-latin chars, use [^a-zA-Z0-9] with re.UNICODE |
| 70 | name = re.sub(r"[^\w]", "_", name, flags=re.UNICODE) |
| 71 | return name |
| 72 | |
| 73 | |
| 74 | def _determine_server_type(config_dict: dict) -> str: |
| 75 | """Determine the server type based on configuration, with backward compatibility.""" |
| 76 | # First check if type is explicitly specified |
| 77 | if "type" in config_dict: |
| 78 | server_type = config_dict["type"].lower() |
| 79 | if server_type in ["sse", "http-stream", "streaming-http", "streamable-http", "http-streaming"]: |
| 80 | return "MCPServerRemote" |
| 81 | elif server_type == "stdio": |
| 82 | return "MCPServerLocal" |
| 83 | # For future types, we could add more cases here |
| 84 | else: |
| 85 | # For unknown types, fall back to URL-based detection |
| 86 | # This allows for graceful handling of new types |
| 87 | pass |
| 88 | |
| 89 | # Backward compatibility: if no type specified, use URL-based detection |
| 90 | if "url" in config_dict or "serverUrl" in config_dict: |
| 91 | return "MCPServerRemote" |
| 92 | else: |
| 93 | return "MCPServerLocal" |
| 94 | |
| 95 | |
| 96 | def _is_streaming_http_type(server_type: str) -> bool: |
| 97 | """Check if the server type is a streaming HTTP variant.""" |
| 98 | return server_type.lower() in ["http-stream", "streaming-http", "streamable-http", "http-streaming"] |
| 99 | |
| 100 | |
| 101 | def _split_qualified_tool_name(tool_name: str) -> tuple[str, str]: |
| 102 | """Split Agent Zero's server.tool MCP name while preserving dots in MCP tool names.""" |
| 103 | if "." not in tool_name: |
| 104 | raise ValueError(f"Tool {tool_name} not found") |
| 105 | server_name_part, tool_name_part = tool_name.split(".", 1) |
| 106 | if not server_name_part or not tool_name_part: |
| 107 | raise ValueError(f"Tool {tool_name} not found") |
| 108 | return server_name_part, tool_name_part |
| 109 | |
| 110 | |
| 111 | def _normalize_disabled_tools(value: Any) -> list[str]: |
| 112 | if not isinstance(value, list): |
| 113 | return [] |
| 114 | return [str(item).strip() for item in value if str(item).strip()] |
| 115 | |
| 116 | |
| 117 | def _split_stdio_command(command: Any) -> tuple[str, list[str]]: |
| 118 | text = str(command or "").strip() |
| 119 | if not text: |
| 120 | return "", [] |
| 121 | try: |
| 122 | parts = shlex.split(text) |
| 123 | except ValueError: |
| 124 | return text, [] |
| 125 | if not parts: |
| 126 | return "", [] |
| 127 | return parts[0], parts[1:] |
| 128 | |
| 129 | |
| 130 | def _split_stdio_arg_fragment(arg: str) -> list[str]: |
| 131 | try: |
| 132 | parts = shlex.split(arg) |
| 133 | except ValueError: |
| 134 | return [arg] |
| 135 | if len(parts) <= 1: |
| 136 | return parts or [] |
| 137 | if parts[0].startswith("-") and "=" not in parts[0]: |
| 138 | return parts |
| 139 | if any(part.startswith("-") for part in parts[1:]): |
| 140 | return parts |
| 141 | return [arg] |
| 142 | |
| 143 | |
| 144 | def _normalize_stdio_args(value: Any) -> list[str]: |
| 145 | if not isinstance(value, list): |
| 146 | return [] |
| 147 | args: list[str] = [] |
| 148 | for item in value: |
| 149 | text = str(item).strip() |
| 150 | if text: |
| 151 | args.extend(_split_stdio_arg_fragment(text)) |
| 152 | return args |
| 153 | |
| 154 | |
| 155 | def initialize_mcp(mcp_servers_config: str): |
| 156 | if not MCPConfig.get_instance().is_initialized(): |
| 157 | try: |
| 158 | MCPConfig.update(mcp_servers_config) |
| 159 | except Exception as e: |
| 160 | from agent import AgentContext |
| 161 | |
| 162 | AgentContext.log_to_all( |
| 163 | type="warning", |
| 164 | content=f"Failed to update MCP settings: {e}", |
| 165 | ) |
| 166 | |
| 167 | PrintStyle( |
| 168 | background_color="black", font_color="red", padding=True |
| 169 | ).print(f"Failed to update MCP settings: {e}") |
| 170 | |
| 171 | |
| 172 | class MCPTool(Tool): |
| 173 | """MCP Tool wrapper""" |
| 174 | |
| 175 | def get_log_object(self) -> LogItem: |
| 176 | return self.agent.context.log.log( |
| 177 | type="mcp", |
| 178 | heading=f"icon://extension {self.agent.agent_name}: Using MCP tool '{self.name}'", |
| 179 | content="", |
| 180 | kvps={"tool_name": self.name, **self.args}, |
| 181 | id=str(uuid.uuid4()), |
| 182 | ) |
| 183 | |
| 184 | def _context_id(self) -> str: |
| 185 | return str(getattr(getattr(self.agent, "context", None), "id", "") or "").strip() |
| 186 | |
| 187 | def _raw_tool_response(self, response: Response) -> str: |
| 188 | raw_tool_response = response.message.strip() if response.message else "" |
| 189 | if not raw_tool_response: |
| 190 | PrintStyle(font_color="red").print( |
| 191 | f"Warning: Tool '{self.name}' returned an empty message." |
| 192 | ) |
| 193 | raw_tool_response = "[Tool returned no textual content]" |
| 194 | return raw_tool_response |
| 195 | |
| 196 | def _coerce_media_token_estimate(self, value: object) -> int: |
| 197 | try: |
| 198 | estimate = int(value or 0) |
| 199 | except (TypeError, ValueError): |
| 200 | estimate = 0 |
| 201 | return estimate if estimate > 0 else MCP_MEDIA_TOKENS_ESTIMATE |
| 202 | |
| 203 | def _format_image_content( |
| 204 | self, |
| 205 | *, |
| 206 | encoded: str, |
| 207 | mime_type: str, |
| 208 | label: str, |
| 209 | index: int, |
| 210 | preferred_name: str = "", |
| 211 | ) -> tuple[str, dict[str, Any] | None, str]: |
| 212 | try: |
| 213 | safe_mime = media_artifacts.normalize_mime( |
| 214 | mime_type, |
| 215 | default="image/png", |
| 216 | required_prefix="image/", |
| 217 | ) |
| 218 | artifact = media_artifacts.save_base64_artifact( |
| 219 | encoded, |
| 220 | mime_type=safe_mime, |
| 221 | directory_parts=self._artifact_directory_parts(), |
| 222 | preferred_name=preferred_name, |
| 223 | default_filename=self._default_artifact_filename( |
| 224 | label=label, |
| 225 | index=index, |
| 226 | mime_type=safe_mime, |
| 227 | ), |
| 228 | ) |
| 229 | except media_artifacts.EmptyBase64Data: |
| 230 | return f"MCP returned an empty {label} attachment.", None, "" |
| 231 | except media_artifacts.InvalidBase64Data: |
| 232 | return ( |
| 233 | f"MCP returned a {label} attachment that could not be decoded.", |
| 234 | None, |
| 235 | "", |
| 236 | ) |
| 237 | |
| 238 | return ( |
| 239 | ( |
| 240 | f"Saved MCP {label} attachment " |
| 241 | f"({artifact.mime}, {artifact.size} bytes) to {artifact.path}." |
| 242 | ), |
| 243 | { |
| 244 | "type": "image_url", |
| 245 | "image_url": {"url": artifact.path}, |
| 246 | }, |
| 247 | artifact.path, |
| 248 | ) |
| 249 | |
| 250 | def _materialize_binary_content( |
| 251 | self, |
| 252 | *, |
| 253 | encoded: str, |
| 254 | mime_type: str, |
| 255 | label: str, |
| 256 | index: int, |
| 257 | preferred_name: str = "", |
| 258 | ) -> str: |
| 259 | try: |
| 260 | safe_mime = media_artifacts.normalize_mime(mime_type) |
| 261 | artifact = media_artifacts.save_base64_artifact( |
| 262 | encoded, |
| 263 | mime_type=safe_mime, |
| 264 | directory_parts=self._artifact_directory_parts(), |
| 265 | preferred_name=preferred_name, |
| 266 | default_filename=self._default_artifact_filename( |
| 267 | label=label, |
| 268 | index=index, |
| 269 | mime_type=safe_mime, |
| 270 | ), |
| 271 | ) |
| 272 | except media_artifacts.EmptyBase64Data: |
| 273 | return f"MCP returned an empty {label} attachment." |
| 274 | except media_artifacts.InvalidBase64Data: |
| 275 | return f"MCP returned a {label} attachment that could not be decoded." |
| 276 | |
| 277 | return f"Saved MCP {label} attachment ({artifact.mime}, {artifact.size} bytes) to {artifact.path}." |
| 278 | |
| 279 | def _artifact_directory_parts(self) -> tuple[str, ...]: |
| 280 | context_id = normalize_name(self._context_id() or "shared") or "shared" |
| 281 | tool_name = normalize_name(self.name or "mcp_tool") or "mcp_tool" |
| 282 | return ("tmp", "mcp", context_id, tool_name) |
| 283 | |
| 284 | def _default_artifact_filename(self, *, label: str, index: int, mime_type: str) -> str: |
| 285 | tool_name = normalize_name(self.name or "mcp_tool") or "mcp_tool" |
| 286 | ext = media_artifacts.guess_extension(mime_type, ".bin") |
| 287 | return f"{tool_name}_{label}_{index}{ext}" |
| 288 | |
| 289 | def _format_resource_text(self, text: str, uri: str = "") -> str: |
| 290 | body = str(text or "").strip() |
| 291 | if not body: |
| 292 | return "" |
| 293 | if len(body) > MAX_MCP_RESOURCE_TEXT_CHARS: |
| 294 | body = body[:MAX_MCP_RESOURCE_TEXT_CHARS].rstrip() + "\n...[truncated]" |
| 295 | if uri: |
| 296 | return f"Resource {uri}:\n{body}" |
| 297 | return body |
| 298 | |
| 299 | def _content_item_dump(self, item: Any) -> dict[str, Any]: |
| 300 | if isinstance(item, dict): |
| 301 | return dict(item) |
| 302 | model_dump = getattr(item, "model_dump", None) |
| 303 | if callable(model_dump): |
| 304 | dumped = model_dump(mode="python") |
| 305 | if isinstance(dumped, dict): |
| 306 | return dumped |
| 307 | item_vars = getattr(item, "__dict__", None) |
| 308 | if isinstance(item_vars, dict): |
| 309 | return dict(item_vars) |
| 310 | return {} |
| 311 | |
| 312 | def _summarize_unknown_item(self, item: Any, item_type: str) -> str: |
| 313 | dumped = self._content_item_dump(item) |
| 314 | if dumped: |
| 315 | dumped.pop("data", None) |
| 316 | resource = dumped.get("resource") |
| 317 | if isinstance(resource, dict): |
| 318 | resource.pop("blob", None) |
| 319 | summary = json.dumps(dumped, ensure_ascii=False) |
| 320 | if len(summary) > 600: |
| 321 | summary = summary[:600] + "...[truncated]" |
| 322 | return f"MCP returned unsupported content item type '{item_type}': {summary}" |
| 323 | return f"MCP returned unsupported content item type '{item_type}'." |
| 324 | |
| 325 | def _format_tool_result( |
| 326 | self, response: CallToolResult |
| 327 | ) -> tuple[str, dict[str, Any] | None]: |
| 328 | text_parts: list[str] = [] |
| 329 | notes: list[str] = [] |
| 330 | raw_images: list[dict[str, Any]] = [] |
| 331 | image_paths: list[str] = [] |
| 332 | content_items = list(getattr(response, "content", []) or []) |
| 333 | |
| 334 | for index, item in enumerate(content_items, start=1): |
| 335 | item_type = str(_mcp_get(item, "type", "") or "").strip().lower() |
| 336 | |
| 337 | if item_type == "text": |
| 338 | text = str(_mcp_get(item, "text", "") or "").strip() |
| 339 | if text: |
| 340 | text_parts.append(text) |
| 341 | continue |
| 342 | |
| 343 | if item_type == "image": |
| 344 | note, raw_content, path = self._format_image_content( |
| 345 | encoded=str(_mcp_get(item, "data", "") or ""), |
| 346 | mime_type=str(_mcp_get(item, "mimeType", "") or "image/png"), |
| 347 | label="image", |
| 348 | index=index, |
| 349 | ) |
| 350 | notes.append(note) |
| 351 | if raw_content: |
| 352 | raw_images.append(raw_content) |
| 353 | if path: |
| 354 | image_paths.append(path) |
| 355 | continue |
| 356 | |
| 357 | if item_type == "audio": |
| 358 | note = self._materialize_binary_content( |
| 359 | encoded=str(_mcp_get(item, "data", "") or ""), |
| 360 | mime_type=str(_mcp_get(item, "mimeType", "") or "audio/wav"), |
| 361 | label="audio", |
| 362 | index=index, |
| 363 | ) |
| 364 | notes.append(note) |
| 365 | continue |
| 366 | |
| 367 | if item_type == "resource": |
| 368 | resource = _mcp_get(item, "resource", None) |
| 369 | uri = str(_mcp_get(resource, "uri", "") or "").strip() |
| 370 | text = _mcp_get(resource, "text", None) |
| 371 | if isinstance(text, str) and text.strip(): |
| 372 | text_parts.append(self._format_resource_text(text, uri)) |
| 373 | continue |
| 374 | |
| 375 | blob = str(_mcp_get(resource, "blob", "") or "").strip() |
| 376 | if blob: |
| 377 | mime_type = str( |
| 378 | _mcp_get(resource, "mimeType", "") or "application/octet-stream" |
| 379 | ).strip().lower() |
| 380 | if mime_type.startswith("image/"): |
| 381 | note, raw_content, path = self._format_image_content( |
| 382 | encoded=blob, |
| 383 | mime_type=mime_type, |
| 384 | label="resource image", |
| 385 | index=index, |
| 386 | preferred_name=uri, |
| 387 | ) |
| 388 | else: |
| 389 | note = self._materialize_binary_content( |
| 390 | encoded=blob, |
| 391 | mime_type=mime_type, |
| 392 | label="resource", |
| 393 | index=index, |
| 394 | preferred_name=uri, |
| 395 | ) |
| 396 | raw_content = None |
| 397 | path = "" |
| 398 | notes.append(note) |
| 399 | if raw_content: |
| 400 | raw_images.append(raw_content) |
| 401 | if path: |
| 402 | image_paths.append(path) |
| 403 | continue |
| 404 | |
| 405 | if uri: |
| 406 | mime_type = str(_mcp_get(resource, "mimeType", "") or "").strip() |
| 407 | details = f" ({mime_type})" if mime_type else "" |
| 408 | notes.append(f"MCP returned a resource reference: {uri}{details}.") |
| 409 | continue |
| 410 | |
| 411 | notes.append("MCP returned a resource item without text or binary data.") |
| 412 | continue |
| 413 | |
| 414 | if item_type: |
| 415 | notes.append(self._summarize_unknown_item(item, item_type)) |
| 416 | continue |
| 417 | |
| 418 | notes.append(self._summarize_unknown_item(item, "unknown")) |
| 419 | |
| 420 | message = "\n\n".join(part for part in [*text_parts, *notes] if part.strip()) |
| 421 | if not message and content_items: |
| 422 | message = "MCP tool returned content that could not be rendered as text." |
| 423 | |
| 424 | additional = None |
| 425 | if raw_images: |
| 426 | additional = { |
| 427 | "raw_content": raw_images, |
| 428 | "preview": f"<MCP image attachments: {len(raw_images)}>", |
| 429 | "_tokens": MCP_MEDIA_TOKENS_ESTIMATE * len(raw_images), |
| 430 | "attachments": image_paths, |
| 431 | "media_paths": image_paths, |
| 432 | } |
| 433 | |
| 434 | return message, additional |
| 435 | |
| 436 | async def execute(self, **kwargs: Any): |
| 437 | from helpers.tool_policy import canonical_mcp_id, ensure_tool_allowed |
| 438 | |
| 439 | if "." in self.name: |
| 440 | ensure_tool_allowed( |
| 441 | self.agent, |
| 442 | self.name, |
| 443 | canonical_id=canonical_mcp_id(self.name), |
| 444 | ) |
| 445 | error = "" |
| 446 | additional: dict[str, Any] | None = None |
| 447 | try: |
| 448 | response: CallToolResult = await MCPConfig.get_for_agent(self.agent).call_tool( |
| 449 | self.name, kwargs |
| 450 | ) |
| 451 | message, additional = self._format_tool_result(response) |
| 452 | if response.isError: |
| 453 | error = message or "MCP tool returned an error without textual content." |
| 454 | except Exception as e: |
| 455 | error = f"MCP Tool Exception: {str(e)}" |
| 456 | message = f"ERROR: {str(e)}" |
| 457 | |
| 458 | if error: |
| 459 | PrintStyle( |
| 460 | background_color="#CC34C3", font_color="white", bold=True, padding=True |
| 461 | ).print(f"MCPTool::Failed to call mcp tool {self.name}:") |
| 462 | PrintStyle( |
| 463 | background_color="#AA4455", font_color="white", padding=False |
| 464 | ).print(error) |
| 465 | |
| 466 | self.agent.context.log.log( |
| 467 | type="warning", |
| 468 | content=f"{self.name}: {error}", |
| 469 | ) |
| 470 | |
| 471 | return Response(message=message, break_loop=False, additional=additional) |
| 472 | |
| 473 | async def before_execution(self, **kwargs: Any): |
| 474 | ( |
| 475 | PrintStyle( |
| 476 | font_color="#1B4F72", padding=True, background_color="white", bold=True |
| 477 | ).print(f"{self.agent.agent_name}: Using tool '{self.name}'") |
| 478 | ) |
| 479 | self.log = self.get_log_object() |
| 480 | |
| 481 | for key, value in self.args.items(): |
| 482 | PrintStyle(font_color="#85C1E9", bold=True).stream( |
| 483 | self.nice_key(key) + ": " |
| 484 | ) |
| 485 | PrintStyle( |
| 486 | font_color="#85C1E9", padding=isinstance(value, str) and "\n" in value |
| 487 | ).stream(value) |
| 488 | PrintStyle().print() |
| 489 | |
| 490 | async def after_execution(self, response: Response, **kwargs: Any): |
| 491 | final_text_for_agent = self._raw_tool_response(response) |
| 492 | additional = dict(response.additional or {}) |
| 493 | raw_content = additional.pop("raw_content", None) |
| 494 | preview = str(additional.pop("preview", "") or "").strip() |
| 495 | token_estimate = self._coerce_media_token_estimate(additional.pop("_tokens", 0)) |
| 496 | |
| 497 | self.agent.hist_add_tool_result( |
| 498 | self.name, |
| 499 | final_text_for_agent, |
| 500 | id=self.log.id if self.log else "", |
| 501 | **additional, |
| 502 | ) |
| 503 | if raw_content: |
| 504 | from helpers import history |
| 505 | |
| 506 | self.agent.hist_add_message( |
| 507 | False, |
| 508 | content=history.RawMessage( |
| 509 | raw_content=raw_content, |
| 510 | preview=preview or final_text_for_agent, |
| 511 | ), |
| 512 | tokens=token_estimate, |
| 513 | ) |
| 514 | ( |
| 515 | PrintStyle( |
| 516 | font_color="#1B4F72", background_color="white", padding=True, bold=True |
| 517 | ).print( |
| 518 | f"{self.agent.agent_name}: Response from tool '{self.name}' (plus context added)" |
| 519 | ) |
| 520 | ) |
| 521 | # Print only the raw response to console for brevity, agent gets the full context. |
| 522 | PrintStyle(font_color="#85C1E9").print( |
| 523 | final_text_for_agent |
| 524 | if final_text_for_agent |
| 525 | else "[No direct textual output from tool]" |
| 526 | ) |
| 527 | if self.log: |
| 528 | self.log.update( |
| 529 | content=final_text_for_agent |
| 530 | ) # Log includes the full context |
| 531 | |
| 532 | |
| 533 | class MCPServerRemote(BaseModel): |
| 534 | name: str = Field(default_factory=str) |
| 535 | description: Optional[str] = Field(default="Remote SSE Server") |
| 536 | type: str = Field(default="sse", description="Server connection type") |
| 537 | url: str = Field(default_factory=str) |
| 538 | headers: dict[str, Any] | None = Field(default_factory=dict[str, Any]) |
| 539 | init_timeout: int = Field(default=0) |
| 540 | tool_timeout: int = Field(default=0) |
| 541 | verify: bool = Field(default=True, description="Verify SSL certificates") |
| 542 | disabled: bool = Field(default=False) |
| 543 | disabled_tools: list[str] = Field(default_factory=list) |
| 544 | scope: str = Field(default="global") |
| 545 | |
| 546 | __lock: ClassVar[threading.Lock] = PrivateAttr(default=threading.Lock()) |
| 547 | __client: Optional["MCPClientRemote"] = PrivateAttr(default=None) |
| 548 | |
| 549 | def __init__(self, config: dict[str, Any]): |
| 550 | super().__init__() |
| 551 | self.__client = MCPClientRemote(self) |
| 552 | self.update(config) |
| 553 | |
| 554 | def get_error(self) -> str: |
| 555 | with self.__lock: |
| 556 | return self.__client.error # type: ignore |
| 557 | |
| 558 | def get_log(self) -> str: |
| 559 | with self.__lock: |
| 560 | return self.__client.get_log() # type: ignore |
| 561 | |
| 562 | def get_tools(self) -> List[dict[str, Any]]: |
| 563 | """Get enabled tools from the server""" |
| 564 | with self.__lock: |
| 565 | tools = self.__client.get_tools() # type: ignore |
| 566 | disabled = set(self.disabled_tools) |
| 567 | return [tool for tool in tools if tool.get("name") not in disabled] |
| 568 | |
| 569 | def get_all_tools(self) -> List[dict[str, Any]]: |
| 570 | """Get all tools from the server and mark disabled tools for UI detail views.""" |
| 571 | with self.__lock: |
| 572 | tools = self.__client.get_tools() # type: ignore |
| 573 | disabled = set(self.disabled_tools) |
| 574 | return [ |
| 575 | {**tool, "disabled": tool.get("name") in disabled} |
| 576 | for tool in tools |
| 577 | ] |
| 578 | |
| 579 | def has_tool(self, tool_name: str) -> bool: |
| 580 | """Check if a tool is available""" |
| 581 | if tool_name in self.disabled_tools: |
| 582 | return False |
| 583 | with self.__lock: |
| 584 | return self.__client.has_tool(tool_name) # type: ignore |
| 585 | |
| 586 | async def call_tool( |
| 587 | self, tool_name: str, input_data: Dict[str, Any] |
| 588 | ) -> CallToolResult: |
| 589 | """Call a tool with the given input data""" |
| 590 | client = self.__client |
| 591 | if client is None: |
| 592 | raise RuntimeError("MCP remote client is not initialized") |
| 593 | if tool_name in self.disabled_tools: |
| 594 | raise ValueError(f"Tool {tool_name} is disabled for server {self.name}.") |
| 595 | return await client.call_tool(tool_name, input_data) |
| 596 | |
| 597 | def update(self, config: dict[str, Any]) -> "MCPServerRemote": |
| 598 | with self.__lock: |
| 599 | for key, value in config.items(): |
| 600 | if key in [ |
| 601 | "name", |
| 602 | "description", |
| 603 | "type", |
| 604 | "url", |
| 605 | "serverUrl", |
| 606 | "headers", |
| 607 | "init_timeout", |
| 608 | "tool_timeout", |
| 609 | "disabled", |
| 610 | "disabled_tools", |
| 611 | "verify", |
| 612 | "scope", |
| 613 | ]: |
| 614 | if key == "name": |
| 615 | value = normalize_name(value) |
| 616 | if key == "serverUrl": |
| 617 | key = "url" # remap serverUrl to url |
| 618 | if key == "disabled_tools": |
| 619 | value = _normalize_disabled_tools(value) |
| 620 | |
| 621 | setattr(self, key, value) |
| 622 | return self |
| 623 | |
| 624 | async def initialize(self) -> "MCPServerRemote": |
| 625 | await self.__client.update_tools() # type: ignore |
| 626 | return self |
| 627 | |
| 628 | |
| 629 | class MCPServerLocal(BaseModel): |
| 630 | name: str = Field(default_factory=str) |
| 631 | description: Optional[str] = Field(default="Local StdIO Server") |
| 632 | type: str = Field(default="stdio", description="Server connection type") |
| 633 | command: str = Field(default_factory=str) |
| 634 | args: list[str] = Field(default_factory=list) |
| 635 | env: dict[str, str] | None = Field(default_factory=dict[str, str]) |
| 636 | encoding: str = Field(default="utf-8") |
| 637 | encoding_error_handler: Literal["strict", "ignore", "replace"] = Field( |
| 638 | default="strict" |
| 639 | ) |
| 640 | init_timeout: int = Field(default=0) |
| 641 | tool_timeout: int = Field(default=0) |
| 642 | verify: bool = Field(default=True, description="Verify SSL certificates") |
| 643 | disabled: bool = Field(default=False) |
| 644 | disabled_tools: list[str] = Field(default_factory=list) |
| 645 | scope: str = Field(default="global") |
| 646 | |
| 647 | __lock: ClassVar[threading.Lock] = PrivateAttr(default=threading.Lock()) |
| 648 | __client: Optional["MCPClientLocal"] = PrivateAttr(default=None) |
| 649 | |
| 650 | def __init__(self, config: dict[str, Any]): |
| 651 | super().__init__() |
| 652 | self.__client = MCPClientLocal(self) |
| 653 | self.update(config) |
| 654 | |
| 655 | def get_error(self) -> str: |
| 656 | with self.__lock: |
| 657 | return self.__client.error # type: ignore |
| 658 | |
| 659 | def get_log(self) -> str: |
| 660 | with self.__lock: |
| 661 | return self.__client.get_log() # type: ignore |
| 662 | |
| 663 | def get_tools(self) -> List[dict[str, Any]]: |
| 664 | """Get enabled tools from the server""" |
| 665 | with self.__lock: |
| 666 | tools = self.__client.get_tools() # type: ignore |
| 667 | disabled = set(self.disabled_tools) |
| 668 | return [tool for tool in tools if tool.get("name") not in disabled] |
| 669 | |
| 670 | def get_all_tools(self) -> List[dict[str, Any]]: |
| 671 | """Get all tools from the server and mark disabled tools for UI detail views.""" |
| 672 | with self.__lock: |
| 673 | tools = self.__client.get_tools() # type: ignore |
| 674 | disabled = set(self.disabled_tools) |
| 675 | return [ |
| 676 | {**tool, "disabled": tool.get("name") in disabled} |
| 677 | for tool in tools |
| 678 | ] |
| 679 | |
| 680 | def has_tool(self, tool_name: str) -> bool: |
| 681 | """Check if a tool is available""" |
| 682 | if tool_name in self.disabled_tools: |
| 683 | return False |
| 684 | with self.__lock: |
| 685 | return self.__client.has_tool(tool_name) # type: ignore |
| 686 | |
| 687 | async def call_tool( |
| 688 | self, tool_name: str, input_data: Dict[str, Any] |
| 689 | ) -> CallToolResult: |
| 690 | """Call a tool with the given input data""" |
| 691 | client = self.__client |
| 692 | if client is None: |
| 693 | raise RuntimeError("MCP local client is not initialized") |
| 694 | if tool_name in self.disabled_tools: |
| 695 | raise ValueError(f"Tool {tool_name} is disabled for server {self.name}.") |
| 696 | return await client.call_tool(tool_name, input_data) |
| 697 | |
| 698 | def update(self, config: dict[str, Any]) -> "MCPServerLocal": |
| 699 | with self.__lock: |
| 700 | command = self.command |
| 701 | command_args: list[str] = [] |
| 702 | args = list(self.args) |
| 703 | if "command" in config: |
| 704 | command, command_args = _split_stdio_command(config.get("command")) |
| 705 | args = [*command_args, *args] |
| 706 | if "args" in config: |
| 707 | args = [*command_args, *_normalize_stdio_args(config.get("args"))] |
| 708 | |
| 709 | for key, value in config.items(): |
| 710 | if key in [ |
| 711 | "name", |
| 712 | "description", |
| 713 | "type", |
| 714 | "command", |
| 715 | "args", |
| 716 | "env", |
| 717 | "encoding", |
| 718 | "encoding_error_handler", |
| 719 | "init_timeout", |
| 720 | "tool_timeout", |
| 721 | "disabled", |
| 722 | "disabled_tools", |
| 723 | "scope", |
| 724 | ]: |
| 725 | if key in ["command", "args"]: |
| 726 | continue |
| 727 | if key == "name": |
| 728 | value = normalize_name(value) |
| 729 | if key == "disabled_tools": |
| 730 | value = _normalize_disabled_tools(value) |
| 731 | setattr(self, key, value) |
| 732 | self.command = command |
| 733 | self.args = args |
| 734 | return self |
| 735 | |
| 736 | async def initialize(self) -> "MCPServerLocal": |
| 737 | await self.__client.update_tools() # type: ignore |
| 738 | return self |
| 739 | |
| 740 | |
| 741 | MCPServer = Annotated[ |
| 742 | Union[ |
| 743 | Annotated[MCPServerRemote, Tag("MCPServerRemote")], |
| 744 | Annotated[MCPServerLocal, Tag("MCPServerLocal")], |
| 745 | ], |
| 746 | Discriminator(_determine_server_type), |
| 747 | ] |
| 748 | |
| 749 | |
| 750 | class MCPConfig(BaseModel): |
| 751 | servers: list[MCPServer] = Field(default_factory=list) |
| 752 | disconnected_servers: list[dict[str, Any]] = Field(default_factory=list) |
| 753 | config_scope: str = Field(default="global") |
| 754 | __lock: ClassVar[threading.Lock] = PrivateAttr(default=threading.Lock()) |
| 755 | __instance: ClassVar[Any] = PrivateAttr(default=None) |
| 756 | __initialized: ClassVar[bool] = PrivateAttr(default=False) |
| 757 | __project_instances: ClassVar[dict[str, tuple[str, "MCPConfig"]]] = {} |
| 758 | |
| 759 | @classmethod |
| 760 | def get_instance(cls) -> "MCPConfig": |
| 761 | with cls.__lock: |
| 762 | if cls.__instance is None: |
| 763 | cls.__instance = cls(servers_list=[], config_scope="global") |
| 764 | return cls.__instance |
| 765 | |
| 766 | @classmethod |
| 767 | def clear_project_instances(cls): |
| 768 | with cls.__lock: |
| 769 | cls.__project_instances = {} |
| 770 | |
| 771 | @classmethod |
| 772 | def parse_config_string(cls, config_str: str) -> List[Dict[str, Any]]: |
| 773 | servers_data: List[Dict[str, Any]] = [] |
| 774 | |
| 775 | if not (config_str and config_str.strip()): |
| 776 | return servers_data |
| 777 | |
| 778 | try: |
| 779 | parsed_value = dirty_json.try_parse(config_str) |
| 780 | normalized = cls.normalize_config(parsed_value) |
| 781 | |
| 782 | if isinstance(normalized, list): |
| 783 | for item in normalized: |
| 784 | if isinstance(item, dict): |
| 785 | servers_data.append(dict(item)) |
| 786 | else: |
| 787 | PrintStyle( |
| 788 | background_color="yellow", |
| 789 | font_color="black", |
| 790 | padding=True, |
| 791 | ).print( |
| 792 | f"Warning: MCP config item was not a dictionary and was ignored: {item}" |
| 793 | ) |
| 794 | else: |
| 795 | PrintStyle( |
| 796 | background_color="red", font_color="white", padding=True |
| 797 | ).print( |
| 798 | f"Error: Parsed MCP config top-level structure is not a list. Config string was: '{config_str}'" |
| 799 | ) |
| 800 | except Exception as e_json: |
| 801 | PrintStyle.error( |
| 802 | f"Error parsing MCP config string: {e_json}. Config string was: '{config_str}'" |
| 803 | ) |
| 804 | |
| 805 | return servers_data |
| 806 | |
| 807 | @classmethod |
| 808 | def merge_config_strings( |
| 809 | cls, global_config: str, project_config: str |
| 810 | ) -> tuple[list[dict[str, Any]], str]: |
| 811 | merged: dict[str, dict[str, Any]] = {} |
| 812 | unnamed: list[dict[str, Any]] = [] |
| 813 | |
| 814 | def add_servers(config_str: str, scope: str): |
| 815 | for server in cls.parse_config_string(config_str): |
| 816 | server_copy = dict(server) |
| 817 | server_copy["scope"] = scope |
| 818 | name = str(server_copy.get("name", "") or "").strip() |
| 819 | if not name: |
| 820 | unnamed.append(server_copy) |
| 821 | continue |
| 822 | normalized_name = normalize_name(name) |
| 823 | server_copy["name"] = normalized_name |
| 824 | merged[normalized_name] = server_copy |
| 825 | |
| 826 | add_servers(global_config or DEFAULT_MCP_SERVERS_CONFIG, "global") |
| 827 | add_servers(project_config or DEFAULT_MCP_SERVERS_CONFIG, "project") |
| 828 | |
| 829 | servers = [*unnamed, *merged.values()] |
| 830 | cache_key = dirty_json.stringify( |
| 831 | { |
| 832 | "mcpServers": { |
| 833 | s.get("name", f"unnamed_{i}"): s |
| 834 | for i, s in enumerate(servers) |
| 835 | } |
| 836 | } |
| 837 | ) |
| 838 | return servers, cache_key |
| 839 | |
| 840 | @classmethod |
| 841 | def get_project_instance(cls, project_name: str | None, *, force: bool = False) -> "MCPConfig": |
| 842 | project_key = str(project_name or "").strip() |
| 843 | if not project_key: |
| 844 | return cls.get_instance() |
| 845 | |
| 846 | from helpers import projects |
| 847 | project_key = projects.validate_project_name(project_key) |
| 848 | |
| 849 | global_config = settings.get_settings().get( |
| 850 | "mcp_servers", DEFAULT_MCP_SERVERS_CONFIG |
| 851 | ) |
| 852 | project_config = projects.load_project_mcp_servers(project_key) |
| 853 | servers_data, cache_key = cls.merge_config_strings(global_config, project_config) |
| 854 | |
| 855 | with cls.__lock: |
| 856 | cached = cls.__project_instances.get(project_key) |
| 857 | if cached and cached[0] == cache_key and not force: |
| 858 | return cached[1] |
| 859 | |
| 860 | instance = cls(servers_list=servers_data, config_scope=f"project:{project_key}") |
| 861 | with cls.__lock: |
| 862 | cls.__project_instances[project_key] = (cache_key, instance) |
| 863 | return instance |
| 864 | |
| 865 | @classmethod |
| 866 | def refresh_project(cls, project_name: str) -> "MCPConfig": |
| 867 | project_key = str(project_name or "").strip() |
| 868 | with cls.__lock: |
| 869 | cls.__project_instances.pop(project_key, None) |
| 870 | return cls.get_project_instance(project_key, force=True) |
| 871 | |
| 872 | @classmethod |
| 873 | def get_for_agent(cls, agent: Any) -> "MCPConfig": |
| 874 | try: |
| 875 | from helpers import projects |
| 876 | |
| 877 | project_name = projects.get_context_project_name(agent.context) |
| 878 | if project_name: |
| 879 | return cls.get_project_instance(project_name) |
| 880 | except Exception: |
| 881 | pass |
| 882 | return cls.get_instance() |
| 883 | |
| 884 | @classmethod |
| 885 | def wait_for_lock(cls): |
| 886 | with cls.__lock: |
| 887 | return |
| 888 | |
| 889 | @classmethod |
| 890 | def update(cls, config_str: str) -> Any: |
| 891 | servers_data = cls.parse_config_string(config_str) |
| 892 | new_instance = cls(servers_list=servers_data, config_scope="global") |
| 893 | with cls.__lock: |
| 894 | # Build and initialize outside the class lock so a slow or wedged MCP |
| 895 | # server cannot freeze status reads, prompts, or later tool calls. |
| 896 | instance = cls.__instance |
| 897 | if instance is None: |
| 898 | instance = new_instance |
| 899 | cls.__instance = instance |
| 900 | else: |
| 901 | instance.servers = new_instance.servers |
| 902 | instance.disconnected_servers = new_instance.disconnected_servers |
| 903 | instance.config_scope = new_instance.config_scope |
| 904 | cls.__project_instances = {} |
| 905 | cls.__initialized = True |
| 906 | return instance |
| 907 | |
| 908 | @classmethod |
| 909 | def normalize_config(cls, servers: Any): |
| 910 | normalized = [] |
| 911 | if isinstance(servers, list): |
| 912 | for server in servers: |
| 913 | if isinstance(server, dict): |
| 914 | normalized.append(dict(server)) |
| 915 | elif isinstance(servers, dict): |
| 916 | if "mcpServers" in servers: |
| 917 | if isinstance(servers["mcpServers"], dict): |
| 918 | for key, value in servers["mcpServers"].items(): |
| 919 | if isinstance(value, dict): |
| 920 | server = dict(value) |
| 921 | server["name"] = key |
| 922 | normalized.append(server) |
| 923 | elif isinstance(servers["mcpServers"], list): |
| 924 | for server in servers["mcpServers"]: |
| 925 | if isinstance(server, dict): |
| 926 | normalized.append(dict(server)) |
| 927 | else: |
| 928 | normalized.append(dict(servers)) # single server? |
| 929 | return normalized |
| 930 | |
| 931 | def __init__(self, servers_list: List[Dict[str, Any]], config_scope: str = "global"): |
| 932 | from collections.abc import Mapping, Iterable |
| 933 | |
| 934 | # # DEBUG: Print the received servers_list |
| 935 | # if servers_list: |
| 936 | # PrintStyle(background_color="blue", font_color="white", padding=True).print( |
| 937 | # f"MCPConfig.__init__ received servers_list: {servers_list}" |
| 938 | # ) |
| 939 | |
| 940 | # This empties the servers list if MCPConfig is a Pydantic model and servers is a field. |
| 941 | # If servers is a field like `servers: List[MCPServer] = Field(default_factory=list)`, |
| 942 | # then super().__init__() might try to initialize it. |
| 943 | # We are re-assigning self.servers later in this __init__. |
| 944 | super().__init__() |
| 945 | |
| 946 | # Clear any servers potentially initialized by super().__init__() before we populate based on servers_list |
| 947 | self.servers = [] |
| 948 | self.config_scope = config_scope |
| 949 | # initialize failed servers list |
| 950 | self.disconnected_servers = [] |
| 951 | |
| 952 | if not isinstance(servers_list, Iterable): |
| 953 | ( |
| 954 | PrintStyle( |
| 955 | background_color="grey", font_color="red", padding=True |
| 956 | ).print("MCPConfig::__init__::servers_list must be a list") |
| 957 | ) |
| 958 | return |
| 959 | |
| 960 | for server_item in servers_list: |
| 961 | if not isinstance(server_item, Mapping): |
| 962 | # log the error |
| 963 | error_msg = "server_item must be a mapping" |
| 964 | ( |
| 965 | PrintStyle( |
| 966 | background_color="grey", font_color="red", padding=True |
| 967 | ).print(f"MCPConfig::__init__::{error_msg}") |
| 968 | ) |
| 969 | # add to failed servers with generic name |
| 970 | self.disconnected_servers.append( |
| 971 | { |
| 972 | "config": ( |
| 973 | server_item |
| 974 | if isinstance(server_item, dict) |
| 975 | else {"raw": str(server_item)} |
| 976 | ), |
| 977 | "error": error_msg, |
| 978 | "name": "invalid_server_config", |
| 979 | } |
| 980 | ) |
| 981 | continue |
| 982 | |
| 983 | server_item = dict(server_item) |
| 984 | server_item["disabled_tools"] = _normalize_disabled_tools( |
| 985 | server_item.get("disabled_tools") |
| 986 | ) |
| 987 | |
| 988 | if server_item.get("disabled", False): |
| 989 | # get server name if available |
| 990 | server_name = server_item.get("name", "unnamed_server") |
| 991 | # normalize server name if it exists |
| 992 | if server_name != "unnamed_server": |
| 993 | server_name = normalize_name(server_name) |
| 994 | |
| 995 | # add to failed servers |
| 996 | self.disconnected_servers.append( |
| 997 | { |
| 998 | "config": server_item, |
| 999 | "error": "Disabled in config", |
| 1000 | "name": server_name, |
| 1001 | } |
| 1002 | ) |
| 1003 | continue |
| 1004 | |
| 1005 | server_name = server_item.get("name", "__not__found__") |
| 1006 | if server_name == "__not__found__": |
| 1007 | # log the error |
| 1008 | error_msg = "server_name is required" |
| 1009 | ( |
| 1010 | PrintStyle( |
| 1011 | background_color="grey", font_color="red", padding=True |
| 1012 | ).print(f"MCPConfig::__init__::{error_msg}") |
| 1013 | ) |
| 1014 | # add to failed servers |
| 1015 | self.disconnected_servers.append( |
| 1016 | { |
| 1017 | "config": server_item, |
| 1018 | "error": error_msg, |
| 1019 | "name": "unnamed_server", |
| 1020 | } |
| 1021 | ) |
| 1022 | continue |
| 1023 | |
| 1024 | try: |
| 1025 | # not generic MCPServer because: "Annotated can not be instatioated" |
| 1026 | if server_item.get("url", None) or server_item.get("serverUrl", None): |
| 1027 | self.servers.append(MCPServerRemote(server_item)) |
| 1028 | else: |
| 1029 | self.servers.append(MCPServerLocal(server_item)) |
| 1030 | except Exception as e: |
| 1031 | # log the error |
| 1032 | error_msg = str(e) |
| 1033 | ( |
| 1034 | PrintStyle( |
| 1035 | background_color="grey", font_color="red", padding=True |
| 1036 | ).print( |
| 1037 | f"MCPConfig::__init__: Failed to create MCPServer '{server_name}': {error_msg}" |
| 1038 | ) |
| 1039 | ) |
| 1040 | # add to failed servers |
| 1041 | self.disconnected_servers.append( |
| 1042 | {"config": server_item, "error": error_msg, "name": server_name} |
| 1043 | ) |
| 1044 | |
| 1045 | # Initialize all servers in parallel (fetch tools concurrently) |
| 1046 | if self.servers: |
| 1047 | async def _init_server(server): |
| 1048 | try: |
| 1049 | await server.initialize() |
| 1050 | except Exception as e: |
| 1051 | error_msg = str(e) |
| 1052 | PrintStyle( |
| 1053 | background_color="grey", font_color="red", padding=True |
| 1054 | ).print( |
| 1055 | f"MCPConfig::__init__: Failed to initialize MCPServer '{server.name}': {error_msg}" |
| 1056 | ) |
| 1057 | |
| 1058 | async def _init_all(): |
| 1059 | await asyncio.gather(*[_init_server(s) for s in self.servers]) |
| 1060 | |
| 1061 | asyncio.run(_init_all()) |
| 1062 | |
| 1063 | def get_server_log(self, server_name: str) -> str: |
| 1064 | with self.__lock: |
| 1065 | for server in self.servers: |
| 1066 | if server.name == server_name: |
| 1067 | return server.get_log() # type: ignore |
| 1068 | return "" |
| 1069 | |
| 1070 | def get_servers_status(self) -> list[dict[str, Any]]: |
| 1071 | """Get status of all servers""" |
| 1072 | result = [] |
| 1073 | with self.__lock: |
| 1074 | # add connected/working servers |
| 1075 | for server in self.servers: |
| 1076 | # get server name |
| 1077 | name = server.name |
| 1078 | # get tool count |
| 1079 | tool_count = len(server.get_tools()) |
| 1080 | # get error message if any |
| 1081 | error = server.get_error() |
| 1082 | # A server object can exist while its initialization failed. |
| 1083 | connected = not bool(error) |
| 1084 | # get log bool |
| 1085 | has_log = server.get_log() != "" |
| 1086 | |
| 1087 | # add server status to result |
| 1088 | result.append( |
| 1089 | { |
| 1090 | "name": name, |
| 1091 | "scope": getattr(server, "scope", self.config_scope), |
| 1092 | "type": getattr(server, "type", ""), |
| 1093 | "description": getattr(server, "description", ""), |
| 1094 | "connected": connected, |
| 1095 | "error": error, |
| 1096 | "tool_count": tool_count, |
| 1097 | "has_log": has_log, |
| 1098 | } |
| 1099 | ) |
| 1100 | |
| 1101 | # add failed servers |
| 1102 | for disconnected in self.disconnected_servers: |
| 1103 | result.append( |
| 1104 | { |
| 1105 | "name": disconnected["name"], |
| 1106 | "scope": disconnected.get("config", {}).get("scope", self.config_scope), |
| 1107 | "type": disconnected.get("config", {}).get("type", ""), |
| 1108 | "description": disconnected.get("config", {}).get("description", ""), |
| 1109 | "connected": False, |
| 1110 | "error": disconnected["error"], |
| 1111 | "tool_count": 0, |
| 1112 | "has_log": False, |
| 1113 | } |
| 1114 | ) |
| 1115 | |
| 1116 | return result |
| 1117 | |
| 1118 | def get_server_detail(self, server_name: str) -> dict[str, Any]: |
| 1119 | with self.__lock: |
| 1120 | for server in self.servers: |
| 1121 | if server.name == server_name: |
| 1122 | try: |
| 1123 | get_all_tools = getattr(server, "get_all_tools", None) |
| 1124 | tools = get_all_tools() if callable(get_all_tools) else server.get_tools() |
| 1125 | except Exception: |
| 1126 | tools = [] |
| 1127 | return { |
| 1128 | "name": server.name, |
| 1129 | "description": server.description, |
| 1130 | "scope": getattr(server, "scope", self.config_scope), |
| 1131 | "type": getattr(server, "type", ""), |
| 1132 | "tools": tools, |
| 1133 | } |
| 1134 | return {} |
| 1135 | |
| 1136 | def is_initialized(self) -> bool: |
| 1137 | """Check if the client is initialized""" |
| 1138 | with self.__lock: |
| 1139 | return self.__initialized |
| 1140 | |
| 1141 | def get_tools(self) -> List[dict[str, dict[str, Any]]]: |
| 1142 | """Get all tools from all servers""" |
| 1143 | with self.__lock: |
| 1144 | tools = [] |
| 1145 | for server in self.servers: |
| 1146 | for tool in server.get_tools(): |
| 1147 | tool_copy = tool.copy() |
| 1148 | tool_copy["server"] = server.name |
| 1149 | tools.append({f"{server.name}.{tool['name']}": tool_copy}) |
| 1150 | return tools |
| 1151 | |
| 1152 | def get_tools_prompt(self, server_name: str = "", agent: Any | None = None) -> str: |
| 1153 | """Get a prompt for all tools""" |
| 1154 | |
| 1155 | # just to wait for pending initialization |
| 1156 | with self.__lock: |
| 1157 | pass |
| 1158 | |
| 1159 | prompt = '## "Remote (MCP Server) Agent Tools" available:\n\n' |
| 1160 | server_names = [] |
| 1161 | for server in self.servers: |
| 1162 | if not server_name or server.name == server_name: |
| 1163 | server_names.append(server.name) |
| 1164 | |
| 1165 | if server_name and server_name not in server_names: |
| 1166 | raise ValueError(f"Server {server_name} not found") |
| 1167 | |
| 1168 | for server in self.servers: |
| 1169 | if server.name in server_names: |
| 1170 | server_name = server.name |
| 1171 | prompt += f"### {server_name}\n" |
| 1172 | prompt += f"{server.description}\n" |
| 1173 | tools = server.get_tools() |
| 1174 | |
| 1175 | for tool in tools: |
| 1176 | qualified_name = f"{server_name}.{tool['name']}" |
| 1177 | if agent is not None: |
| 1178 | from helpers.tool_policy import canonical_mcp_id, resolve_tool |
| 1179 | |
| 1180 | if not resolve_tool( |
| 1181 | agent, |
| 1182 | qualified_name, |
| 1183 | canonical_id=canonical_mcp_id(qualified_name), |
| 1184 | ).allowed: |
| 1185 | continue |
| 1186 | prompt += ( |
| 1187 | f"\n### {qualified_name}:\n" |
| 1188 | f"{tool['description']}\n\n" |
| 1189 | # f"#### Categories:\n" |
| 1190 | # f"* kind: MCP Server Tool\n" |
| 1191 | # f'* server: "{server_name}" ({server.description})\n\n' |
| 1192 | # f"#### Arguments:\n" |
| 1193 | ) |
| 1194 | |
| 1195 | input_schema = ( |
| 1196 | json.dumps(tool["input_schema"]) if tool["input_schema"] else "" |
| 1197 | ) |
| 1198 | |
| 1199 | prompt += f"#### Input schema for tool_args:\n{input_schema}\n" |
| 1200 | |
| 1201 | prompt += "\n" |
| 1202 | |
| 1203 | prompt += ( |
| 1204 | f"#### Usage:\n" |
| 1205 | f"{{\n" |
| 1206 | # f' "observations": ["..."],\n' # TODO: this should be a prompt file with placeholders |
| 1207 | f' "thoughts": ["..."],\n' |
| 1208 | # f' "reflection": ["..."],\n' # TODO: this should be a prompt file with placeholders |
| 1209 | f" \"tool_name\": \"{qualified_name}\",\n" |
| 1210 | f' "tool_args": !follow schema above\n' |
| 1211 | f"}}\n" |
| 1212 | ) |
| 1213 | |
| 1214 | return prompt |
| 1215 | |
| 1216 | def has_tool(self, tool_name: str) -> bool: |
| 1217 | """Check if a tool is available""" |
| 1218 | try: |
| 1219 | server_name_part, tool_name_part = _split_qualified_tool_name(tool_name) |
| 1220 | except ValueError: |
| 1221 | return False |
| 1222 | with self.__lock: |
| 1223 | for server in self.servers: |
| 1224 | if server.name == server_name_part: |
| 1225 | return server.has_tool(tool_name_part) |
| 1226 | return False |
| 1227 | |
| 1228 | def get_tool(self, agent: Any, tool_name: str) -> MCPTool | None: |
| 1229 | effective_config = MCPConfig.get_for_agent(agent) |
| 1230 | if effective_config is not self: |
| 1231 | return effective_config.get_tool(agent, tool_name) |
| 1232 | if not self.has_tool(tool_name): |
| 1233 | get_data = getattr(agent, "get_data", None) |
| 1234 | name_map_key = getattr(agent, "DATA_NAME_RESPONSES_TOOL_NAME_MAP", "") |
| 1235 | tool_name = original_tool_name( |
| 1236 | tool_name, |
| 1237 | get_data(name_map_key) |
| 1238 | if name_map_key and callable(get_data) |
| 1239 | else None, |
| 1240 | ) |
| 1241 | if not self.has_tool(tool_name): |
| 1242 | return None |
| 1243 | return MCPTool(agent=agent, name=tool_name, method=None, args={}, message="", loop_data=None) |
| 1244 | |
| 1245 | async def call_tool( |
| 1246 | self, tool_name: str, input_data: Dict[str, Any] |
| 1247 | ) -> CallToolResult: |
| 1248 | """Call a tool with the given input data""" |
| 1249 | server_name_part, tool_name_part = _split_qualified_tool_name(tool_name) |
| 1250 | matched_server = None |
| 1251 | with self.__lock: |
| 1252 | for server in self.servers: |
| 1253 | if server.name == server_name_part and server.has_tool(tool_name_part): |
| 1254 | matched_server = server |
| 1255 | break |
| 1256 | if matched_server is None: |
| 1257 | raise ValueError(f"Tool {tool_name} not found") |
| 1258 | return await matched_server.call_tool(tool_name_part, input_data) |
| 1259 | |
| 1260 | |
| 1261 | T = TypeVar("T") |
| 1262 | |
| 1263 | |
| 1264 | class MCPClientBase(ABC): |
| 1265 | # server: Union[MCPServerLocal, MCPServerRemote] # Defined in __init__ |
| 1266 | # tools: List[dict[str, Any]] # Defined in __init__ |
| 1267 | # No self.session, self.exit_stack, self.stdio, self.write as persistent instance fields |
| 1268 | |
| 1269 | __lock: ClassVar[threading.Lock] = threading.Lock() |
| 1270 | |
| 1271 | def __init__(self, server: Union[MCPServerLocal, MCPServerRemote]): |
| 1272 | self.server = server |
| 1273 | self.tools: List[dict[str, Any]] = [] # Tools are cached on the client instance |
| 1274 | self.error: str = "" |
| 1275 | self.log: List[str] = [] |
| 1276 | self.log_file: Optional[TextIO] = None |
| 1277 | |
| 1278 | def _operation_timeout_seconds(self, read_timeout_seconds: float) -> float: |
| 1279 | try: |
| 1280 | seconds = float(read_timeout_seconds) |
| 1281 | except (TypeError, ValueError): |
| 1282 | seconds = 60.0 |
| 1283 | if seconds <= 0: |
| 1284 | seconds = 60.0 |
| 1285 | return seconds + MCP_OPERATION_TIMEOUT_GRACE_SECONDS |
| 1286 | |
| 1287 | def _operation_thread_name(self, operation_name: str) -> str: |
| 1288 | server_name = normalize_name(str(getattr(self.server, "name", "") or "server")) |
| 1289 | return f"MCPClient-{server_name[:32] or 'server'}-{operation_name}-{uuid.uuid4().hex[:8]}" |
| 1290 | |
| 1291 | async def _run_isolated_operation( |
| 1292 | self, |
| 1293 | operation_name: str, |
| 1294 | operation: Callable[[], Awaitable[T]], |
| 1295 | timeout_seconds: float, |
| 1296 | ) -> T: |
| 1297 | worker = DeferredTask(thread_name=self._operation_thread_name(operation_name)) |
| 1298 | timed_out = False |
| 1299 | try: |
| 1300 | return await asyncio.wait_for( |
| 1301 | worker.execute_inside(operation), |
| 1302 | timeout=timeout_seconds, |
| 1303 | ) |
| 1304 | except asyncio.TimeoutError as exc: |
| 1305 | timed_out = True |
| 1306 | message = ( |
| 1307 | f"MCPClientBase ({self.server.name} - {operation_name}): " |
| 1308 | f"operation did not finish within {timeout_seconds:.1f}s; " |
| 1309 | "abandoning the isolated worker so Agent Zero can continue." |
| 1310 | ) |
| 1311 | PrintStyle.warning(message) |
| 1312 | with self.__lock: |
| 1313 | self.error = message |
| 1314 | raise TimeoutError(message) from exc |
| 1315 | finally: |
| 1316 | if timed_out: |
| 1317 | worker.kill(terminate_thread=False) |
| 1318 | else: |
| 1319 | worker.kill(terminate_thread=True) |
| 1320 | |
| 1321 | # Protected method |
| 1322 | @abstractmethod |
| 1323 | async def _create_stdio_transport( |
| 1324 | self, current_exit_stack: AsyncExitStack |
| 1325 | ) -> tuple[ |
| 1326 | MemoryObjectReceiveStream[SessionMessage | Exception], |
| 1327 | MemoryObjectSendStream[SessionMessage], |
| 1328 | ]: |
| 1329 | """Create stdio/write streams using the provided exit_stack.""" |
| 1330 | ... |
| 1331 | |
| 1332 | async def _execute_with_session( |
| 1333 | self, |
| 1334 | coro_func: Callable[[ClientSession], Awaitable[T]], |
| 1335 | read_timeout_seconds=60, |
| 1336 | ) -> T: |
| 1337 | """ |
| 1338 | Manages the lifecycle of an MCP session for a single operation. |
| 1339 | Creates a temporary session, executes coro_func with it, and ensures cleanup. |
| 1340 | """ |
| 1341 | operation_name = coro_func.__name__ # For logging |
| 1342 | # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Creating new session for operation '{operation_name}'...") |
| 1343 | original_exception = None |
| 1344 | result: T | None = None |
| 1345 | has_result = False |
| 1346 | temp_stack = AsyncExitStack() |
| 1347 | try: |
| 1348 | stdio, write = await self._create_stdio_transport(temp_stack) |
| 1349 | # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...") |
| 1350 | session = await temp_stack.enter_async_context( |
| 1351 | ClientSession( |
| 1352 | stdio, # type: ignore |
| 1353 | write, # type: ignore |
| 1354 | read_timeout_seconds=timedelta( |
| 1355 | seconds=read_timeout_seconds |
| 1356 | ), |
| 1357 | ) |
| 1358 | ) |
| 1359 | await session.initialize() |
| 1360 | |
| 1361 | result = await coro_func(session) |
| 1362 | has_result = True |
| 1363 | except Exception as e: |
| 1364 | excs = getattr(e, "exceptions", None) # Python 3.11+ ExceptionGroup |
| 1365 | if excs: |
| 1366 | original_exception = excs[0] |
| 1367 | else: |
| 1368 | original_exception = e |
| 1369 | try: |
| 1370 | await asyncio.wait_for( |
| 1371 | temp_stack.aclose(), |
| 1372 | timeout=MCP_SESSION_CLEANUP_TIMEOUT_SECONDS, |
| 1373 | ) |
| 1374 | except asyncio.TimeoutError: |
| 1375 | PrintStyle.warning( |
| 1376 | f"MCPClientBase ({self.server.name} - {operation_name}): " |
| 1377 | f"session cleanup exceeded {MCP_SESSION_CLEANUP_TIMEOUT_SECONDS:.1f}s." |
| 1378 | ) |
| 1379 | except Exception as cleanup_exception: |
| 1380 | PrintStyle.warning( |
| 1381 | f"MCPClientBase ({self.server.name} - {operation_name}): " |
| 1382 | f"session cleanup failed: {type(cleanup_exception).__name__}: {cleanup_exception}" |
| 1383 | ) |
| 1384 | if original_exception is not None: |
| 1385 | PrintStyle( |
| 1386 | background_color="#AA4455", font_color="white", padding=False |
| 1387 | ).print( |
| 1388 | f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(original_exception).__name__}: {original_exception}" |
| 1389 | ) |
| 1390 | raise original_exception |
| 1391 | if has_result: |
| 1392 | return cast(T, result) |
| 1393 | raise RuntimeError( |
| 1394 | f"MCPClientBase ({self.server.name} - {operation_name}): _execute_with_session exited 'async with' block unexpectedly." |
| 1395 | ) |
| 1396 | |
| 1397 | async def update_tools(self) -> "MCPClientBase": |
| 1398 | # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Starting 'update_tools' operation...") |
| 1399 | |
| 1400 | async def list_tools_op(current_session: ClientSession): |
| 1401 | response: ListToolsResult = await current_session.list_tools() |
| 1402 | with self.__lock: |
| 1403 | self.tools = [ |
| 1404 | { |
| 1405 | "name": tool.name, |
| 1406 | "description": tool.description, |
| 1407 | "input_schema": tool.inputSchema, |
| 1408 | } |
| 1409 | for tool in response.tools |
| 1410 | ] |
| 1411 | self.error = "" |
| 1412 | PrintStyle(font_color="green").print( |
| 1413 | f"MCPClientBase ({self.server.name}): Tools updated. Found {len(self.tools)} tools." |
| 1414 | ) |
| 1415 | |
| 1416 | try: |
| 1417 | current_settings = settings.get_settings() |
| 1418 | init_timeout = ( |
| 1419 | self.server.init_timeout |
| 1420 | or current_settings.get("mcp_client_init_timeout", 10) |
| 1421 | or 10 |
| 1422 | ) |
| 1423 | await self._run_isolated_operation( |
| 1424 | "update_tools", |
| 1425 | lambda: self._execute_with_session( |
| 1426 | list_tools_op, |
| 1427 | read_timeout_seconds=init_timeout, |
| 1428 | ), |
| 1429 | timeout_seconds=self._operation_timeout_seconds(init_timeout), |
| 1430 | ) |
| 1431 | except Exception as e: |
| 1432 | # e = eg.exceptions[0] |
| 1433 | error_text = errors.format_error(e, 0, 0) |
| 1434 | # Error already logged by _execute_with_session, this is for specific handling if needed |
| 1435 | PrintStyle( |
| 1436 | background_color="#CC34C3", font_color="white", bold=True, padding=True |
| 1437 | ).print( |
| 1438 | f"MCPClientBase ({self.server.name}): 'update_tools' operation failed: {error_text}" |
| 1439 | ) |
| 1440 | with self.__lock: |
| 1441 | self.tools = [] # Ensure tools are cleared on failure |
| 1442 | self.error = f"Failed to initialize. {error_text[:200]}{'...' if len(error_text) > 200 else ''}" # store error from tools fetch |
| 1443 | return self |
| 1444 | |
| 1445 | def has_tool(self, tool_name: str) -> bool: |
| 1446 | """Check if a tool is available (uses cached tools)""" |
| 1447 | with self.__lock: |
| 1448 | for tool in self.tools: |
| 1449 | if tool["name"] == tool_name: |
| 1450 | return True |
| 1451 | return False |
| 1452 | |
| 1453 | def get_tools(self) -> List[dict[str, Any]]: |
| 1454 | """Get all tools from the server (uses cached tools)""" |
| 1455 | with self.__lock: |
| 1456 | return [dict(tool) for tool in self.tools] |
| 1457 | |
| 1458 | async def call_tool( |
| 1459 | self, tool_name: str, input_data: Dict[str, Any] |
| 1460 | ) -> CallToolResult: |
| 1461 | # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Preparing for 'call_tool' operation for tool '{tool_name}'.") |
| 1462 | if not self.has_tool(tool_name): |
| 1463 | PrintStyle(font_color="orange").print( |
| 1464 | f"MCPClientBase ({self.server.name}): Tool '{tool_name}' not in cache for 'call_tool', refreshing tools..." |
| 1465 | ) |
| 1466 | await self.update_tools() # This will use its own properly managed session |
| 1467 | if not self.has_tool(tool_name): |
| 1468 | PrintStyle(font_color="red").print( |
| 1469 | f"MCPClientBase ({self.server.name}): Tool '{tool_name}' not found after refresh. Raising ValueError." |
| 1470 | ) |
| 1471 | raise ValueError( |
| 1472 | f"Tool {tool_name} not found after refreshing tool list for server {self.server.name}." |
| 1473 | ) |
| 1474 | PrintStyle(font_color="green").print( |
| 1475 | f"MCPClientBase ({self.server.name}): Tool '{tool_name}' found after updating tools." |
| 1476 | ) |
| 1477 | |
| 1478 | current_settings = settings.get_settings() |
| 1479 | tool_timeout = ( |
| 1480 | self.server.tool_timeout |
| 1481 | or current_settings.get("mcp_client_tool_timeout", 120) |
| 1482 | or 120 |
| 1483 | ) |
| 1484 | |
| 1485 | async def call_tool_op(current_session: ClientSession): |
| 1486 | # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Executing 'call_tool' for '{tool_name}' via MCP session...") |
| 1487 | response: CallToolResult = await current_session.call_tool( |
| 1488 | tool_name, |
| 1489 | input_data, |
| 1490 | read_timeout_seconds=timedelta(seconds=tool_timeout), |
| 1491 | ) |
| 1492 | # PrintStyle(font_color="green").print(f"MCPClientBase ({self.server.name}): Tool '{tool_name}' call successful via session.") |
| 1493 | return response |
| 1494 | |
| 1495 | try: |
| 1496 | response = await self._run_isolated_operation( |
| 1497 | "call_tool", |
| 1498 | lambda: self._execute_with_session( |
| 1499 | call_tool_op, |
| 1500 | read_timeout_seconds=tool_timeout, |
| 1501 | ), |
| 1502 | timeout_seconds=self._operation_timeout_seconds(tool_timeout), |
| 1503 | ) |
| 1504 | with self.__lock: |
| 1505 | self.error = "" |
| 1506 | return response |
| 1507 | except Exception as e: |
| 1508 | # Error logged by _execute_with_session. Re-raise a specific error for the caller. |
| 1509 | PrintStyle( |
| 1510 | background_color="#AA4455", font_color="white", padding=True |
| 1511 | ).print( |
| 1512 | f"MCPClientBase ({self.server.name}): 'call_tool' operation for '{tool_name}' failed: {type(e).__name__}: {e}" |
| 1513 | ) |
| 1514 | raise ConnectionError( |
| 1515 | f"MCPClientBase::Failed to call tool '{tool_name}' on server '{self.server.name}'. Original error: {type(e).__name__}: {e}" |
| 1516 | ) |
| 1517 | |
| 1518 | def get_log(self): |
| 1519 | # read and return lines from self.log_file, do not close it |
| 1520 | if not hasattr(self, "log_file") or self.log_file is None: |
| 1521 | return "" |
| 1522 | self.log_file.seek(0) |
| 1523 | try: |
| 1524 | log = self.log_file.read() |
| 1525 | except Exception: |
| 1526 | log = "" |
| 1527 | return log |
| 1528 | |
| 1529 | |
| 1530 | class MCPClientLocal(MCPClientBase): |
| 1531 | def __del__(self): |
| 1532 | # close the log file if it exists |
| 1533 | if hasattr(self, "log_file") and self.log_file is not None: |
| 1534 | try: |
| 1535 | self.log_file.close() |
| 1536 | except Exception: |
| 1537 | pass |
| 1538 | self.log_file = None |
| 1539 | |
| 1540 | async def _create_stdio_transport( |
| 1541 | self, current_exit_stack: AsyncExitStack |
| 1542 | ) -> tuple[ |
| 1543 | MemoryObjectReceiveStream[SessionMessage | Exception], |
| 1544 | MemoryObjectSendStream[SessionMessage], |
| 1545 | ]: |
| 1546 | """Connect to an MCP server, init client and save stdio/write streams""" |
| 1547 | server: MCPServerLocal = cast(MCPServerLocal, self.server) |
| 1548 | |
| 1549 | if not server.command: |
| 1550 | raise ValueError("Command not specified") |
| 1551 | if not which(server.command): |
| 1552 | raise ValueError(f"Command '{server.command}' not found") |
| 1553 | |
| 1554 | server_params = StdioServerParameters( |
| 1555 | command=server.command, |
| 1556 | args=server.args, |
| 1557 | env=server.env, |
| 1558 | encoding=server.encoding, |
| 1559 | encoding_error_handler=server.encoding_error_handler, |
| 1560 | ) |
| 1561 | # create a custom error log handler that will capture error output |
| 1562 | import tempfile |
| 1563 | |
| 1564 | # use a temporary file for error logging (text mode) if not already present |
| 1565 | if not hasattr(self, "log_file") or self.log_file is None: |
| 1566 | self.log_file = tempfile.TemporaryFile(mode="w+", encoding="utf-8") |
| 1567 | |
| 1568 | # use the stdio_client with our error log file |
| 1569 | stdio_transport = await current_exit_stack.enter_async_context( |
| 1570 | stdio_client(server_params, errlog=self.log_file) |
| 1571 | ) |
| 1572 | # do not read or close the file here, as stdio is async |
| 1573 | return stdio_transport |
| 1574 | |
| 1575 | class CustomHTTPClientFactory(ABC): |
| 1576 | def __init__(self, verify: bool = True): |
| 1577 | self.verify = verify |
| 1578 | |
| 1579 | def __call__( |
| 1580 | self, |
| 1581 | headers: dict[str, str] | None = None, |
| 1582 | timeout: httpx.Timeout | None = None, |
| 1583 | auth: httpx.Auth | None = None, |
| 1584 | ) -> httpx.AsyncClient: |
| 1585 | # Set MCP defaults |
| 1586 | kwargs: dict[str, Any] = { |
| 1587 | "follow_redirects": True, |
| 1588 | } |
| 1589 | |
| 1590 | # Handle timeout |
| 1591 | if timeout is None: |
| 1592 | kwargs["timeout"] = httpx.Timeout(30.0) |
| 1593 | else: |
| 1594 | kwargs["timeout"] = timeout |
| 1595 | |
| 1596 | # Handle headers |
| 1597 | if headers is not None: |
| 1598 | kwargs["headers"] = headers |
| 1599 | |
| 1600 | # Handle authentication |
| 1601 | if auth is not None: |
| 1602 | kwargs["auth"] = auth |
| 1603 | |
| 1604 | return httpx.AsyncClient(**kwargs, verify=self.verify) |
| 1605 | |
| 1606 | class MCPClientRemote(MCPClientBase): |
| 1607 | |
| 1608 | def __init__(self, server: Union[MCPServerLocal, MCPServerRemote]): |
| 1609 | super().__init__(server) |
| 1610 | self.session_id: Optional[str] = None # Track session ID for streaming HTTP clients |
| 1611 | self.session_id_callback: Optional[Callable[[], Optional[str]]] = None |
| 1612 | |
| 1613 | async def _create_stdio_transport( |
| 1614 | self, current_exit_stack: AsyncExitStack |
| 1615 | ) -> tuple[ |
| 1616 | MemoryObjectReceiveStream[SessionMessage | Exception], |
| 1617 | MemoryObjectSendStream[SessionMessage], |
| 1618 | ]: |
| 1619 | """Connect to an MCP server, init client and save stdio/write streams""" |
| 1620 | server: MCPServerRemote = cast(MCPServerRemote, self.server) |
| 1621 | current_settings = settings.get_settings() |
| 1622 | |
| 1623 | # Resolve timeout: check server config first, then settings, defaulting to 5s/10s |
| 1624 | init_timeout = ( |
| 1625 | server.init_timeout |
| 1626 | or current_settings.get("mcp_client_init_timeout", 10) |
| 1627 | or 10 |
| 1628 | ) |
| 1629 | tool_timeout = ( |
| 1630 | server.tool_timeout |
| 1631 | or current_settings.get("mcp_client_tool_timeout", 120) |
| 1632 | or 120 |
| 1633 | ) |
| 1634 | |
| 1635 | client_factory = CustomHTTPClientFactory(verify=server.verify) |
| 1636 | # Check if this is a streaming HTTP type |
| 1637 | if _is_streaming_http_type(server.type): |
| 1638 | # Use streamable HTTP client |
| 1639 | transport_result = await current_exit_stack.enter_async_context( |
| 1640 | streamablehttp_client( |
| 1641 | url=server.url, |
| 1642 | headers=server.headers, |
| 1643 | timeout=timedelta(seconds=init_timeout), |
| 1644 | sse_read_timeout=timedelta(seconds=tool_timeout), |
| 1645 | httpx_client_factory=client_factory, |
| 1646 | ) |
| 1647 | ) |
| 1648 | # streamablehttp_client returns (read_stream, write_stream, get_session_id_callback) |
| 1649 | read_stream, write_stream, get_session_id_callback = transport_result |
| 1650 | |
| 1651 | # Store session ID callback for potential future use |
| 1652 | self.session_id_callback = get_session_id_callback |
| 1653 | |
| 1654 | return read_stream, write_stream |
| 1655 | else: |
| 1656 | # Use traditional SSE client (default behavior) |
| 1657 | stdio_transport = await current_exit_stack.enter_async_context( |
| 1658 | sse_client( |
| 1659 | url=server.url, |
| 1660 | headers=server.headers, |
| 1661 | timeout=init_timeout, |
| 1662 | sse_read_timeout=tool_timeout, |
| 1663 | httpx_client_factory=client_factory, |
| 1664 | ) |
| 1665 | ) |
| 1666 | return stdio_transport |
| 1667 | |
| 1668 | def get_session_id(self) -> Optional[str]: |
| 1669 | """Get the current session ID if available (for streaming HTTP clients).""" |
| 1670 | if self.session_id_callback is not None: |
| 1671 | return self.session_id_callback() |
| 1672 | return None |