This is an automated email from the ASF dual-hosted git repository.
FreeOnePlus pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
The following commit(s) were added to refs/heads/master by this push:
new f3d43fa feat: add allowlisted custom tool providers (#162)
f3d43fa is described below
commit f3d43fa36b9f86038a1a0034db0162aa5167626c
Author: Yijia Su <[email protected]>
AuthorDate: Thu Jul 30 20:02:42 2026 +0800
feat: add allowlisted custom tool providers (#162)
---
.env.example | 6 +
CHANGELOG.md | 3 +
README.md | 10 +
README.zh-CN.md | 6 +
docs/custom-tool-providers.md | 199 +++++++++++++
doris_mcp_server/auth/operation_policy.py | 14 +-
doris_mcp_server/main.py | 1 +
doris_mcp_server/tools/__init__.py | 14 +
doris_mcp_server/tools/tool_catalog.py | 9 +-
doris_mcp_server/tools/tool_provider.py | 371 ++++++++++++++++++++++++
doris_mcp_server/tools/tool_registry.py | 36 ++-
doris_mcp_server/tools/tools_manager.py | 56 +++-
doris_mcp_server/utils/config.py | 23 ++
test/protocol/test_multiworker_config.py | 29 ++
test/security/test_operation_policy.py | 31 ++
test/test_product_identity.py | 16 ++
test/tools/test_custom_tool_provider.py | 456 ++++++++++++++++++++++++++++++
17 files changed, 1271 insertions(+), 9 deletions(-)
diff --git a/.env.example b/.env.example
index 0f731a3..84b4686 100644
--- a/.env.example
+++ b/.env.example
@@ -67,6 +67,12 @@ DORIS_MAX_CONNECTION_AGE=3600
# FE_ARROW_FLIGHT_SQL_PORT=
# BE_ARROW_FLIGHT_SQL_PORT=
+# Explicit allowlist of installed Python entry points from the
+# doris_mcp_server.tool_providers group. Empty means no custom code is loaded.
+# Custom providers are supported with local/stdio, static token, and JWT auth;
+# they are hidden from OAuth modes until they have a reviewed OAuth policy.
+MCP_TOOL_PROVIDERS=
+
# ===================================================================
# Security Configuration
# ===================================================================
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 18b3470..3e8daf8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -42,6 +42,9 @@ under **Unreleased** until a new version is selected and
published.
- W3C `traceparent`, `tracestate`, and `baggage` propagation from request
`_meta`, with value-safe validation, credential-like baggage redaction,
per-request isolation, and no trace metadata in model-facing results.
+- Explicitly installed and allowlisted custom tool providers for existing
+ business APIs, with lifecycle management, schema validation, safe audit
+ metadata, and bounded process-local rate limits.
- Bounded query result streaming with deployment and absolute ceilings for
rows, serialized bytes, and execution time, plus cancellation-safe database
connection disposal.
diff --git a/README.md b/README.md
index dad596a..a8a88b8 100644
--- a/README.md
+++ b/README.md
@@ -220,6 +220,8 @@ export ENABLE_LEGACY_HTTP_ADAPTER=false
# Bound each resources/list, tools/list, and prompts/list response.
export MCP_LIST_PAGE_SIZE=100
+# Load only these installed custom tool providers. Empty disables extensions.
+export MCP_TOOL_PROVIDERS="orders_api"
# A launch-local key is generated automatically. Configure one shared
# high-entropy value when independently launched replicas share traffic.
export MCP_STATE_HANDLE_SECRET="$(python -c 'import secrets;
print(secrets.token_urlsafe(32))')"
@@ -316,6 +318,8 @@ cp .env.example .env
modern traffic always uses `POST /mcp`
* `MCP_LIST_PAGE_SIZE`: Maximum resources, tools, or prompts returned
per protocol page (default: 100; range: 1-1000)
+ * `MCP_TOOL_PROVIDERS`: Comma-separated allowlist of installed
+ `doris_mcp_server.tool_providers` entry points (default: empty)
* `MCP_STATE_HANDLE_SECRET`: Optional shared high-entropy key (at least
32 bytes) used to authenticate explicit cross-call state handles
* `MCP_STATE_HANDLE_TTL_SECONDS`: Lifetime of an explicit state handle
@@ -1498,6 +1502,12 @@ doris-mcp-server/
This section outlines the process for adding new MCP tools to the Doris MCP
Server, based on the unified modular architecture with centralized tool
management.
+Existing business APIs do not need to be built into this repository. Package
+them as explicitly installed, allowlisted custom tool providers instead. The
+[custom tool provider guide](docs/custom-tool-providers.md) defines the entry
+point contract, lifecycle, process-local QPS limits, authentication boundary,
+FastGPT integration, and production security checklist.
+
### 1. Leverage Existing Utility Modules
The server provides comprehensive utility modules for common database
operations:
diff --git a/README.zh-CN.md b/README.zh-CN.md
index 68f097d..ab05c1b 100644
--- a/README.zh-CN.md
+++ b/README.zh-CN.md
@@ -174,6 +174,7 @@ curl --fail http://127.0.0.1:3000/ready
| `MCP_STATE_HANDLE_SECRET` | 多副本共享的状态句柄签名密钥 | 启动时生成 |
| `MCP_ALLOWED_HOSTS` | HTTP Host 白名单 | 回环地址 |
| `ENABLE_LEGACY_HTTP_ADAPTER` | 启用 `/mcp/legacy` 迁移端点 | `false` |
+| `MCP_TOOL_PROVIDERS` | 已安装自定义工具 Provider 的显式白名单 | 空 |
多副本经负载均衡提供分页时，应配置同一个高强度
`MCP_STATE_HANDLE_SECRET`。分页 Cursor 是带签名的状态句柄，客户端不应解析或
@@ -446,6 +447,11 @@ uv run bandit -q -c pyproject.toml -r doris_mcp_server
4. 权限与审计元数据；
5. 单元测试及必要的真实 Doris 集成测试。
+已有业务 API 不必直接合入本仓库。可以把它封装为安装在服务运行环境中的
+自定义工具 Provider，再通过 `MCP_TOOL_PROVIDERS` 显式启用。入口协议、生命周期、
+进程内 QPS 限制、FastGPT 接入方式和生产安全边界见
+[自定义工具 Provider 指南](docs/custom-tool-providers.md)。
+
详细流程见[英文开发指南](README.md#developing-new-tools)。
## 常见问题
diff --git a/docs/custom-tool-providers.md b/docs/custom-tool-providers.md
new file mode 100644
index 0000000..54f1179
--- /dev/null
+++ b/docs/custom-tool-providers.md
@@ -0,0 +1,199 @@
+<!--
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements. See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership. The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied. See the License for the
+specific language governing permissions and limitations
+under the License.
+-->
+
+# Custom tool providers
+
+Doris MCP Server can expose an existing business API as an MCP tool without
+adding business-specific code to this repository. The extension is a trusted
+Python package installed in the same environment as the server and enabled by
+an explicit deploy-time allowlist.
+
+This boundary intentionally does not accept an arbitrary target URL, method,
+headers, or credentials from MCP clients. Provider code executes inside the
+server process and therefore has the same operating-system privileges as the
+server. Install and enable only packages that the operator has reviewed.
+
+## Provider contract
+
+Declare a named entry point in the provider package:
+
+```toml
+[project.entry-points."doris_mcp_server.tool_providers"]
+orders_api = "acme_doris_tools.provider:create_provider"
+```
+
+The entry-point target must be a zero-argument callable. It returns an object
+whose `name` exactly matches the entry-point name and whose `tools()` method
+returns one or more `CustomTool` definitions.
+
+```python
+from __future__ import annotations
+
+import os
+from typing import Any
+from urllib.parse import quote
+
+import httpx
+from mcp.types import Tool
+
+from doris_mcp_server.tools import CustomTool, ToolRateLimit
+
+
+class OrdersApiProvider:
+ name = "orders_api"
+
+ def __init__(self) -> None:
+ self._base_url = os.environ["ORDERS_API_BASE_URL"].rstrip("/")
+ if not self._base_url.startswith("https://";):
+ raise ValueError("ORDERS_API_BASE_URL must use HTTPS")
+ self._token = os.environ["ORDERS_API_TOKEN"]
+ self._client: httpx.AsyncClient | None = None
+
+ async def start(self) -> None:
+ self._client = httpx.AsyncClient(
+ base_url=self._base_url,
+ headers={"Authorization": f"Bearer {self._token}"},
+ timeout=httpx.Timeout(5.0, connect=2.0),
+ follow_redirects=False,
+ )
+
+ async def close(self) -> None:
+ if self._client is not None:
+ await self._client.aclose()
+ self._client = None
+
+ def tools(self) -> tuple[CustomTool, ...]:
+ return (
+ CustomTool(
+ tool=Tool(
+ name="lookup_business_order",
+ description="Look up one order in the reviewed business
API",
+ input_schema={
+ "type": "object",
+ "properties": {
+ "order_id": {
+ "type": "string",
+ "minLength": 1,
+ "maxLength": 128,
+ }
+ },
+ "required": ["order_id"],
+ "additionalProperties": False,
+ },
+ ),
+ handler=self._lookup_order,
+ risk="medium",
+ rate_limit=ToolRateLimit(
+ max_calls=10,
+ period_seconds=1,
+ scope="principal",
+ ),
+ ),
+ )
+
+ async def _lookup_order(
+ self,
+ arguments: dict[str, Any],
+ ) -> dict[str, Any]:
+ if self._client is None:
+ raise RuntimeError("provider is not started")
+ order_id = arguments["order_id"]
+ if not isinstance(order_id, str) or not order_id.strip():
+ raise ValueError("order_id must be a non-empty string")
+
+ response = await self._client.get(
+ f"/v1/orders/{quote(order_id, safe='')}"
+ )
+ response.raise_for_status()
+ if len(response.content) > 1_048_576:
+ raise ValueError("upstream response is too large")
+ payload = response.json()
+ if not isinstance(payload, dict):
+ raise ValueError("upstream response must be a JSON object")
+ return payload
+
+
+def create_provider() -> OrdersApiProvider:
+ return OrdersApiProvider()
+```
+
+Install the provider package in the same virtual environment, then enable only
+the reviewed entry-point name:
+
+```bash
+export MCP_TOOL_PROVIDERS=orders_api
+doris-mcp-server --transport http --host 127.0.0.1 --port 3000
+```
+
+An empty `MCP_TOOL_PROVIDERS` value loads no custom code. Startup fails if a
+configured provider is missing, has a mismatched name, returns an invalid
+definition, or shadows a built-in or another custom tool.
+
+Optional synchronous or asynchronous `start()` and `close()` methods can own
+HTTP clients or other runtime resources. A partial startup failure closes
+providers that were already started.
+
+## Authentication and rate limits
+
+Custom provider tools are available to local/stdio, static-token, and JWT
+deployments under their existing authorization boundary. External OAuth and
+Doris-backed OAuth hide custom tools and reject direct calls: they fail closed
+until the project defines a reviewed dynamic scope and policy model.
+
+`ToolRateLimit` is a bounded, process-local fixed-window limiter:
+
+- `scope="principal"` separates callers by the authenticated user or token
+ identity;
+- `scope="tool"` applies one shared limit to the tool in that server process;
+- rejected calls return `TOOL_RATE_LIMITED` and a bounded retry delay;
+- the provider handler is not invoked for a rejected call.
+
+Each worker and each replica has an independent limiter. For a global QPS,
+concurrency, quota, or cost limit, enforce the limit in an API gateway or
shared
+rate-limit service. With four workers, a process-local limit can allow roughly
+four times the configured rate.
+
+## FastGPT and other MCP clients
+
+Expose Doris MCP Server through its normal Streamable HTTP or stdio transport.
+FastGPT and other MCP clients discover the custom schema through `tools/list`
+and invoke it through `tools/call`; no provider-specific protocol is required.
+The same MCP `2026-07-28` request metadata and HTTP headers described in the
+main README still apply.
+
+## Production checklist
+
+- Keep target hosts and API paths fixed in reviewed provider code. Never accept
+ an arbitrary URL, scheme, host, redirect target, or authorization header from
+ MCP arguments.
+- Use HTTPS and a network egress allowlist. Disable redirects unless every
+ redirect target is separately validated.
+- Read credentials from environment variables or a secret manager. Never place
+ secrets in tool schemas, results, audit fields, or logs.
+- Apply bounded connect/read/total timeouts and response-size limits. Validate
+ the upstream content type and response structure.
+- Return only fields the MCP caller is authorized to see. Treat upstream error
+ bodies as sensitive; the server intentionally sanitizes uncaught exceptions.
+- Define retry and idempotency behavior explicitly. Do not automatically retry
+ unsafe mutations.
+- Add upstream concurrency limits, circuit breaking, and global rate limiting
+ outside the process when availability or cost requires them.
+- Test provider lifecycle, input validation, tool-name collisions, rate-limit
+ behavior, authentication visibility, audit metadata, and the real upstream
+ failure modes before production deployment.
diff --git a/doris_mcp_server/auth/operation_policy.py
b/doris_mcp_server/auth/operation_policy.py
index 7c5fba1..8bfad49 100644
--- a/doris_mcp_server/auth/operation_policy.py
+++ b/doris_mcp_server/auth/operation_policy.py
@@ -418,8 +418,18 @@ def authorize_operation(auth_context: Any | None,
operation: str) -> None:
def filter_tools_for_auth_context(auth_context: Any | None, tools: list[Any])
-> list[Any]:
- """Filter visible tools for Doris OAuth users."""
- if auth_context is None or auth_context.auth_method != "doris_oauth":
+ """Filter tools that do not have a reviewed policy for the OAuth mode."""
+ if auth_context is None:
+ return tools
+
+ if auth_context.auth_method == "external_oauth":
+ return [
+ tool
+ for tool in tools
+ if policy_definition_for_tool(getattr(tool, "name", "")) is not
None
+ ]
+
+ if auth_context.auth_method != "doris_oauth":
return tools
filtered = []
diff --git a/doris_mcp_server/main.py b/doris_mcp_server/main.py
index 6337fbb..40017fa 100644
--- a/doris_mcp_server/main.py
+++ b/doris_mcp_server/main.py
@@ -84,6 +84,7 @@ def _multiworker_environment(
config.enable_legacy_http_adapter
).lower(),
"MCP_LIST_PAGE_SIZE": str(config.mcp_list_page_size),
+ "MCP_TOOL_PROVIDERS": ",".join(config.mcp_tool_providers),
"MCP_STATE_HANDLE_SECRET": config.mcp_state_handle_secret,
"MCP_STATE_HANDLE_TTL_SECONDS":
str(config.mcp_state_handle_ttl_seconds),
"SERVER_NAME": config.server_name,
diff --git a/doris_mcp_server/tools/__init__.py
b/doris_mcp_server/tools/__init__.py
index 0c32176..70de15f 100644
--- a/doris_mcp_server/tools/__init__.py
+++ b/doris_mcp_server/tools/__init__.py
@@ -23,3 +23,17 @@ This package includes:
- Prompt managers
- Tool registration and initialization
"""
+
+from .tool_provider import (
+ CustomTool,
+ CustomToolProvider,
+ ToolProviderError,
+ ToolRateLimit,
+)
+
+__all__ = [
+ "CustomTool",
+ "CustomToolProvider",
+ "ToolProviderError",
+ "ToolRateLimit",
+]
diff --git a/doris_mcp_server/tools/tool_catalog.py
b/doris_mcp_server/tools/tool_catalog.py
index 8c26208..8956203 100644
--- a/doris_mcp_server/tools/tool_catalog.py
+++ b/doris_mcp_server/tools/tool_catalog.py
@@ -16,18 +16,21 @@
# under the License.
"""Immutable MCP tool catalog, separate from manager lifecycle and routing."""
+from collections.abc import Iterable
from typing import Any
from mcp.types import Tool
from ..result_limits import configured_default_result_rows,
configured_result_limits
from ..utils.config import ADBCConfig
+from .tool_provider import CustomTool
from .tool_registry import ToolDefinitionRegistry
def build_tool_registry(
handler_owner: Any,
config: Any | None,
+ custom_tools: Iterable[tuple[str, CustomTool]] = (),
) -> ToolDefinitionRegistry:
"""Build immutable tool metadata and bind it to a handler owner."""
adbc_config = getattr(config, "adbc", None) or ADBCConfig()
@@ -936,4 +939,8 @@ No parameters required. Returns connection status,
configuration, and diagnostic
),
]
- return ToolDefinitionRegistry.from_tools(tools, handler_owner)
+ return ToolDefinitionRegistry.from_tools(
+ tools,
+ handler_owner,
+ custom_tools=custom_tools,
+ )
diff --git a/doris_mcp_server/tools/tool_provider.py
b/doris_mcp_server/tools/tool_provider.py
new file mode 100644
index 0000000..49fe168
--- /dev/null
+++ b/doris_mcp_server/tools/tool_provider.py
@@ -0,0 +1,371 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Explicit, deploy-time extension boundary for trusted custom MCP tools."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+import re
+import time
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from dataclasses import dataclass
+from importlib import metadata
+from typing import Any, Literal, Protocol
+
+from mcp.types import Tool
+
+CUSTOM_TOOL_PROVIDER_ENTRY_POINT = "doris_mcp_server.tool_providers"
+MAX_PROVIDER_NAME_LENGTH = 128
+MAX_RATE_LIMIT_KEYS = 10_000
+_PROVIDER_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
+logger = logging.getLogger(__name__)
+
+ToolHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
+ToolRisk = Literal["low", "medium", "high"]
+RateLimitScope = Literal["tool", "principal"]
+
+
+class ToolProviderError(ValueError):
+ """Raised when an explicitly configured tool provider is invalid."""
+
+
+@dataclass(frozen=True)
+class ToolRateLimit:
+ """Process-local fixed-window limit for one custom tool."""
+
+ max_calls: int
+ period_seconds: float
+ scope: RateLimitScope = "principal"
+
+ def __post_init__(self) -> None:
+ if isinstance(self.max_calls, bool) or not 1 <= self.max_calls <=
1_000_000:
+ raise ToolProviderError(
+ "Custom tool rate limit max_calls must be in the range
1-1000000"
+ )
+ if (
+ isinstance(self.period_seconds, bool)
+ or not 0.1 <= self.period_seconds <= 3600
+ ):
+ raise ToolProviderError(
+ "Custom tool rate limit period_seconds must be in the range
0.1-3600"
+ )
+ if self.scope not in {"tool", "principal"}:
+ raise ToolProviderError(
+ "Custom tool rate limit scope must be tool or principal"
+ )
+
+
+@dataclass(frozen=True)
+class CustomTool:
+ """One schema, handler, and bounded runtime policy supplied by a
provider."""
+
+ tool: Tool
+ handler: ToolHandler
+ risk: ToolRisk = "high"
+ rate_limit: ToolRateLimit | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.tool, Tool):
+ raise ToolProviderError("Custom tool schema must be an MCP Tool")
+ if not callable(self.handler):
+ raise ToolProviderError(
+ f"Custom tool {self.tool.name!r} handler must be callable"
+ )
+ if self.risk not in {"low", "medium", "high"}:
+ raise ToolProviderError(
+ f"Custom tool {self.tool.name!r} risk must be low, medium, or
high"
+ )
+
+
+class CustomToolProvider(Protocol):
+ """Provider contract returned by an installed Python entry point."""
+
+ name: str
+
+ def tools(self) -> Iterable[CustomTool]:
+ """Return immutable tool definitions without performing network I/O."""
+
+
+@dataclass(frozen=True)
+class LoadedToolProvider:
+ """Validated provider and its materialized custom tool definitions."""
+
+ name: str
+ provider: CustomToolProvider
+ tools: tuple[CustomTool, ...]
+
+
+def normalize_tool_provider_names(raw_names: Iterable[str]) -> tuple[str, ...]:
+ """Validate and deduplicate an explicit provider allowlist."""
+ normalized: list[str] = []
+ seen: set[str] = set()
+ for raw_name in raw_names:
+ if not isinstance(raw_name, str):
+ raise ToolProviderError("Custom tool provider names must be
strings")
+ name = raw_name.strip()
+ if not _PROVIDER_NAME.fullmatch(name):
+ raise ToolProviderError(
+ "Custom tool provider names must contain only letters, digits,
"
+ "dot, underscore, or hyphen and be at most "
+ f"{MAX_PROVIDER_NAME_LENGTH} characters"
+ )
+ if name in seen:
+ raise ToolProviderError(
+ f"Duplicate custom tool provider in allowlist: {name}"
+ )
+ normalized.append(name)
+ seen.add(name)
+ return tuple(normalized)
+
+
+def prepare_tool_provider(
+ provider: CustomToolProvider,
+ *,
+ expected_name: str | None = None,
+) -> LoadedToolProvider:
+ """Validate one trusted provider before it reaches the tool registry."""
+ provider_name = getattr(provider, "name", None)
+ if not isinstance(provider_name, str):
+ raise ToolProviderError("Custom tool provider must define a string
name")
+ name = normalize_tool_provider_names([provider_name])[0]
+ if expected_name is not None and name != expected_name:
+ raise ToolProviderError(
+ f"Custom tool provider {expected_name!r} returned mismatched name
{name!r}"
+ )
+
+ tools_factory = getattr(provider, "tools", None)
+ if not callable(tools_factory):
+ raise ToolProviderError(
+ f"Custom tool provider {name!r} must define a tools() method"
+ )
+ try:
+ tools = tuple(tools_factory())
+ except Exception as exc:
+ raise ToolProviderError(
+ f"Custom tool provider {name!r} failed to define tools"
+ ) from exc
+ if not tools:
+ raise ToolProviderError(
+ f"Custom tool provider {name!r} must define at least one tool"
+ )
+ for tool in tools:
+ if not isinstance(tool, CustomTool):
+ raise ToolProviderError(
+ f"Custom tool provider {name!r} returned an invalid tool
definition"
+ )
+ return LoadedToolProvider(name=name, provider=provider, tools=tools)
+
+
+def load_tool_providers(
+ provider_names: Sequence[str],
+) -> tuple[LoadedToolProvider, ...]:
+ """Load only explicitly enabled installed entry points."""
+ names = normalize_tool_provider_names(provider_names)
+ if not names:
+ return ()
+
+ installed: dict[str, metadata.EntryPoint] = {}
+ for entry_point in metadata.entry_points(
+ group=CUSTOM_TOOL_PROVIDER_ENTRY_POINT
+ ):
+ if entry_point.name in installed:
+ raise ToolProviderError(
+ f"Duplicate installed custom tool provider: {entry_point.name}"
+ )
+ installed[entry_point.name] = entry_point
+
+ loaded: list[LoadedToolProvider] = []
+ for name in names:
+ selected_entry_point = installed.get(name)
+ if selected_entry_point is None:
+ raise ToolProviderError(
+ f"Configured custom tool provider is not installed: {name}"
+ )
+ try:
+ factory = selected_entry_point.load()
+ if not callable(factory):
+ raise TypeError("entry point target is not callable")
+ provider = factory()
+ except Exception as exc:
+ raise ToolProviderError(
+ f"Unable to load custom tool provider: {name}"
+ ) from exc
+ loaded.append(prepare_tool_provider(provider, expected_name=name))
+ return tuple(loaded)
+
+
+async def call_provider_lifecycle_hook(
+ loaded_provider: LoadedToolProvider,
+ hook_name: Literal["start", "close"],
+) -> None:
+ """Call one optional synchronous or asynchronous provider lifecycle
hook."""
+ hook = getattr(loaded_provider.provider, hook_name, None)
+ if hook is None:
+ return
+ if not callable(hook):
+ raise ToolProviderError(
+ f"Custom tool provider {loaded_provider.name!r} "
+ f"{hook_name} attribute must be callable"
+ )
+ result = hook()
+ if inspect.isawaitable(result):
+ await result
+
+
+@dataclass
+class _RateWindow:
+ calls: int
+ expires_at: float
+
+
+class LocalToolRateLimiter:
+ """Bounded process-local fixed-window limiter for custom tools."""
+
+ def __init__(
+ self,
+ *,
+ clock: Callable[[], float] = time.monotonic,
+ max_keys: int = MAX_RATE_LIMIT_KEYS,
+ ) -> None:
+ self._clock = clock
+ self._max_keys = max_keys
+ self._windows: dict[tuple[str, str], _RateWindow] = {}
+ self._lock = asyncio.Lock()
+
+ async def retry_after(
+ self,
+ *,
+ tool_name: str,
+ principal: str,
+ limit: ToolRateLimit,
+ ) -> float | None:
+ """Record an allowed call or return the retry delay for a rejected
one."""
+ key_principal = principal if limit.scope == "principal" else "*"
+ key = (tool_name, key_principal)
+ now = self._clock()
+
+ async with self._lock:
+ window = self._windows.get(key)
+ if window is not None and window.expires_at <= now:
+ self._windows.pop(key, None)
+ window = None
+
+ if window is None:
+ if len(self._windows) >= self._max_keys:
+ self._discard_expired_keys(now)
+ if len(self._windows) >= self._max_keys:
+ return limit.period_seconds
+ self._windows[key] = _RateWindow(
+ calls=1,
+ expires_at=now + limit.period_seconds,
+ )
+ return None
+
+ if window.calls >= limit.max_calls:
+ return max(0.0, window.expires_at - now)
+ window.calls += 1
+ return None
+
+ def _discard_expired_keys(self, now: float) -> None:
+ empty_keys = [
+ key
+ for key, window in self._windows.items()
+ if window.expires_at <= now
+ ]
+ for key in empty_keys:
+ self._windows.pop(key, None)
+
+
+class ToolProviderRuntime:
+ """Own provider discovery, lifecycle, and process-local rate-limit
state."""
+
+ def __init__(self, providers: Sequence[LoadedToolProvider]) -> None:
+ self.providers = tuple(providers)
+ self._started: list[LoadedToolProvider] = []
+ self._rate_limiter = LocalToolRateLimiter()
+
+ @classmethod
+ def create(
+ cls,
+ config: Any,
+ providers: Iterable[CustomToolProvider] | None,
+ ) -> ToolProviderRuntime:
+ """Build a runtime from configuration or explicitly injected
providers."""
+ if providers is None:
+ configured_names = getattr(config, "mcp_tool_providers", ())
+ if not isinstance(configured_names, list | tuple):
+ configured_names = ()
+ loaded = load_tool_providers(configured_names)
+ else:
+ loaded = tuple(prepare_tool_provider(provider) for provider in
providers)
+ return cls(loaded)
+
+ @property
+ def provider_count(self) -> int:
+ return len(self.providers)
+
+ def custom_tools(self) -> Iterable[tuple[str, CustomTool]]:
+ for loaded_provider in self.providers:
+ for custom_tool in loaded_provider.tools:
+ yield loaded_provider.name, custom_tool
+
+ async def start(self) -> None:
+ """Start providers in order and roll back completed starts on
failure."""
+ try:
+ for loaded_provider in self.providers:
+ await call_provider_lifecycle_hook(loaded_provider, "start")
+ self._started.append(loaded_provider)
+ except Exception:
+ await self.close()
+ raise
+
+ async def close(self) -> None:
+ """Close every started provider in reverse order."""
+ for loaded_provider in reversed(self._started):
+ try:
+ await call_provider_lifecycle_hook(loaded_provider, "close")
+ except Exception as exc:
+ logger.error(
+ "Custom tool provider shutdown failed (%s)",
+ type(exc).__name__,
+ )
+ self._started.clear()
+
+ async def retry_after(
+ self,
+ *,
+ tool_name: str,
+ auth_context: Any | None,
+ limit: ToolRateLimit,
+ ) -> float | None:
+ """Apply one tool limit without retaining bearer credentials."""
+ principal = "anonymous"
+ if auth_context is not None:
+ principal = (
+ getattr(auth_context, "user_id", None)
+ or getattr(auth_context, "token_id", None)
+ or getattr(auth_context, "oauth_client_id", None)
+ or getattr(auth_context, "auth_method", None)
+ or principal
+ )
+ return await self._rate_limiter.retry_after(
+ tool_name=tool_name,
+ principal=principal,
+ limit=limit,
+ )
diff --git a/doris_mcp_server/tools/tool_registry.py
b/doris_mcp_server/tools/tool_registry.py
index 2f742ef..a8f460d 100644
--- a/doris_mcp_server/tools/tool_registry.py
+++ b/doris_mcp_server/tools/tool_registry.py
@@ -19,12 +19,14 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Iterable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Any, Literal, Protocol, cast
from mcp.types import Tool
+from .tool_provider import CustomTool, ToolRateLimit
+
ToolPolicyClass = Literal["metadata", "query", "explain", "restricted"]
ToolHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
@@ -127,8 +129,17 @@ class ToolDefinition:
audit: ToolAuditDefinition
advertised: bool = True
argument_overrides: tuple[tuple[str, Any], ...] = ()
+ provider_name: str | None = None
+ direct_handler: ToolHandler | None = field(
+ default=None,
+ repr=False,
+ compare=False,
+ )
+ rate_limit: ToolRateLimit | None = None
def bind_handler(self, owner: ToolHandlerOwner) -> ToolHandler:
+ if self.direct_handler is not None:
+ return self.direct_handler
handler = getattr(owner, self.handler_name, None)
if not callable(handler):
raise ToolRegistryError(
@@ -229,6 +240,7 @@ class ToolDefinitionRegistry:
cls,
tools: Iterable[Tool],
owner: ToolHandlerOwner,
+ custom_tools: Iterable[tuple[str, CustomTool]] = (),
) -> ToolDefinitionRegistry:
tools_by_name: dict[str, Tool] = {}
definitions: list[ToolDefinition] = []
@@ -272,6 +284,28 @@ class ToolDefinitionRegistry:
)
)
+ for provider_name, custom_tool in custom_tools:
+ tool = custom_tool.tool
+ policy = ToolPolicyDefinition(
+ "restricted",
+ "custom_provider",
+ custom_tool.risk,
+ "UNSUPPORTED_FOR_OAUTH",
+ )
+ definitions.append(
+ ToolDefinition(
+ name=tool.name,
+ canonical_name=tool.name,
+ tool=tool,
+ handler_name=f"provider:{provider_name}",
+ policy=policy,
+ audit=_audit_for_tool(tool.name, tool, policy),
+ provider_name=provider_name,
+ direct_handler=custom_tool.handler,
+ rate_limit=custom_tool.rate_limit,
+ )
+ )
+
return cls(definitions, owner)
@property
diff --git a/doris_mcp_server/tools/tools_manager.py
b/doris_mcp_server/tools/tools_manager.py
index 0242197..6f054c1 100644
--- a/doris_mcp_server/tools/tools_manager.py
+++ b/doris_mcp_server/tools/tools_manager.py
@@ -20,7 +20,9 @@ Responsible for tool registration, management, scheduling and
routing, does not
"""
import json
+import math
import time
+from collections.abc import Iterable
from datetime import datetime
from typing import Any
@@ -47,6 +49,7 @@ from ..utils.schema_extractor import MetadataExtractor
from ..utils.security import get_current_auth_context
from ..utils.security_analytics_tools import SecurityAnalyticsTools
from .tool_catalog import build_tool_registry
+from .tool_provider import CustomToolProvider, ToolProviderRuntime
from .tool_registry import ToolDefinition, ToolDefinitionRegistry
logger = get_logger(__name__)
@@ -55,10 +58,15 @@ logger = get_logger(__name__)
class DorisToolsManager:
"""Apache Doris Tools Manager"""
- def __init__(self, connection_manager: DorisConnectionManager) -> None:
+ def __init__(
+ self,
+ connection_manager: DorisConnectionManager,
+ *,
+ tool_providers: Iterable[CustomToolProvider] | None = None,
+ ) -> None:
self.connection_manager = connection_manager
-
- # Initialize business logic processors
+ config = getattr(connection_manager, "config", None)
+ self._tool_provider_runtime = ToolProviderRuntime.create(config,
tool_providers)
self.query_executor = DorisQueryExecutor(connection_manager)
self.table_analyzer = TableAnalyzer(connection_manager)
self.sql_analyzer = SQLAnalyzer(connection_manager)
@@ -83,15 +91,23 @@ class DorisToolsManager:
self._tool_registry = self._build_tool_registry()
logger.info(
- "DorisToolsManager initialized with business logic processors,
v0.5.0 analytics tools, and ADBC query tools"
+ "DorisToolsManager initialized with business logic processors,
v0.5.0 "
+ "analytics tools, ADBC query tools, and %d custom tool providers",
+ self._tool_provider_runtime.provider_count,
)
async def start(self) -> None:
"""Start runtime resources owned by the tools manager."""
await self.query_executor.start()
+ try:
+ await self._tool_provider_runtime.start()
+ except Exception:
+ await self.query_executor.close()
+ raise
async def close(self) -> None:
"""Stop runtime resources owned by the tools manager."""
+ await self._tool_provider_runtime.close()
await self.query_executor.close()
@staticmethod
@@ -117,6 +133,8 @@ class DorisToolsManager:
def _build_tool_registry(self) -> ToolDefinitionRegistry:
"""Build the registry from the standalone immutable catalog."""
+ provider_runtime = getattr(self, "_tool_provider_runtime", None)
+ custom_tools = provider_runtime.custom_tools() if provider_runtime
else ()
return build_tool_registry(
self,
getattr(
@@ -124,6 +142,7 @@ class DorisToolsManager:
"config",
None,
),
+ custom_tools=custom_tools,
)
@property
@@ -151,13 +170,40 @@ class DorisToolsManager:
start_time = time.time()
try:
definition = self.tool_registry.resolve(name)
+ if definition.rate_limit is not None:
+ retry_after = await self._tool_provider_runtime.retry_after(
+ tool_name=definition.name,
+ auth_context=get_current_auth_context(),
+ limit=definition.rate_limit,
+ )
+ if retry_after is not None:
+ execution_time = time.time() - start_time
+ self._audit_tool_call(
+ definition,
+ arguments=arguments,
+ status="rate_limited",
+ execution_time=execution_time,
+ )
+ return json.dumps(
+ {
+ "error": "Tool rate limit exceeded",
+ "error_code": "TOOL_RATE_LIMITED",
+ "retry_after_seconds": max(
+ 1,
+ math.ceil(retry_after),
+ ),
+ "timestamp": datetime.now().isoformat(),
+ },
+ ensure_ascii=False,
+ indent=2,
+ )
handler = definition.bind_handler(self)
prepared_arguments = definition.prepare_arguments(arguments)
result = await handler(prepared_arguments)
execution_time = time.time() - start_time
# Add execution information
- if isinstance(result, dict):
+ if isinstance(result, dict) and definition.provider_name is None:
result["_execution_info"] = {
"tool_name": name,
"canonical_tool_name": definition.canonical_name,
diff --git a/doris_mcp_server/utils/config.py b/doris_mcp_server/utils/config.py
index ff68dc3..42f237d 100644
--- a/doris_mcp_server/utils/config.py
+++ b/doris_mcp_server/utils/config.py
@@ -43,6 +43,10 @@ from ..result_limits import (
DEFAULT_RESULT_ROWS,
MIN_RESULT_BYTES,
)
+from ..tools.tool_provider import (
+ ToolProviderError,
+ normalize_tool_provider_names,
+)
from ..tools.tool_registry import (
DORIS_OAUTH_EXPLAIN_TOOL_SET,
DORIS_OAUTH_METADATA_TOOL_NAMES,
@@ -841,6 +845,7 @@ class DorisConfig:
mcp_allowed_origins: list[str] = field(default_factory=list)
enable_legacy_http_adapter: bool = False
mcp_list_page_size: int = 100
+ mcp_tool_providers: list[str] = field(default_factory=list)
mcp_state_handle_secret: str = field(
default_factory=lambda: secrets.token_urlsafe(32),
repr=False,
@@ -1461,6 +1466,13 @@ class DorisConfig:
config.mcp_list_page_size,
)
_mark_source(config, "mcp_list_page_size", "env")
+ if "MCP_TOOL_PROVIDERS" in os.environ:
+ config.mcp_tool_providers = [
+ value.strip()
+ for value in os.getenv("MCP_TOOL_PROVIDERS", "").split(",")
+ if value.strip()
+ ]
+ _mark_source(config, "mcp_tool_providers", "env")
if "MCP_STATE_HANDLE_SECRET" in os.environ:
config.mcp_state_handle_secret = os.getenv(
"MCP_STATE_HANDLE_SECRET",
@@ -1491,6 +1503,7 @@ class DorisConfig:
"mcp_allowed_origins",
"enable_legacy_http_adapter",
"mcp_list_page_size",
+ "mcp_tool_providers",
"mcp_state_handle_ttl_seconds",
"temp_files_dir",
"transport",
@@ -1567,6 +1580,7 @@ class DorisConfig:
"mcp_allowed_origins": self.mcp_allowed_origins,
"enable_legacy_http_adapter": self.enable_legacy_http_adapter,
"mcp_list_page_size": self.mcp_list_page_size,
+ "mcp_tool_providers": self.mcp_tool_providers,
"mcp_state_handle_ttl_seconds": self.mcp_state_handle_ttl_seconds,
"temp_files_dir": self.temp_files_dir,
"database": {
@@ -1774,6 +1788,15 @@ class DorisConfig:
if not 1 <= self.mcp_list_page_size <= 1000:
errors.append("MCP list page size must be in the range 1-1000")
+ raw_tool_providers: Any = self.mcp_tool_providers
+ if not isinstance(raw_tool_providers, list):
+ errors.append("MCP tool providers must be a list")
+ else:
+ try:
+ normalize_tool_provider_names(raw_tool_providers)
+ except ToolProviderError as exc:
+ errors.append(str(exc))
+
if len(self.mcp_state_handle_secret.encode("utf-8")) < 32:
errors.append("MCP state handle secret must contain at least 32
bytes")
diff --git a/test/protocol/test_multiworker_config.py
b/test/protocol/test_multiworker_config.py
index 9d05131..25996b2 100644
--- a/test/protocol/test_multiworker_config.py
+++ b/test/protocol/test_multiworker_config.py
@@ -43,6 +43,33 @@ def
test_mcp_list_page_size_is_configurable_and_bounded(monkeypatch):
assert "MCP list page size must be in the range 1-1000" in
configured.validate()
+def test_custom_tool_provider_allowlist_is_explicit_and_validated(monkeypatch):
+ monkeypatch.delenv("MCP_TOOL_PROVIDERS", raising=False)
+ assert DorisConfig.from_env().mcp_tool_providers == []
+
+ monkeypatch.setenv("MCP_TOOL_PROVIDERS", "orders_api, customer-tools")
+ configured = DorisConfig.from_env()
+ assert configured.mcp_tool_providers == ["orders_api", "customer-tools"]
+ assert configured.to_dict()["mcp_tool_providers"] == [
+ "orders_api",
+ "customer-tools",
+ ]
+ assert configured.validate() == []
+
+ configured.mcp_tool_providers = ["orders_api", "orders_api"]
+ assert (
+ "Duplicate custom tool provider in allowlist: orders_api"
+ in configured.validate()
+ )
+ configured.mcp_tool_providers = ["bad/provider"]
+ assert any(
+ error.startswith("Custom tool provider names must contain")
+ for error in configured.validate()
+ )
+ configured.mcp_tool_providers = "orders_api" # type: ignore[assignment]
+ assert "MCP tool providers must be a list" in configured.validate()
+
+
def
test_state_handle_secret_and_ttl_are_configurable_without_serializing_secret(
monkeypatch,
):
@@ -82,6 +109,7 @@ def
test_multiworker_environment_preserves_resolved_parent_config(monkeypatch):
config.mcp_allowed_origins = ["https://client.example.test";]
config.enable_legacy_http_adapter = True
config.mcp_list_page_size = 17
+ config.mcp_tool_providers = ["orders_api", "customer-tools"]
config.mcp_state_handle_secret = "parent-shared-state-handle-secret-value"
config.mcp_state_handle_ttl_seconds = 45
@@ -111,6 +139,7 @@ def
test_multiworker_environment_preserves_resolved_parent_config(monkeypatch):
assert child_config.mcp_allowed_origins == ["https://client.example.test";]
assert child_config.enable_legacy_http_adapter is True
assert child_config.mcp_list_page_size == 17
+ assert child_config.mcp_tool_providers == ["orders_api", "customer-tools"]
assert (
child_config.mcp_state_handle_secret
== "parent-shared-state-handle-secret-value"
diff --git a/test/security/test_operation_policy.py
b/test/security/test_operation_policy.py
index 6e9101d..5d99366 100644
--- a/test/security/test_operation_policy.py
+++ b/test/security/test_operation_policy.py
@@ -305,3 +305,34 @@ def
test_filter_tools_for_doris_oauth_hides_disabled_or_denied_tools():
)
assert [tool.name for tool in filtered] == ["get_db_list", "exec_query"]
+
+
+def test_external_oauth_hides_custom_tools_without_reviewed_policy():
+ context = AuthContext(
+ auth_method="external_oauth",
+ oauth_scopes=[
+ "tool:list",
+ "tool:call:get_db_list",
+ "tool:call:custom_business_api",
+ ],
+ )
+ tools = [
+ SimpleNamespace(name="get_db_list"),
+ SimpleNamespace(name="custom_business_api"),
+ ]
+
+ filtered = filter_tools_for_auth_context(context, tools)
+
+ assert [tool.name for tool in filtered] == ["get_db_list"]
+
+
+def test_external_oauth_rejects_custom_tool_call_without_reviewed_policy():
+ context = AuthContext(
+ auth_method="external_oauth",
+ oauth_scopes=["tool:call:custom_business_api"],
+ )
+
+ with pytest.raises(OperationAuthorizationError) as exc:
+ authorize_operation(context, "tool:custom_business_api")
+
+ assert exc.value.error_code == "UNKNOWN_OPERATION"
diff --git a/test/test_product_identity.py b/test/test_product_identity.py
index e234ca6..4ca5b2f 100644
--- a/test/test_product_identity.py
+++ b/test/test_product_identity.py
@@ -108,3 +108,19 @@ def
test_fine_grained_access_control_guide_is_linked_and_complete():
assert "root` and `admin" in guide
assert "Doris-backed OAuth" in guide
assert "Token-bound database configuration" in guide
+
+
+def test_custom_tool_provider_guide_is_linked_and_complete():
+ guide_path = PROJECT_ROOT / "docs" / "custom-tool-providers.md"
+ guide = guide_path.read_text(encoding="utf-8")
+ readme = (PROJECT_ROOT / "README.md").read_text(encoding="utf-8")
+ chinese_readme = (PROJECT_ROOT /
"README.zh-CN.md").read_text(encoding="utf-8")
+ link = "docs/custom-tool-providers.md"
+
+ assert f"]({link})" in readme
+ assert f"]({link})" in chinese_readme
+ assert "doris_mcp_server.tool_providers" in guide
+ assert "MCP_TOOL_PROVIDERS" in guide
+ assert "ToolRateLimit" in guide
+ assert "fail closed" in guide
+ assert "FastGPT" in guide
diff --git a/test/tools/test_custom_tool_provider.py
b/test/tools/test_custom_tool_provider.py
new file mode 100644
index 0000000..3ffd7c2
--- /dev/null
+++ b/test/tools/test_custom_tool_provider.py
@@ -0,0 +1,456 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for the explicit trusted custom-tool provider boundary."""
+
+from __future__ import annotations
+
+import json
+import logging
+from importlib import metadata
+from unittest.mock import AsyncMock, Mock
+
+import httpx2
+import pytest
+from mcp.types import Tool
+
+from doris_mcp_server import __version__
+from doris_mcp_server.protocol import (
+ create_doris_mcp_server,
+ create_transport_security,
+)
+from doris_mcp_server.tools.tool_provider import (
+ CustomTool,
+ LocalToolRateLimiter,
+ ToolProviderError,
+ ToolRateLimit,
+ load_tool_providers,
+)
+from doris_mcp_server.tools.tool_registry import ToolRegistryError
+from doris_mcp_server.tools.tools_manager import DorisToolsManager
+from doris_mcp_server.utils.config import DorisConfig
+from doris_mcp_server.utils.security import (
+ AuthContext,
+ reset_auth_context,
+ set_current_auth_context,
+)
+
+
+def _connection_manager() -> Mock:
+ manager = Mock()
+ manager.config = DorisConfig()
+ manager.get_connection = AsyncMock()
+ return manager
+
+
+class RecordingProvider:
+ name = "orders_api"
+
+ def __init__(
+ self,
+ *,
+ tool_name: str = "lookup_business_order",
+ rate_limit: ToolRateLimit | None = None,
+ ) -> None:
+ self.handler = AsyncMock(
+ return_value={"ok": True, "source": "business-api"}
+ )
+ self.start = AsyncMock()
+ self.close = AsyncMock()
+ self._tool_name = tool_name
+ self._rate_limit = rate_limit
+
+ def tools(self) -> tuple[CustomTool, ...]:
+ return (
+ CustomTool(
+ tool=Tool(
+ name=self._tool_name,
+ description="Look up one order through a trusted business
API",
+ input_schema={
+ "type": "object",
+ "properties": {
+ "order_id": {"type": "string"},
+ },
+ "required": ["order_id"],
+ "additionalProperties": False,
+ },
+ output_schema={
+ "type": "object",
+ "properties": {
+ "ok": {"type": "boolean"},
+ "source": {"type": "string"},
+ },
+ "required": ["ok", "source"],
+ "additionalProperties": False,
+ },
+ ),
+ handler=self.handler,
+ rate_limit=self._rate_limit,
+ ),
+ )
+
+
[email protected]
+async def
test_provider_tool_is_listed_dispatched_audited_and_lifecycle_managed():
+ provider = RecordingProvider()
+ manager = DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[provider],
+ )
+ manager.query_executor.start = AsyncMock()
+ manager.query_executor.close = AsyncMock()
+
+ await manager.start()
+ listed = await manager.list_tools()
+ payload = json.loads(
+ await manager.call_tool(
+ "lookup_business_order",
+ {"order_id": "order-42"},
+ )
+ )
+ await manager.close()
+
+ assert "lookup_business_order" in {tool.name for tool in listed}
+ provider.handler.assert_awaited_once_with({"order_id": "order-42"})
+ assert payload["ok"] is True
+ assert payload["source"] == "business-api"
+ assert "_execution_info" not in payload
+ definition = manager.tool_registry.resolve("lookup_business_order")
+ assert definition.provider_name == "orders_api"
+ assert definition.policy.channel == "custom_provider"
+ provider.start.assert_awaited_once_with()
+ provider.close.assert_awaited_once_with()
+
+
[email protected]
+async def test_provider_start_failure_rolls_back_started_providers():
+ first = RecordingProvider()
+ second = RecordingProvider(tool_name="lookup_customer")
+ second.name = "customer_api"
+ second.start.side_effect = RuntimeError("upstream unavailable")
+ manager = DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[first, second],
+ )
+ manager.query_executor.start = AsyncMock()
+ manager.query_executor.close = AsyncMock()
+
+ with pytest.raises(RuntimeError, match="upstream unavailable"):
+ await manager.start()
+
+ first.start.assert_awaited_once_with()
+ first.close.assert_awaited_once_with()
+ second.start.assert_awaited_once_with()
+ second.close.assert_not_awaited()
+ manager.query_executor.close.assert_awaited_once_with()
+
+
[email protected]
+async def test_custom_tool_rate_limit_is_fail_closed_and_principal_scoped():
+ provider = RecordingProvider(
+ rate_limit=ToolRateLimit(
+ max_calls=1,
+ period_seconds=60,
+ scope="principal",
+ )
+ )
+ manager = DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[provider],
+ )
+
+ first = json.loads(
+ await manager.call_tool(
+ "lookup_business_order",
+ {"order_id": "first"},
+ )
+ )
+ limited = json.loads(
+ await manager.call_tool(
+ "lookup_business_order",
+ {"order_id": "second"},
+ )
+ )
+
+ context_token = set_current_auth_context(
+ AuthContext(
+ auth_method="token",
+ token_id="another-principal",
+ user_id="another-principal",
+ )
+ )
+ try:
+ another_principal = json.loads(
+ await manager.call_tool(
+ "lookup_business_order",
+ {"order_id": "third"},
+ )
+ )
+ finally:
+ reset_auth_context(context_token)
+
+ assert first["ok"] is True
+ assert limited["error_code"] == "TOOL_RATE_LIMITED"
+ assert limited["retry_after_seconds"] >= 1
+ assert another_principal["ok"] is True
+ assert provider.handler.await_count == 2
+
+
+def test_custom_tool_cannot_shadow_a_builtin_tool():
+ provider = RecordingProvider(tool_name="exec_query")
+
+ with pytest.raises(ToolRegistryError, match="Duplicate tool definition"):
+ DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[provider],
+ )
+
+
+def test_custom_tools_cannot_shadow_each_other():
+ first = RecordingProvider()
+ second = RecordingProvider()
+ second.name = "customer_api"
+
+ with pytest.raises(ToolRegistryError, match="Duplicate tool definition"):
+ DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[first, second],
+ )
+
+
+class FakeEntryPoint:
+ def __init__(self, name: str, factory: object) -> None:
+ self.name = name
+ self._factory = factory
+
+ def load(self) -> object:
+ return self._factory
+
+
+def test_loader_loads_only_explicit_installed_entry_points(monkeypatch):
+ provider = RecordingProvider()
+ unselected_factory = Mock(side_effect=AssertionError("must not be loaded"))
+ monkeypatch.setattr(
+ metadata,
+ "entry_points",
+ lambda **kwargs: [
+ FakeEntryPoint("orders_api", lambda: provider),
+ FakeEntryPoint("unselected", unselected_factory),
+ ],
+ )
+
+ loaded = load_tool_providers(["orders_api"])
+
+ assert [item.name for item in loaded] == ["orders_api"]
+ assert loaded[0].provider is provider
+ assert loaded[0].tools[0].tool.name == "lookup_business_order"
+ unselected_factory.assert_not_called()
+
+
+def test_loader_rejects_unknown_or_mismatched_providers(monkeypatch):
+ monkeypatch.setattr(metadata, "entry_points", lambda **kwargs: [])
+ with pytest.raises(ToolProviderError, match="is not installed"):
+ load_tool_providers(["missing"])
+
+ mismatched = RecordingProvider()
+ mismatched.name = "different_name"
+ monkeypatch.setattr(
+ metadata,
+ "entry_points",
+ lambda **kwargs: [
+ FakeEntryPoint("orders_api", lambda: mismatched),
+ ],
+ )
+ with pytest.raises(ToolProviderError, match="mismatched name"):
+ load_tool_providers(["orders_api"])
+
+
[email protected](
+ ("max_calls", "period_seconds", "scope"),
+ [
+ (0, 1, "principal"),
+ (1, 0, "principal"),
+ (1, 1, "unknown"),
+ ],
+)
+def test_custom_tool_rate_limit_rejects_invalid_bounds(
+ max_calls,
+ period_seconds,
+ scope,
+):
+ with pytest.raises(ToolProviderError):
+ ToolRateLimit(
+ max_calls=max_calls,
+ period_seconds=period_seconds,
+ scope=scope,
+ )
+
+
[email protected]
+async def test_rate_limiter_recycles_expired_principal_keys():
+ now = 0.0
+ limiter = LocalToolRateLimiter(clock=lambda: now, max_keys=1)
+ limit = ToolRateLimit(max_calls=1, period_seconds=1)
+
+ assert (
+ await limiter.retry_after(
+ tool_name="lookup",
+ principal="first",
+ limit=limit,
+ )
+ is None
+ )
+ assert (
+ await limiter.retry_after(
+ tool_name="lookup",
+ principal="second",
+ limit=limit,
+ )
+ == 1
+ )
+
+ now = 1.1
+ assert (
+ await limiter.retry_after(
+ tool_name="lookup",
+ principal="second",
+ limit=limit,
+ )
+ is None
+ )
+
+
+def _modern_request(
+ request_id: int,
+ method: str,
+ *,
+ name: str | None = None,
+ arguments: dict | None = None,
+) -> dict:
+ params: dict = {
+ "_meta": {
+ "io.modelcontextprotocol/protocolVersion": "2026-07-28",
+ "io.modelcontextprotocol/clientCapabilities": {},
+ "io.modelcontextprotocol/clientInfo": {
+ "name": "custom-provider-test",
+ "version": "1.0.0",
+ },
+ }
+ }
+ if name is not None:
+ params.update({"name": name, "arguments": arguments or {}})
+ return {
+ "jsonrpc": "2.0",
+ "id": request_id,
+ "method": method,
+ "params": params,
+ }
+
+
+def _modern_headers(method: str, *, name: str | None = None) -> dict[str, str]:
+ headers = {
+ "Accept": "application/json, text/event-stream",
+ "Content-Type": "application/json",
+ "Mcp-Protocol-Version": "2026-07-28",
+ "Mcp-Method": method,
+ }
+ if name is not None:
+ headers["Mcp-Name"] = name
+ return headers
+
+
[email protected]
+async def
test_custom_tool_uses_real_streamable_http_list_validation_and_call():
+ provider = RecordingProvider()
+ manager = DorisToolsManager(
+ _connection_manager(),
+ tool_providers=[provider],
+ )
+ resources_manager = Mock()
+ resources_manager.list_resources = AsyncMock(return_value=[])
+ resources_manager.read_resource = AsyncMock()
+ prompts_manager = Mock()
+ prompts_manager.list_prompts = AsyncMock(return_value=[])
+ prompts_manager.get_prompt = AsyncMock()
+ server = create_doris_mcp_server(
+ resources_manager=resources_manager,
+ tools_manager=manager,
+ prompts_manager=prompts_manager,
+ name="doris-mcp-custom-provider-test",
+ version=__version__,
+ logger=logging.getLogger(__name__),
+ )
+ app = server.streamable_http_app(
+ json_response=True,
+ stateless_http=True,
+ host="127.0.0.1",
+ transport_security=create_transport_security("127.0.0.1"),
+ )
+
+ async with (
+ app.router.lifespan_context(app),
+ httpx2.ASGITransport(app) as transport,
+ httpx2.AsyncClient(
+ transport=transport,
+ base_url="http://127.0.0.1:3000";,
+ ) as client,
+ ):
+ listed = await client.post(
+ "/mcp",
+ json=_modern_request(1, "tools/list"),
+ headers=_modern_headers("tools/list"),
+ )
+ invalid = await client.post(
+ "/mcp",
+ json=_modern_request(
+ 2,
+ "tools/call",
+ name="lookup_business_order",
+ ),
+ headers=_modern_headers(
+ "tools/call",
+ name="lookup_business_order",
+ ),
+ )
+ called = await client.post(
+ "/mcp",
+ json=_modern_request(
+ 3,
+ "tools/call",
+ name="lookup_business_order",
+ arguments={"order_id": "order-42"},
+ ),
+ headers=_modern_headers(
+ "tools/call",
+ name="lookup_business_order",
+ ),
+ )
+
+ assert listed.status_code == 200
+ listed_tools = {
+ tool["name"]: tool for tool in listed.json()["result"]["tools"]
+ }
+ assert "lookup_business_order" in listed_tools
+ assert listed_tools["lookup_business_order"]["inputSchema"]["required"] ==
[
+ "order_id"
+ ]
+ assert invalid.status_code == 400
+ assert invalid.json()["error"]["code"] == -32602
+ assert called.status_code == 200
+ structured = called.json()["result"]["structuredContent"]
+ assert structured["ok"] is True
+ assert structured["source"] == "business-api"
+ assert "_execution_info" not in structured
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
