fix(extension): 修复 MCP Host 资源与协议边界

This commit is contained in:
2026-09-01 12:11:30 +08:00
parent aedb1c1267
commit 1132a4cece
15 changed files with 138 additions and 24 deletions
+1 -1
View File
@@ -23,7 +23,7 @@ uv run uvicorn app.main:app --reload --host 127.0.0.1 --port 8000
uv run pytest
```
当前基线为 87 项测试通过。Provider API Key 可通过前端设置页写入,也可用 `OPENAI_API_KEY``DEEPSEEK_API_KEY``AINOTE_CREDENTIAL_<ID>` 注入;不要把真实密钥写入仓库。
当前基线为 91 项测试通过。Provider API Key 可通过前端设置页写入,也可用 `OPENAI_API_KEY``DEEPSEEK_API_KEY``AINOTE_CREDENTIAL_<ID>` 注入;不要把真实密钥写入仓库。
团队接口清单见 `../docs/contracts/后端接口契约-开发版.md`,机器可读契约以运行时的 `/openapi.json` 为准。
+33 -6
View File
@@ -188,7 +188,7 @@ class McpStdioClient:
result = response.get("result")
if not isinstance(result, dict):
raise McpBridgeError(
"MCP_TOOL_CALL_FAILED", "MCP response result must be an object."
response_error_code, "MCP response result must be an object."
)
return result
@@ -207,9 +207,18 @@ class McpStdioClient:
except McpBridgeError:
pass
def abandon(self, request_id: int) -> None:
def abandon(
self, request_id: int, wake_error: BaseException | None = None
) -> None:
with self._pending_lock:
self._pending.pop(request_id, None)
pending = self._pending.pop(request_id, None)
# asyncio.to_thread 被取消时不会停止底层线程;主动唤醒 Queue,避免线程
# 一直占用默认线程池直至远端超时。
if pending is not None and wake_error is not None:
try:
pending.response.put_nowait(wake_error)
except queue.Full:
pass
def stop(self) -> None:
process = self.process
@@ -258,7 +267,15 @@ class McpStdioClient:
assert process is not None and process.stdout is not None
failure: str | None = None
try:
for raw_line in process.stdout:
while True:
# readline(size) 在换行缺失时仍有硬上限,不能先把任意大的
# 第三方 stdout 行完整读入宿主内存再检查。
raw_line = process.stdout.readline(MAX_MCP_MESSAGE_BYTES + 1)
if raw_line == "":
break
if not raw_line.endswith("\n"):
failure = "MCP server emitted an oversized or unterminated message."
break
if len(raw_line.encode("utf-8")) > MAX_MCP_MESSAGE_BYTES:
failure = "MCP server emitted an oversized protocol message."
break
@@ -313,7 +330,12 @@ class McpStdioClient:
process = self.process
assert process is not None and process.stderr is not None
try:
for line in process.stderr:
while True:
# stderr 不是协议通道,但同样按块读取,避免无换行日志造成
# 宿主侧的无界字符串分配。
line = process.stderr.readline(1025)
if line == "":
break
self._stderr_tail.append(line.rstrip()[:1024])
except (OSError, ValueError):
return
@@ -489,7 +511,12 @@ class McpBridge:
)
except asyncio.CancelledError:
host.client.cancel(rpc_id)
host.client.abandon(rpc_id)
host.client.abandon(
rpc_id,
McpBridgeError(
"MCP_TOOL_CALL_FAILED", "MCP request was cancelled."
),
)
raise
except McpBridgeError as exc:
raise ToolExecutionError(exc.code, exc.message) from exc
+14 -1
View File
@@ -479,6 +479,17 @@ class PluginRuntime:
status_code=409,
details={"plugin_id": plugin_id},
)
if record.plugin.status in {
PluginStatus.installed,
PluginStatus.disabled,
PluginStatus.permission_required,
}:
raise ExtensionError(
"PLUGIN_HOST_UNAVAILABLE",
"Disabled or inactive MCP Plugins must be started with Enable.",
status_code=409,
details={"plugin_id": plugin_id, "status": record.plugin.status.value},
)
for name in record.registered_tools:
self.registry.unregister(name)
record.registered_tools.clear()
@@ -664,7 +675,9 @@ def _arguments_model_from_schema(
annotation = types.get(field_schema.get("type"), Any)
fields[name] = (annotation, ... if name in required else None)
model_name = "PluginArgs_" + re.sub(r"\W+", "_", tool_name)
return create_model(model_name, __config__=ConfigDict(extra="forbid"), **fields)
# 完整 JSON Schema 已在 ToolRegistry 中先行校验。这里允许额外字段,避免
# Pydantic 再次拒绝 additionalProperties/patternProperties 接受的合法参数。
return create_model(model_name, __config__=ConfigDict(extra="allow"), **fields)
def _validate_tool_schema(spec: DeclarativeToolSpec) -> None:
@@ -135,6 +135,15 @@ def main() -> None:
request_id = message.get("id")
params = message.get("params") or {}
if method == "initialize" and isinstance(request_id, int):
if MODE == "invalid-result":
send({"jsonrpc": "2.0", "id": request_id, "result": None})
continue
if MODE == "oversized-stdout":
# 不带换行,验证 Host 在读取完整内容前执行硬上限。
sys.stdout.write("x" * (2 * 1024 * 1024 + 1))
sys.stdout.flush()
time.sleep(10)
return
respond(
request_id,
{
+61 -1
View File
@@ -1,5 +1,6 @@
import asyncio
import shutil
import threading
import time
import pytest
@@ -14,6 +15,8 @@ from app.contracts import (
ToolCall,
)
from app.extensions import ExtensionError
from app.extensions.mcp import McpStdioClient
from app.extensions.runtime import _arguments_model_from_schema
from app.services import note_service
from app.config import BACKEND_DIR, get_settings
@@ -362,6 +365,11 @@ def test_mcp_stdio_host_discovers_namespaced_tools_and_maps_results(
assert disabled.status == "disabled"
assert mcp_container.plugins.get_host_status("mcp-fixture").status == "stopped"
assert not mcp_container.tools.contains("mcp-fixture.echo")
with pytest.raises(ExtensionError) as exc:
mcp_container.plugins.restart_host("mcp-fixture")
assert exc.value.code == "PLUGIN_HOST_UNAVAILABLE"
assert mcp_container.plugins.get("mcp-fixture").status == "disabled"
assert not mcp_container.tools.contains("mcp-fixture.echo")
run(scenario())
@@ -438,6 +446,56 @@ def test_mcp_business_error_size_limit_and_timeout_are_structured(mcp_container)
run(scenario())
def test_mcp_cancel_releases_blocking_response_thread(
mcp_container, monkeypatch
) -> None:
async def scenario() -> None:
mcp_container.plugins.enable("mcp-fixture")
released = threading.Event()
original_wait = McpStdioClient.wait_response
def tracked_wait(self, *args, **kwargs):
try:
return original_wait(self, *args, **kwargs)
finally:
released.set()
monkeypatch.setattr(McpStdioClient, "wait_response", tracked_wait)
task = asyncio.create_task(
mcp_container.plugins.mcp.call_tool(
"mcp-fixture",
"sleep",
{"seconds": 5},
request_id="call_cancel_release",
)
)
await asyncio.sleep(0.05)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
deadline = time.monotonic() + 0.5
while not released.is_set() and time.monotonic() < deadline:
await asyncio.sleep(0.01)
assert released.is_set(), "cancelled MCP wait must not occupy a worker until timeout"
run(scenario())
def test_mcp_argument_model_preserves_json_schema_additional_properties() -> None:
arguments_model = _arguments_model_from_schema(
"mcp-fixture.dynamic",
{
"type": "object",
"additionalProperties": {"type": "string"},
},
)
arguments = arguments_model.model_validate({"dynamic_key": "value"})
assert arguments.model_dump() == {"dynamic_key": "value"}
def test_mcp_abnormal_exit_unregisters_tools_and_restart_recovers(mcp_container) -> None:
async def scenario() -> None:
mcp_container.plugins.enable("mcp-fixture")
@@ -471,9 +529,11 @@ def test_mcp_abnormal_exit_unregisters_tools_and_restart_recovers(mcp_container)
[
("no-tools", "[]", "MCP_CAPABILITY_UNSUPPORTED"),
("invalid-schema", "[mcp-invalid.broken]", "MCP_TOOL_SCHEMA_INVALID"),
("invalid-result", "[]", "MCP_INITIALIZE_FAILED"),
("oversized-stdout", "[]", "PLUGIN_HOST_UNAVAILABLE"),
],
)
def test_mcp_rejects_missing_capability_and_invalid_discovery(
def test_mcp_rejects_invalid_initialization_and_discovery(
tmp_path, mode, contributions, expected_code
) -> None:
package = tmp_path / f"mcp-{mode}"