-
-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Expand file tree
/
Copy pathfunc_tool_manager.py
More file actions
1029 lines (889 loc) · 36.8 KB
/
func_tool_manager.py
File metadata and controls
1029 lines (889 loc) · 36.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from __future__ import annotations
import asyncio
import copy
import json
import os
import threading
import urllib.parse
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any
import aiohttp
from astrbot import logger
from astrbot.core import sp
from astrbot.core.agent.mcp_client import MCPClient, MCPTool
from astrbot.core.agent.tool import FunctionTool, ToolSet
from astrbot.core.tools.registry import (
ensure_builtin_tools_loaded,
get_builtin_tool_class,
get_builtin_tool_name,
iter_builtin_tool_classes,
)
from astrbot.core.utils.astrbot_path import get_astrbot_data_path
DEFAULT_MCP_CONFIG = {"mcpServers": {}}
DEFAULT_MCP_INIT_TIMEOUT_SECONDS = 180.0
DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS = 180.0
MCP_INIT_TIMEOUT_ENV = "ASTRBOT_MCP_INIT_TIMEOUT"
ENABLE_MCP_TIMEOUT_ENV = "ASTRBOT_MCP_ENABLE_TIMEOUT"
MAX_MCP_TIMEOUT_SECONDS = 300.0
class MCPInitError(Exception):
"""Base exception for MCP initialization failures."""
class MCPInitTimeoutError(asyncio.TimeoutError, MCPInitError):
"""Raised when MCP client initialization exceeds the configured timeout."""
class MCPAllServicesFailedError(MCPInitError):
"""Raised when all configured MCP services fail to initialize."""
class MCPShutdownTimeoutError(asyncio.TimeoutError):
"""Raised when MCP shutdown exceeds the configured timeout."""
def __init__(self, names: list[str], timeout: float) -> None:
self.names = names
self.timeout = timeout
message = f"MCP 服务关闭超时({timeout:g} 秒):{', '.join(names)}"
super().__init__(message)
@dataclass
class MCPInitSummary:
total: int
success: int
failed: list[str]
@dataclass
class _MCPServerRuntime:
name: str
client: MCPClient
shutdown_event: asyncio.Event
lifecycle_task: asyncio.Task[None]
class _MCPClientDictView(Mapping[str, MCPClient]):
"""Read-only view of MCP clients derived from runtime state."""
def __init__(self, runtime: dict[str, _MCPServerRuntime]) -> None:
self._runtime = runtime
def __getitem__(self, key: str) -> MCPClient:
return self._runtime[key].client
def __iter__(self):
return iter(self._runtime)
def __len__(self) -> int:
return len(self._runtime)
def _resolve_timeout(
timeout: float | int | str | None = None,
*,
env_name: str = MCP_INIT_TIMEOUT_ENV,
default: float = DEFAULT_MCP_INIT_TIMEOUT_SECONDS,
) -> float:
"""Resolve timeout with precedence: explicit argument > env value > default."""
source = f"环境变量 {env_name}"
if timeout is None:
timeout = os.getenv(env_name, str(default))
else:
source = "显式参数 timeout"
try:
timeout_value = float(timeout)
except (TypeError, ValueError):
logger.warning(
f"超时配置({source})={timeout!r} 无效,使用默认值 {default:g} 秒。"
)
return default
if timeout_value <= 0:
logger.warning(
f"超时配置({source})={timeout_value:g} 必须大于 0,使用默认值 {default:g} 秒。"
)
return default
if timeout_value > MAX_MCP_TIMEOUT_SECONDS:
logger.warning(
f"超时配置({source})={timeout_value:g} 过大,已限制为最大值 "
f"{MAX_MCP_TIMEOUT_SECONDS:g} 秒,以避免长时间等待。"
)
return MAX_MCP_TIMEOUT_SECONDS
return timeout_value
SUPPORTED_TYPES = [
"string",
"number",
"object",
"array",
"boolean",
] # json schema 支持的数据类型
PY_TO_JSON_TYPE = {
"int": "number",
"float": "number",
"bool": "boolean",
"str": "string",
"dict": "object",
"list": "array",
"tuple": "array",
"set": "array",
}
# alias
FuncTool = FunctionTool
def _prepare_config(config: dict) -> dict:
"""准备配置,处理嵌套格式"""
if config.get("mcpServers"):
first_key = next(iter(config["mcpServers"]))
config = config["mcpServers"][first_key]
config.pop("active", None)
return config
async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]:
"""快速测试 MCP 服务器可达性"""
import aiohttp
cfg = _prepare_config(config.copy())
url = cfg["url"]
headers = cfg.get("headers", {})
timeout = cfg.get("timeout", 10)
try:
async with aiohttp.ClientSession() as session:
if cfg.get("transport") == "streamable_http":
test_payload = {
"jsonrpc": "2.0",
"method": "initialize",
"id": 0,
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test-client", "version": "1.2.3"},
},
}
async with session.post(
url,
headers={
**headers,
"Content-Type": "application/json",
"Accept": "application/json, text/event-stream",
},
json=test_payload,
timeout=aiohttp.ClientTimeout(total=timeout),
) as response:
if response.status == 200:
return True, ""
return False, f"HTTP {response.status}: {response.reason}"
else:
async with session.get(
url,
headers={
**headers,
"Accept": "application/json, text/event-stream",
},
timeout=aiohttp.ClientTimeout(total=timeout),
) as response:
if response.status == 200:
return True, ""
return False, f"HTTP {response.status}: {response.reason}"
except asyncio.TimeoutError:
return False, f"连接超时: {timeout}秒"
except Exception as e:
return False, f"{e!s}"
class FunctionToolManager:
def __init__(self) -> None:
self.func_list: list[FuncTool] = []
"""All tools include mcp tools and plugin tools, except astrbot builtin tools."""
self.builtin_func_list: dict[type[FuncTool], FuncTool] = {}
"""All astrbot builtin tools, keyed by their class. Values are instantiated tool objects, created on demand."""
self._mcp_server_runtime: dict[str, _MCPServerRuntime] = {}
"""MCP runtime metadata, keyed by server name. Updated atomically on MCP lifecycle changes."""
self._mcp_server_runtime_view = MappingProxyType(self._mcp_server_runtime)
self._mcp_client_dict_view = _MCPClientDictView(self._mcp_server_runtime)
self._timeout_mismatch_warned = False
self._timeout_warn_lock = threading.Lock()
self._runtime_lock = asyncio.Lock()
self._mcp_starting: set[str] = set()
self._init_timeout_default = _resolve_timeout(
timeout=None,
env_name=MCP_INIT_TIMEOUT_ENV,
default=DEFAULT_MCP_INIT_TIMEOUT_SECONDS,
)
self._enable_timeout_default = _resolve_timeout(
timeout=None,
env_name=ENABLE_MCP_TIMEOUT_ENV,
default=DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS,
)
self._warn_on_timeout_mismatch(
self._init_timeout_default,
self._enable_timeout_default,
)
@property
def mcp_client_dict(self) -> Mapping[str, MCPClient]:
"""Read-only compatibility view for external callers that still read mcp_client_dict.
Note: Mutating this mapping is unsupported and will raise TypeError.
"""
return self._mcp_client_dict_view
@property
def mcp_server_runtime_view(self) -> Mapping[str, _MCPServerRuntime]:
"""Read-only view of MCP runtime metadata for external callers."""
return self._mcp_server_runtime_view
@property
def mcp_server_runtime(self) -> Mapping[str, _MCPServerRuntime]:
"""Backward-compatible read-only view (deprecated). Do not mutate.
Note: Mutations are not supported and will raise TypeError.
"""
return self._mcp_server_runtime_view
def empty(self) -> bool:
return len(self.func_list) == 0
def spec_to_func(
self,
name: str,
func_args: list[dict],
desc: str,
handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]],
) -> FuncTool:
params = {
"type": "object", # hard-coded here
"properties": {},
}
for param in func_args:
p = copy.deepcopy(param)
p.pop("name", None)
params["properties"][param["name"]] = p
return FuncTool(
name=name,
parameters=params,
description=desc,
handler=handler,
)
def add_func(
self,
name: str,
func_args: list,
desc: str,
handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]],
) -> None:
"""添加函数调用工具
@param name: 函数名
@param func_args: 函数参数列表,格式为 [{"type": "string", "name": "arg_name", "description": "arg_description"}, ...]
@param desc: 函数描述
@param func_obj: 处理函数
"""
# check if the tool has been added before
self.remove_func(name)
self.func_list.append(
self.spec_to_func(
name=name,
func_args=func_args,
desc=desc,
handler=handler,
),
)
logger.info(f"添加函数调用工具: {name}")
def remove_func(self, name: str) -> None:
"""删除一个函数调用工具。"""
for i, f in enumerate(self.func_list):
if f.name == name:
self.func_list.pop(i)
break
def _find_mcp_tool_by_original_name(self, name: str) -> MCPTool | None:
"""按原始(命名空间前)名查找 MCP 工具,active 优先。"""
matches = [
f
for f in self.func_list
if isinstance(f, MCPTool) and f.original_tool_name == name
]
if not matches:
return None
active_matches = [f for f in matches if getattr(f, "active", True)]
candidates = active_matches or matches
if len(candidates) > 1:
candidates.sort(key=lambda f: f.name)
logger.warning(
f"Multiple MCP tools found with original name '{name}': "
f"{[f.name for f in candidates]}. Using {candidates[0].name}"
)
return candidates[0]
def get_func(self, name) -> FuncTool | None:
# 优先返回已激活的工具(后加载的覆盖前面的,与 ToolSet.add_tool 保持一致)
# 使用 getattr(..., True) 与 ToolSet.add_tool 保持一致:没有 active 属性的工具视为已激活
for f in reversed(self.func_list):
if f.name == name and getattr(f, "active", True):
return f
# 退化则拿最后一个同名工具
for f in reversed(self.func_list):
if f.name == name:
return f
if isinstance(name, str):
# 老 persona 配置按原始 MCP 名回查
mcp_tool = self._find_mcp_tool_by_original_name(name)
if mcp_tool is not None:
return mcp_tool
try:
builtin_tool = self.get_builtin_tool(name)
except KeyError:
return None
if getattr(builtin_tool, "active", True):
return builtin_tool
return builtin_tool
return None
def get_builtin_tool(self, tool: str | type[FuncTool]) -> FuncTool:
ensure_builtin_tools_loaded()
if isinstance(tool, str):
tool_cls = get_builtin_tool_class(tool)
if tool_cls is None:
raise KeyError(f"Builtin tool {tool} is not registered.")
elif isinstance(tool, type) and issubclass(tool, FunctionTool):
tool_cls = tool
if get_builtin_tool_name(tool_cls) is None:
raise KeyError(
f"Builtin tool class {tool_cls.__module__}.{tool_cls.__name__} is not registered.",
)
else:
raise TypeError("tool must be a builtin tool name or FunctionTool class.")
cached_tool = self.builtin_func_list.get(tool_cls)
if cached_tool is not None:
return cached_tool
builtin_tool = tool_cls() # type: ignore
self.builtin_func_list[tool_cls] = builtin_tool
return builtin_tool
def iter_builtin_tools(self) -> list[FuncTool]:
ensure_builtin_tools_loaded()
return [
self.get_builtin_tool(tool_cls) for tool_cls in iter_builtin_tool_classes()
]
def is_builtin_tool(self, name: str) -> bool:
ensure_builtin_tools_loaded()
return get_builtin_tool_class(name) is not None
def get_full_tool_set(self) -> ToolSet:
"""获取完整工具集
使用 ToolSet.add_tool 进行填充。对于同名工具,去重规则为:
- 优先保留 active=True 的工具;
- 当 active 状态相同时,后加载的工具会覆盖前面的工具。
因此,后加载的 inactive 工具不会覆盖已激活的工具;
同时,MCP 工具在需要时仍可覆盖被禁用的内置工具。
"""
tool_set = ToolSet()
for tool in self.func_list:
tool_set.add_tool(tool)
return tool_set
@staticmethod
def _log_safe_mcp_debug_config(cfg: dict) -> None:
# 仅记录脱敏后的摘要,避免泄露 command/args/url 中的敏感信息
if "command" in cfg:
cmd = cfg["command"]
executable = str(cmd[0] if isinstance(cmd, (list, tuple)) and cmd else cmd)
args_val = cfg.get("args", [])
args_count = (
len(args_val)
if isinstance(args_val, (list, tuple))
else (0 if args_val is None else 1)
)
logger.debug(f" 命令可执行文件: {executable}, 参数数量: {args_count}")
return
if "url" in cfg:
parsed = urllib.parse.urlparse(str(cfg["url"]))
host = parsed.hostname or ""
scheme = parsed.scheme or "unknown"
try:
port = f":{parsed.port}" if parsed.port else ""
except ValueError:
port = ""
logger.debug(f" 主机: {scheme}://{host}{port}")
async def init_mcp_clients(
self, raise_on_all_failed: bool = False
) -> MCPInitSummary:
"""从项目根目录读取 mcp_server.json 文件,初始化 MCP 服务列表。文件格式如下:
```
{
"mcpServers": {
"weather": {
"command": "uv",
"args": [
"--directory",
"/ABSOLUTE/PATH/TO/PARENT/FOLDER/weather",
"run",
"weather.py"
]
}
}
...
}
```
Timeout behavior:
- 初始化超时使用环境变量 ASTRBOT_MCP_INIT_TIMEOUT 或默认值。
- 动态启用超时使用 ASTRBOT_MCP_ENABLE_TIMEOUT(独立于初始化超时)。
"""
data_dir = get_astrbot_data_path()
mcp_json_file = os.path.join(data_dir, "mcp_server.json")
if not os.path.exists(mcp_json_file):
# 配置文件不存在错误处理
with open(mcp_json_file, "w", encoding="utf-8") as f:
json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4)
logger.info(f"未找到 MCP 服务配置文件,已创建默认配置文件 {mcp_json_file}")
return MCPInitSummary(total=0, success=0, failed=[])
with open(mcp_json_file, encoding="utf-8") as f:
mcp_server_json_obj: dict[str, dict] = json.load(f)["mcpServers"]
init_timeout = self._init_timeout_default
timeout_display = f"{init_timeout:g}"
active_configs: list[tuple[str, dict, asyncio.Event]] = []
for name, cfg in mcp_server_json_obj.items():
if cfg.get("active", True):
shutdown_event = asyncio.Event()
active_configs.append((name, cfg, shutdown_event))
if not active_configs:
return MCPInitSummary(total=0, success=0, failed=[])
logger.info(f"等待 {len(active_configs)} 个 MCP 服务初始化...")
init_tasks = [
asyncio.create_task(
self._start_mcp_server(
name=name,
cfg=cfg,
shutdown_event=shutdown_event,
timeout=init_timeout,
),
name=f"mcp-init:{name}",
)
for (name, cfg, shutdown_event) in active_configs
]
results = await asyncio.gather(*init_tasks, return_exceptions=True)
success_count = 0
failed_services: list[str] = []
for (name, cfg, _), result in zip(active_configs, results, strict=False):
if isinstance(result, Exception):
if isinstance(result, MCPInitTimeoutError):
logger.error(
f"Connected to MCP server {name} timeout ({timeout_display} seconds)"
)
else:
logger.error(f"Failed to initialize MCP server {name}: {result}")
self._log_safe_mcp_debug_config(cfg)
failed_services.append(name)
async with self._runtime_lock:
self._mcp_server_runtime.pop(name, None)
continue
success_count += 1
if failed_services:
logger.warning(
f"The following MCP services failed to initialize: {', '.join(failed_services)}. "
f"Please check the mcp_server.json file and server availability."
)
summary = MCPInitSummary(
total=len(active_configs), success=success_count, failed=failed_services
)
logger.info(
f"MCP services initialization completed: {summary.success}/{summary.total} successful, {len(summary.failed)} failed."
)
if summary.total > 0 and summary.success == 0:
msg = "All MCP services failed to initialize, please check the mcp_server.json and server availability."
if raise_on_all_failed:
raise MCPAllServicesFailedError(msg)
logger.error(msg)
return summary
async def _start_mcp_server(
self,
name: str,
cfg: dict,
*,
shutdown_event: asyncio.Event | None = None,
timeout: float,
) -> None:
"""Initialize MCP server with timeout and register task/event together.
This method is idempotent. If the server is already running, the existing
runtime is kept and the new config is ignored.
"""
async with self._runtime_lock:
if name in self._mcp_server_runtime or name in self._mcp_starting:
logger.warning(
f"Connected to MCP server {name}, ignoring this startup request (timeout={timeout:g})."
)
self._log_safe_mcp_debug_config(cfg)
return
self._mcp_starting.add(name)
if shutdown_event is None:
shutdown_event = asyncio.Event()
mcp_client: MCPClient | None = None
try:
mcp_client = await asyncio.wait_for(
self._init_mcp_client(name, cfg),
timeout=timeout,
)
except asyncio.TimeoutError as exc:
raise MCPInitTimeoutError(
f"Connected to MCP server {name} timeout ({timeout:g} seconds)"
) from exc
except Exception:
logger.error(f"Failed to initialize MCP client {name}", exc_info=True)
raise
finally:
if mcp_client is None:
async with self._runtime_lock:
self._mcp_starting.discard(name)
async def lifecycle() -> None:
try:
await shutdown_event.wait()
logger.info(f"Received shutdown signal for MCP client {name}")
except asyncio.CancelledError:
logger.debug(f"MCP client {name} task was cancelled")
raise
finally:
await self._terminate_mcp_client(name)
lifecycle_task = asyncio.create_task(lifecycle(), name=f"mcp-client:{name}")
async with self._runtime_lock:
# After successful initialization mcp_client is logically non-None;
# raise explicitly so this invariant still holds under `python -O`
# (which strips `assert`).
if mcp_client is None:
raise RuntimeError(
f"MCP client {name} unexpectedly None after successful initialization"
)
self._mcp_server_runtime[name] = _MCPServerRuntime(
name=name,
client=mcp_client,
shutdown_event=shutdown_event,
lifecycle_task=lifecycle_task,
)
self._mcp_starting.discard(name)
async def _shutdown_runtimes(
self,
runtimes: list[_MCPServerRuntime],
timeout: float,
*,
strict: bool = True,
) -> list[str]:
"""Shutdown runtimes and wait for lifecycle tasks to complete."""
lifecycle_tasks = [
runtime.lifecycle_task
for runtime in runtimes
if not runtime.lifecycle_task.done()
]
if not lifecycle_tasks:
return []
for runtime in runtimes:
runtime.shutdown_event.set()
try:
results = await asyncio.wait_for(
asyncio.gather(*lifecycle_tasks, return_exceptions=True),
timeout=timeout,
)
except asyncio.TimeoutError:
pending_names = [
runtime.name
for runtime in runtimes
if not runtime.lifecycle_task.done()
]
for task in lifecycle_tasks:
if not task.done():
task.cancel()
await asyncio.gather(*lifecycle_tasks, return_exceptions=True)
if strict:
raise MCPShutdownTimeoutError(pending_names, timeout)
logger.warning(
"MCP server shutdown timeout (%s seconds), the following servers were not fully closed: %s",
f"{timeout:g}",
", ".join(pending_names),
)
return pending_names
else:
for result in results:
if isinstance(result, asyncio.CancelledError):
logger.debug("MCP lifecycle task was cancelled during shutdown.")
elif isinstance(result, Exception):
logger.error(
"MCP lifecycle task failed during shutdown.",
exc_info=(type(result), result, result.__traceback__),
)
return []
async def _cleanup_mcp_client_safely(
self, mcp_client: MCPClient, name: str
) -> None:
"""安全清理单个 MCP 客户端,避免清理异常中断主流程。"""
try:
await mcp_client.cleanup()
except Exception as cleanup_exc: # noqa: BLE001 - only log here
logger.error(
f"Failed to cleanup MCP client resources {name}: {cleanup_exc}"
)
async def _init_mcp_client(self, name: str, config: dict) -> MCPClient:
"""初始化单个MCP客户端"""
mcp_client = MCPClient()
mcp_client.name = name
try:
await mcp_client.connect_to_server(config, name)
tools_res = await mcp_client.list_tools_and_save()
except asyncio.CancelledError:
await self._cleanup_mcp_client_safely(mcp_client, name)
raise
except Exception:
await self._cleanup_mcp_client_safely(mcp_client, name)
raise
logger.debug(f"MCP server {name} list tools response: {tools_res}")
tool_names = [tool.name for tool in tools_res.tools]
# 移除该MCP服务之前的工具(如有)
self.func_list = [
f
for f in self.func_list
if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
]
# 将 MCP 工具转换为 FuncTool 并添加到 func_list
for tool in mcp_client.tools:
func_tool = MCPTool(
mcp_tool=tool,
mcp_client=mcp_client,
mcp_server_name=name,
)
self.func_list.append(func_tool)
logger.info(f"Connected to MCP server {name}, Tools: {tool_names}")
return mcp_client
async def _terminate_mcp_client(self, name: str) -> None:
"""关闭并清理MCP客户端"""
async with self._runtime_lock:
runtime = self._mcp_server_runtime.get(name)
if runtime:
client = runtime.client
# 关闭MCP连接
await self._cleanup_mcp_client_safely(client, name)
# 移除关联的FuncTool
self.func_list = [
f
for f in self.func_list
if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
]
async with self._runtime_lock:
self._mcp_server_runtime.pop(name, None)
self._mcp_starting.discard(name)
logger.info(f"Disconnected from MCP server {name}")
return
# Runtime missing but stale tools may still exist after failed flows.
self.func_list = [
f
for f in self.func_list
if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
]
async with self._runtime_lock:
self._mcp_starting.discard(name)
@staticmethod
async def test_mcp_server_connection(config: dict) -> list[str]:
if "url" in config:
success, error_msg = await _quick_test_mcp_connection(config)
if not success:
raise Exception(error_msg)
mcp_client = MCPClient()
try:
logger.debug(f"testing MCP server connection with config: {config}")
await mcp_client.connect_to_server(config, "test")
tools_res = await mcp_client.list_tools_and_save()
tool_names = [tool.name for tool in tools_res.tools]
finally:
logger.debug("Cleaning up MCP client after testing connection.")
await mcp_client.cleanup()
return tool_names
async def enable_mcp_server(
self,
name: str,
config: dict,
shutdown_event: asyncio.Event | None = None,
timeout: float | int | str | None = None,
) -> None:
"""Enable a new MCP server and initialize it.
Args:
name: The name of the MCP server.
config: Configuration for the MCP server.
shutdown_event: Event to signal when the MCP client should shut down.
timeout: Timeout in seconds for initialization.
Uses ASTRBOT_MCP_ENABLE_TIMEOUT by default (separate from init timeout).
Raises:
MCPInitTimeoutError: If initialization does not complete within timeout.
Exception: If there is an error during initialization.
"""
if timeout is None:
timeout_value = self._enable_timeout_default
else:
timeout_value = _resolve_timeout(
timeout=timeout,
env_name=ENABLE_MCP_TIMEOUT_ENV,
default=self._enable_timeout_default,
)
await self._start_mcp_server(
name=name,
cfg=config,
shutdown_event=shutdown_event,
timeout=timeout_value,
)
async def disable_mcp_server(
self,
name: str | None = None,
timeout: float = 10,
) -> None:
"""Disable an MCP server by its name.
Args:
name (str): The name of the MCP server to disable. If None, ALL MCP servers will be disabled.
timeout (int): Timeout.
Raises:
MCPShutdownTimeoutError: If shutdown does not complete within timeout.
Only raised when disabling a specific server (name is not None).
"""
if name:
async with self._runtime_lock:
runtime = self._mcp_server_runtime.get(name)
if runtime is None:
return
await self._shutdown_runtimes([runtime], timeout, strict=True)
else:
async with self._runtime_lock:
runtimes = list(self._mcp_server_runtime.values())
await self._shutdown_runtimes(runtimes, timeout, strict=False)
def _warn_on_timeout_mismatch(
self,
init_timeout: float,
enable_timeout: float,
) -> None:
if init_timeout == enable_timeout:
return
with self._timeout_warn_lock:
if self._timeout_mismatch_warned:
return
logger.info(
"检测到 MCP 初始化超时与动态启用超时配置不同:"
"初始化使用 %s 秒,动态启用使用 %s 秒。如需一致,请设置相同值。",
f"{init_timeout:g}",
f"{enable_timeout:g}",
)
self._timeout_mismatch_warned = True
def get_func_desc_openai_style(self, omit_empty_parameter_field=False) -> list:
"""获得 OpenAI API 风格的**已经激活**的工具描述"""
tools = [f for f in self.func_list if f.active]
toolset = ToolSet(tools)
return toolset.openai_schema(
omit_empty_parameter_field=omit_empty_parameter_field,
)
def get_func_desc_anthropic_style(self) -> list:
"""获得 Anthropic API 风格的**已经激活**的工具描述"""
tools = [f for f in self.func_list if f.active]
toolset = ToolSet(tools)
return toolset.anthropic_schema()
def get_func_desc_google_genai_style(self) -> dict:
"""获得 Google GenAI API 风格的**已经激活**的工具描述"""
tools = [f for f in self.func_list if f.active]
toolset = ToolSet(tools)
return toolset.google_schema()
def deactivate_llm_tool(self, name: str) -> bool:
"""停用一个已经注册的函数调用工具。
Returns:
如果没找到,会返回 False
"""
func_tool = self.get_func(name)
if func_tool is not None:
func_tool.active = False
inactivated_llm_tools: list = sp.get(
"inactivated_llm_tools",
[],
scope="global",
scope_id="global",
)
if name not in inactivated_llm_tools:
inactivated_llm_tools.append(name)
sp.put(
"inactivated_llm_tools",
inactivated_llm_tools,
scope="global",
scope_id="global",
)
return True
return False
# 因为不想解决循环引用,所以这里直接传入 star_map 先了...
def activate_llm_tool(self, name: str, star_map: dict) -> bool:
func_tool = self.get_func(name)
if func_tool is not None:
if func_tool.handler_module_path in star_map:
if not star_map[func_tool.handler_module_path].activated:
raise ValueError(
f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用,请先在管理面板启用再激活此工具。",
)
func_tool.active = True
inactivated_llm_tools: list = sp.get(
"inactivated_llm_tools",
[],
scope="global",
scope_id="global",
)
if name in inactivated_llm_tools:
inactivated_llm_tools.remove(name)
sp.put(
"inactivated_llm_tools",
inactivated_llm_tools,
scope="global",
scope_id="global",
)
return True
return False
@property
def mcp_config_path(self):
data_dir = get_astrbot_data_path()
return os.path.join(data_dir, "mcp_server.json")
def load_mcp_config(self):
if not os.path.exists(self.mcp_config_path):
# 配置文件不存在,创建默认配置
os.makedirs(os.path.dirname(self.mcp_config_path), exist_ok=True)
with open(self.mcp_config_path, "w", encoding="utf-8") as f:
json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4)
return DEFAULT_MCP_CONFIG
try:
with open(self.mcp_config_path, encoding="utf-8") as f:
return json.load(f)
except Exception as e:
logger.error(f"加载 MCP 配置失败: {e}")
return DEFAULT_MCP_CONFIG
def save_mcp_config(self, config: dict) -> bool:
try:
with open(self.mcp_config_path, "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=4)
return True
except Exception as e:
logger.error(f"保存 MCP 配置失败: {e}")
return False
async def sync_modelscope_mcp_servers(self, access_token: str) -> None:
"""从 ModelScope 平台同步 MCP 服务器配置"""
base_url = "https://www.modelscope.cn/openapi/v1"
url = f"{base_url}/mcp/servers/operational"
headers = {
"Authorization": f"Bearer {access_token.strip()}",
"Content-Type": "application/json",
}
try:
async with aiohttp.ClientSession() as session:
async with session.get(url, headers=headers) as response:
if response.status == 200:
data = await response.json()
mcp_server_list = data.get("data", {}).get(
"mcp_server_list",
[],
)
local_mcp_config = self.load_mcp_config()
synced_count = 0
for server in mcp_server_list:
server_name = server["name"]
operational_urls = server.get("operational_urls", [])
if not operational_urls:
continue
url_info = operational_urls[0]
server_url = url_info.get("url")
if not server_url:
continue
# 添加到配置中(同名会覆盖)
local_mcp_config["mcpServers"][server_name] = {
"url": server_url,
"transport": "sse",
"active": True,
"provider": "modelscope",
}
synced_count += 1
if synced_count > 0:
self.save_mcp_config(local_mcp_config)
tasks = []
for server in mcp_server_list:
name = server["name"]
tasks.append(
self.enable_mcp_server(