Files
DocsGPT/docsgpt/api/user/tools/mcp.py
T
arc53-machine 56ababae71 Stop a removed MCP connection from saving its server as a custom tool
After removing a Linear connection, connecting again could skip signing in
and save Linear as an unconnected custom tool: a client cached before the
removal still held the old tokens, and a late token write re-created a
connection for them. A token write for a named connection no longer creates
one, removing or disconnecting a connection drops its cached clients, and a
sign-in server is never saved without its connection.
2026-09-29 13:13:27 +01:00

632 lines
27 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.tool_pins import carry_pins_between
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:
# Fail closed: a server whose admin switch cannot be read is not contacted.
current_app.logger.warning("Could not read connector policies", exc_info=True)
return make_response(
jsonify({"success": False, "error": "Could not check whether this MCP server is allowed"}),
503,
)
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", "task_id")
}
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" and not config.get("oauth_task_id"):
# No new handshake: the caller signed in to this server
# before, so its stored tokens answer the discovery.
try:
mcp_tool = MCPTool(config=mcp_config, user_id=user)
mcp_tool.discover_tools()
actions_metadata = mcp_tool.get_actions_metadata()
except Exception:
return make_response(
jsonify(
{
"success": False,
"error": "Connection not authorized. Please complete the OAuth authorization first.",
}
),
400,
)
elif auth_type == "oauth":
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 auth_type == "oauth" and not connection_id:
# A sign-in server's tokens live on its connection. Without one
# (it was removed, and a client cached before that answered the
# discovery) the tool would be saved unconnected.
return make_response(
jsonify({
"success": False,
"error": "Not signed in to this server. Sign in again to connect it.",
}),
400,
)
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,
# Fixed values the user set survive a re-save.
"actions": carry_pins_between(existing_doc.get("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": carry_pins_between(existing_by_name.get("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, iss=request.args.get("iss"))
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,
)