mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-28 12:16:36 +00:00
arc53-machine
473e8ec15c
- Editing a connection-backed tool validates the config before writing new credentials, so a rejected edit no longer rotates the connection's key and resumes its sources. - An MCP edit keeps its previous connection only when the caller owns it and it signs in the way the tool now does.
609 lines
25 KiB
Python
609 lines
25 KiB
Python
"""Tool management MCP server integration."""
|
|
|
|
from urllib.parse import urlencode, urlparse
|
|
|
|
from flask import current_app, jsonify, make_response, redirect, request
|
|
from flask_restx import Namespace, Resource, fields
|
|
|
|
from docsgpt.agents.tools.mcp_tool import MCPOAuthManager, MCPTool
|
|
from docsgpt.api import api
|
|
from docsgpt.api.user.tools.routes import transform_actions
|
|
from docsgpt.cache import get_redis_instance
|
|
from docsgpt.core.url_validation import SSRFError, validate_url
|
|
from docsgpt.security.encryption import decrypt_credentials, encrypt_credentials
|
|
from docsgpt.storage.db.repositories.connector_sessions import (
|
|
ConnectorSessionsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
from docsgpt.storage.db.session import db_readonly, db_session
|
|
from docsgpt.utils import check_required_fields
|
|
|
|
tools_mcp_ns = Namespace("tools", description="Tool management operations", path="/api")
|
|
|
|
_ALLOWED_TRANSPORTS = {"auto", "sse", "http"}
|
|
|
|
|
|
def _sanitize_mcp_transport(config):
|
|
"""Normalise and validate the transport_type field.
|
|
|
|
Strips ``command`` / ``args`` keys that are only valid for local STDIO
|
|
transports, and ``connection_id``, which only the tool executor sets
|
|
(it picks whose MCP tokens the tool uses). Returns the cleaned
|
|
transport type string.
|
|
"""
|
|
transport_type = (config.get("transport_type") or "auto").lower()
|
|
if transport_type not in _ALLOWED_TRANSPORTS:
|
|
raise ValueError(f"Unsupported transport_type: {transport_type}")
|
|
config.pop("command", None)
|
|
config.pop("args", None)
|
|
config.pop("connection_id", None)
|
|
config["transport_type"] = transport_type
|
|
return transport_type
|
|
|
|
|
|
def _extract_auth_credentials(config):
|
|
"""Build an ``auth_credentials`` dict from the raw MCP config."""
|
|
auth_credentials = {}
|
|
auth_type = config.get("auth_type", "none")
|
|
|
|
if auth_type == "api_key":
|
|
if config.get("api_key"):
|
|
auth_credentials["api_key"] = config["api_key"]
|
|
if config.get("api_key_header"):
|
|
auth_credentials["api_key_header"] = config["api_key_header"]
|
|
elif auth_type == "bearer":
|
|
if config.get("bearer_token"):
|
|
auth_credentials["bearer_token"] = config["bearer_token"]
|
|
elif auth_type == "basic":
|
|
if config.get("username"):
|
|
auth_credentials["username"] = config["username"]
|
|
if config.get("password"):
|
|
auth_credentials["password"] = config["password"]
|
|
|
|
return auth_credentials
|
|
|
|
|
|
def _validate_mcp_server_url(config: dict) -> None:
|
|
"""Validate the server_url in an MCP config to prevent SSRF.
|
|
|
|
Raises:
|
|
ValueError: If the URL is missing or points to a blocked address.
|
|
"""
|
|
server_url = (config.get("server_url") or "").strip()
|
|
if not server_url:
|
|
raise ValueError("server_url is required")
|
|
try:
|
|
validate_url(server_url)
|
|
except SSRFError as exc:
|
|
raise ValueError(f"Invalid server URL: {exc}") from exc
|
|
|
|
|
|
def _mcp_connection(user, config, auth_type, auth_credentials, display_name):
|
|
"""The connection an MCP tool runs with; created on first save.
|
|
|
|
OAuth servers already have one (the sign-in stored its tokens there).
|
|
Key, bearer and basic auth store their secret on a connection; servers
|
|
with no auth get a credential-less connection so they still appear on
|
|
the Connectors page. Returns None when a multi-user install runs on the
|
|
default encryption key, which keeps the legacy per-tool secret.
|
|
"""
|
|
from docsgpt.connectors import catalog, service
|
|
|
|
base = catalog.base_url(config.get("server_url"))
|
|
if not base:
|
|
return None
|
|
if auth_type == "oauth":
|
|
with db_readonly() as conn:
|
|
row = service._mcp_row(conn, user, base, None)
|
|
return str(row["id"]) if row else None
|
|
definition = catalog.get_definition("custom_mcp")
|
|
host = base.split("://")[-1]
|
|
try:
|
|
with db_session() as conn:
|
|
if auth_credentials:
|
|
row, _ = service.create_api_key_connection(
|
|
conn, user, definition, auth_credentials, server_url=base, display_name=display_name,
|
|
)
|
|
else:
|
|
repo = ConnectorSessionsRepository(conn)
|
|
row = repo.find_account(user, "custom_mcp", server_url=base, account_label=host) or repo.create(
|
|
user, "custom_mcp", connector_key="custom_mcp", auth_kind="none",
|
|
display_name=display_name, account_label=host, server_url=base,
|
|
)
|
|
except service.EncryptionKeyNotConfigured:
|
|
return None
|
|
return str(row["id"]) if row else None
|
|
|
|
|
|
def _previous_connection(existing_doc, config, user) -> str | None:
|
|
"""The saved tool's connection, kept only while the server is unchanged.
|
|
|
|
An edit that cannot resolve a connection of its own (the default
|
|
encryption key blocks storing a new secret) must not carry the old
|
|
one over to a different server (its key would be sent there), to a
|
|
connection the caller does not own, or to one that signs in another way.
|
|
"""
|
|
from docsgpt.connectors import catalog
|
|
|
|
connection_id = (existing_doc or {}).get("connection_id")
|
|
base = catalog.base_url(config.get("server_url"))
|
|
if not connection_id or not base:
|
|
return None
|
|
with db_readonly() as conn:
|
|
row = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user)
|
|
if row is None or catalog.base_url(row.get("server_url")) != base:
|
|
return None
|
|
wanted = {"oauth": "mcp_oauth", "none": "none"}.get(config.get("auth_type") or "none", "api_key")
|
|
if (row.get("auth_kind") or "") != wanted:
|
|
return None
|
|
return str(connection_id)
|
|
|
|
|
|
def _mcp_policy_error(config: dict):
|
|
"""A 403 when an admin turned this MCP server's connector off, else None.
|
|
|
|
A preset's own switch applies to its server; any other server is a
|
|
custom connector and needs "Allow custom MCP servers".
|
|
"""
|
|
from docsgpt.connectors import catalog, service
|
|
|
|
preset = catalog.preset_for_url(config.get("server_url"))
|
|
key = preset.key if preset else "custom_mcp"
|
|
try:
|
|
with db_readonly() as conn:
|
|
service.ensure_connector_allowed(conn, key)
|
|
except service.ConnectorDisabled:
|
|
return make_response(
|
|
jsonify({"success": False, "error": "This MCP server is turned off by an admin", "code": "disabled"}),
|
|
403,
|
|
)
|
|
except Exception:
|
|
# The switch is an admin preference; when the policy table cannot be
|
|
# read, saving the tool (which needs the database) fails on its own.
|
|
current_app.logger.warning("Could not read connector policies", exc_info=True)
|
|
return None
|
|
|
|
|
|
@tools_mcp_ns.route("/mcp_server/test")
|
|
class TestMCPServerConfig(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"MCPServerTestModel",
|
|
{
|
|
"config": fields.Raw(
|
|
required=True, description="MCP server configuration to test"
|
|
),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(description="Test MCP server connection with provided configuration")
|
|
def post(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
data = request.get_json()
|
|
|
|
required_fields = ["config"]
|
|
missing_fields = check_required_fields(data, required_fields)
|
|
if missing_fields:
|
|
return missing_fields
|
|
try:
|
|
config = data["config"]
|
|
try:
|
|
_sanitize_mcp_transport(config)
|
|
except ValueError:
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Unsupported transport_type"}),
|
|
400,
|
|
)
|
|
|
|
_validate_mcp_server_url(config)
|
|
policy_error = _mcp_policy_error(config)
|
|
if policy_error is not None:
|
|
return policy_error
|
|
|
|
auth_credentials = _extract_auth_credentials(config)
|
|
test_config = config.copy()
|
|
test_config["auth_credentials"] = auth_credentials
|
|
|
|
mcp_tool = MCPTool(config=test_config, user_id=user)
|
|
result = mcp_tool.test_connection()
|
|
|
|
if result.get("requires_oauth"):
|
|
safe_result = {
|
|
k: v
|
|
for k, v in result.items()
|
|
if k in ("success", "requires_oauth", "auth_url")
|
|
}
|
|
return make_response(jsonify(safe_result), 200)
|
|
|
|
if not result.get("success"):
|
|
current_app.logger.error(
|
|
f"MCP connection test failed: {result.get('message')}"
|
|
)
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "Connection test failed",
|
|
"tools_count": 0,
|
|
}
|
|
),
|
|
200,
|
|
)
|
|
|
|
safe_result = {
|
|
"success": True,
|
|
"message": result.get("message", "Connection successful"),
|
|
"tools_count": result.get("tools_count", 0),
|
|
"tools": result.get("tools", []),
|
|
}
|
|
return make_response(jsonify(safe_result), 200)
|
|
except ValueError as e:
|
|
current_app.logger.warning(f"Invalid MCP server test request: {e}")
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Invalid MCP server configuration"}),
|
|
400,
|
|
)
|
|
except Exception as e:
|
|
current_app.logger.error(f"Error testing MCP server: {e}", exc_info=True)
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Connection test failed"}),
|
|
500,
|
|
)
|
|
|
|
|
|
@tools_mcp_ns.route("/mcp_server/save")
|
|
class MCPServerSave(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"MCPServerSaveModel",
|
|
{
|
|
"id": fields.String(
|
|
required=False, description="Tool ID for updates (optional)"
|
|
),
|
|
"displayName": fields.String(
|
|
required=True, description="Display name for the MCP server"
|
|
),
|
|
"config": fields.Raw(
|
|
required=True, description="MCP server configuration"
|
|
),
|
|
"status": fields.Boolean(
|
|
required=False, default=True, description="Tool status"
|
|
),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(description="Create or update MCP server with automatic tool discovery")
|
|
def post(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
data = request.get_json()
|
|
|
|
required_fields = ["displayName", "config"]
|
|
missing_fields = check_required_fields(data, required_fields)
|
|
if missing_fields:
|
|
return missing_fields
|
|
try:
|
|
config = data["config"]
|
|
try:
|
|
_sanitize_mcp_transport(config)
|
|
except ValueError:
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Unsupported transport_type"}),
|
|
400,
|
|
)
|
|
|
|
_validate_mcp_server_url(config)
|
|
policy_error = _mcp_policy_error(config)
|
|
if policy_error is not None:
|
|
return policy_error
|
|
|
|
auth_credentials = _extract_auth_credentials(config)
|
|
auth_type = config.get("auth_type", "none")
|
|
mcp_config = config.copy()
|
|
mcp_config["auth_credentials"] = auth_credentials
|
|
|
|
if auth_type == "oauth":
|
|
if not config.get("oauth_task_id"):
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"error": "Connection not authorized. Please complete the OAuth authorization first.",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
redis_client = get_redis_instance()
|
|
manager = MCPOAuthManager(redis_client)
|
|
result = manager.get_oauth_status(
|
|
config["oauth_task_id"], user
|
|
)
|
|
if not result.get("status") == "completed":
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"error": "OAuth failed or not completed. Please try authorizing again.",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
actions_metadata = result.get("tools", [])
|
|
elif auth_type == "none" or auth_credentials:
|
|
mcp_tool = MCPTool(config=mcp_config, user_id=user)
|
|
mcp_tool.discover_tools()
|
|
actions_metadata = mcp_tool.get_actions_metadata()
|
|
else:
|
|
raise Exception(
|
|
"No valid credentials provided for the selected authentication type"
|
|
)
|
|
storage_config = config.copy()
|
|
|
|
tool_id = data.get("id")
|
|
existing_doc = None
|
|
existing_encrypted = None
|
|
if tool_id:
|
|
with db_readonly() as conn:
|
|
repo = UserToolsRepository(conn)
|
|
existing_doc = repo.get_any(tool_id, user)
|
|
if existing_doc and existing_doc.get("name") == "mcp_tool":
|
|
existing_encrypted = (existing_doc.get("config") or {}).get(
|
|
"encrypted_credentials"
|
|
)
|
|
else:
|
|
existing_doc = None
|
|
|
|
if auth_credentials:
|
|
if existing_encrypted:
|
|
existing_secrets = decrypt_credentials(existing_encrypted, user)
|
|
existing_secrets.update(auth_credentials)
|
|
auth_credentials = existing_secrets
|
|
storage_config["encrypted_credentials"] = encrypt_credentials(
|
|
auth_credentials, user
|
|
)
|
|
elif existing_encrypted:
|
|
storage_config["encrypted_credentials"] = existing_encrypted
|
|
|
|
for field in [
|
|
"api_key",
|
|
"bearer_token",
|
|
"username",
|
|
"password",
|
|
"api_key_header",
|
|
"redirect_uri",
|
|
]:
|
|
storage_config.pop(field, None)
|
|
from docsgpt.connectors.permissions import apply_default_permissions
|
|
|
|
transformed_actions = apply_default_permissions(
|
|
"mcp_tool", transform_actions(actions_metadata),
|
|
)
|
|
|
|
display_name = data["displayName"]
|
|
connection_id = _mcp_connection(
|
|
user, storage_config, auth_type, auth_credentials, display_name,
|
|
) or _previous_connection(existing_doc, storage_config, user)
|
|
if connection_id and auth_type != "oauth":
|
|
# The secret lives on the connection only.
|
|
storage_config.pop("encrypted_credentials", None)
|
|
description = f"MCP Server: {storage_config.get('server_url', 'Unknown')}"
|
|
status_bool = bool(data.get("status", True))
|
|
|
|
with db_session() as conn:
|
|
repo = UserToolsRepository(conn)
|
|
if existing_doc:
|
|
repo.update(
|
|
str(existing_doc["id"]), user,
|
|
{
|
|
"display_name": display_name,
|
|
"custom_name": display_name,
|
|
"description": description,
|
|
"config": storage_config,
|
|
"actions": transformed_actions,
|
|
"status": status_bool,
|
|
"connection_id": connection_id,
|
|
},
|
|
)
|
|
saved_id = str(existing_doc["id"])
|
|
response_data = {
|
|
"success": True,
|
|
"id": saved_id,
|
|
"message": f"MCP server updated successfully! Discovered {len(transformed_actions)} tools.",
|
|
"tools_count": len(transformed_actions),
|
|
}
|
|
else:
|
|
# Fall back to find_by_user_and_name — the original
|
|
# dual-write path also ran an existence check before
|
|
# deciding between insert and update.
|
|
existing_by_name = repo.find_by_user_and_name(user, "mcp_tool")
|
|
if tool_id is None and existing_by_name and (
|
|
(existing_by_name.get("config") or {}).get("server_url")
|
|
== storage_config.get("server_url")
|
|
):
|
|
repo.update(
|
|
str(existing_by_name["id"]), user,
|
|
{
|
|
"display_name": display_name,
|
|
"custom_name": display_name,
|
|
"description": description,
|
|
"config": storage_config,
|
|
"actions": transformed_actions,
|
|
"status": status_bool,
|
|
"connection_id": connection_id,
|
|
},
|
|
)
|
|
saved_id = str(existing_by_name["id"])
|
|
response_data = {
|
|
"success": True,
|
|
"id": saved_id,
|
|
"message": f"MCP server updated successfully! Discovered {len(transformed_actions)} tools.",
|
|
"tools_count": len(transformed_actions),
|
|
}
|
|
else:
|
|
created = repo.create(
|
|
user, "mcp_tool",
|
|
config=storage_config,
|
|
custom_name=display_name,
|
|
display_name=display_name,
|
|
description=description,
|
|
config_requirements={},
|
|
actions=transformed_actions,
|
|
status=status_bool,
|
|
connection_id=connection_id,
|
|
)
|
|
saved_id = str(created["id"])
|
|
response_data = {
|
|
"success": True,
|
|
"id": saved_id,
|
|
"message": f"MCP server created successfully! Discovered {len(transformed_actions)} tools.",
|
|
"tools_count": len(transformed_actions),
|
|
}
|
|
if tool_id and existing_doc is None:
|
|
# Client requested update on a non-existent tool id.
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"error": "Tool not found or access denied",
|
|
}
|
|
),
|
|
404,
|
|
)
|
|
return make_response(jsonify(response_data), 200)
|
|
except ValueError as e:
|
|
current_app.logger.warning(f"Invalid MCP server save request: {e}")
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Invalid MCP server configuration"}),
|
|
400,
|
|
)
|
|
except Exception as e:
|
|
current_app.logger.error(f"Error saving MCP server: {e}", exc_info=True)
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Failed to save MCP server"}),
|
|
500,
|
|
)
|
|
|
|
|
|
@tools_mcp_ns.route("/mcp_server/callback")
|
|
class MCPOAuthCallback(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"MCPServerCallbackModel",
|
|
{
|
|
"code": fields.String(required=True, description="Authorization code"),
|
|
"state": fields.String(required=True, description="State parameter"),
|
|
"error": fields.String(
|
|
required=False, description="Error message (if any)"
|
|
),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(
|
|
description="Handle OAuth callback by providing the authorization code and state"
|
|
)
|
|
def get(self):
|
|
code = request.args.get("code")
|
|
state = request.args.get("state")
|
|
error = request.args.get("error")
|
|
|
|
if error:
|
|
params = {
|
|
"status": "error",
|
|
"message": f"OAuth error: {error}. Please try again and make sure to grant all requested permissions, including offline access.",
|
|
"provider": "mcp_tool",
|
|
}
|
|
return redirect(f"/api/connectors/callback-status?{urlencode(params)}")
|
|
if not code or not state:
|
|
return redirect(
|
|
"/api/connectors/callback-status?status=error&message=Authorization+code+or+state+not+provided.+Please+complete+the+authorization+process+and+make+sure+to+grant+offline+access.&provider=mcp_tool"
|
|
)
|
|
try:
|
|
redis_client = get_redis_instance()
|
|
if not redis_client:
|
|
return redirect(
|
|
"/api/connectors/callback-status?status=error&message=Internal+server+error:+Redis+not+available.&provider=mcp_tool"
|
|
)
|
|
manager = MCPOAuthManager(redis_client)
|
|
success = manager.handle_oauth_callback(state, code, error)
|
|
if success:
|
|
return redirect(
|
|
"/api/connectors/callback-status?status=success&message=Authorization+code+received+successfully.+You+can+close+this+window.&provider=mcp_tool"
|
|
)
|
|
else:
|
|
return redirect(
|
|
"/api/connectors/callback-status?status=error&message=OAuth+callback+failed.&provider=mcp_tool"
|
|
)
|
|
except Exception as e:
|
|
current_app.logger.error(
|
|
f"Error handling MCP OAuth callback: {str(e)}", exc_info=True
|
|
)
|
|
return redirect(
|
|
"/api/connectors/callback-status?status=error&message=Internal+server+error.&provider=mcp_tool"
|
|
)
|
|
|
|
|
|
@tools_mcp_ns.route("/mcp_server/auth_status")
|
|
class MCPAuthStatus(Resource):
|
|
@api.doc(
|
|
description="Batch check auth status for all MCP tools. "
|
|
"Lightweight DB-only check — no network calls to MCP servers."
|
|
)
|
|
def get(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
try:
|
|
with db_readonly() as conn:
|
|
tools_repo = UserToolsRepository(conn)
|
|
sessions_repo = ConnectorSessionsRepository(conn)
|
|
all_tools = tools_repo.list_for_user(user)
|
|
mcp_tools = [t for t in all_tools if t.get("name") == "mcp_tool"]
|
|
if not mcp_tools:
|
|
return make_response(
|
|
jsonify({"success": True, "statuses": {}}), 200
|
|
)
|
|
|
|
from docsgpt.connectors import service
|
|
|
|
# Read from connection status alone: status checks never
|
|
# decrypt credentials.
|
|
statuses: dict = {}
|
|
for tool in mcp_tools:
|
|
tool_id = str(tool["id"])
|
|
config = tool.get("config") or {}
|
|
auth_type = config.get("auth_type", "none")
|
|
row = None
|
|
if tool.get("connection_id"):
|
|
row = sessions_repo.get(str(tool["connection_id"]))
|
|
elif auth_type == "oauth" and config.get("server_url"):
|
|
parsed = urlparse(config["server_url"])
|
|
row = sessions_repo.get_by_user_provider(
|
|
user, service.mcp_provider(f"{parsed.scheme}://{parsed.netloc}"),
|
|
)
|
|
if row is not None:
|
|
connected = service.normalize_status(row) == service.STATUS_CONNECTED
|
|
if auth_type == "oauth" or not connected:
|
|
statuses[tool_id] = "connected" if connected else "needs_auth"
|
|
else:
|
|
statuses[tool_id] = "configured"
|
|
elif auth_type == "oauth":
|
|
statuses[tool_id] = "needs_auth"
|
|
else:
|
|
statuses[tool_id] = "configured"
|
|
|
|
return make_response(jsonify({"success": True, "statuses": statuses}), 200)
|
|
except Exception as e:
|
|
current_app.logger.error(
|
|
"Error checking MCP auth status: %s", e, exc_info=True
|
|
)
|
|
return make_response(
|
|
jsonify({"success": False, "error": "Failed to check auth status"}),
|
|
500,
|
|
)
|