Files
DocsGPT/tests/api/user/test_tools_mcp_pg.py
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

601 lines
22 KiB
Python

"""Tests for docsgpt/api/user/tools/mcp.py using real PG."""
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
@pytest.fixture(autouse=True)
def _mcp_servers_allowed():
"""The admin switch fails closed without a database; these tests are not about it.
The policy itself is covered in ``tests/connectors``.
"""
with patch("docsgpt.api.user.tools.mcp._mcp_policy_error", return_value=None):
yield
@pytest.fixture
def app():
return Flask(__name__)
@contextmanager
def _patch_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"docsgpt.api.user.tools.mcp.db_session", _yield
), patch(
"docsgpt.api.user.tools.mcp.db_readonly", _yield
):
yield
class TestSanitizeMcpTransport:
def test_defaults_to_auto(self):
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
cfg = {}
got = _sanitize_mcp_transport(cfg)
assert got == "auto"
assert cfg["transport_type"] == "auto"
def test_accepts_supported_transports(self):
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
for t in ("auto", "sse", "http"):
cfg = {"transport_type": t}
assert _sanitize_mcp_transport(cfg) == t
def test_strips_command_and_args(self):
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
cfg = {"transport_type": "http", "command": "/bin/x", "args": ["a"]}
_sanitize_mcp_transport(cfg)
assert "command" not in cfg
assert "args" not in cfg
def test_unsupported_transport_raises(self):
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
with pytest.raises(ValueError):
_sanitize_mcp_transport({"transport_type": "websocket"})
class TestExtractAuthCredentials:
def test_api_key_auth(self):
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
got = _extract_auth_credentials({
"auth_type": "api_key",
"api_key": "secret",
"api_key_header": "X-API-Key",
})
assert got == {"api_key": "secret", "api_key_header": "X-API-Key"}
def test_bearer_auth(self):
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
got = _extract_auth_credentials({
"auth_type": "bearer",
"bearer_token": "my-token",
})
assert got == {"bearer_token": "my-token"}
def test_basic_auth(self):
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
got = _extract_auth_credentials({
"auth_type": "basic",
"username": "u", "password": "p",
})
assert got == {"username": "u", "password": "p"}
def test_none_auth_empty_creds(self):
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
assert _extract_auth_credentials({"auth_type": "none"}) == {}
class TestValidateMcpServerUrl:
def test_empty_url_raises(self):
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
with pytest.raises(ValueError):
_validate_mcp_server_url({})
def test_missing_server_url(self):
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
with pytest.raises(ValueError):
_validate_mcp_server_url({"server_url": ""})
def test_ssrf_url_raises(self):
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
with pytest.raises(ValueError):
_validate_mcp_server_url({"server_url": "http://127.0.0.1"})
def test_valid_public_url_passes(self):
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
# Should not raise for a public-ish URL
try:
_validate_mcp_server_url({"server_url": "https://example.com/mcp"})
except ValueError as e:
# If SSRF rules reject example.com for some reason, accept that
if "Invalid" not in str(e):
raise
class TestTestMCPServerConfig:
def test_returns_401_unauthenticated(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
with app.test_request_context(
"/api/mcp_server/test", method="POST",
json={"config": {}},
):
from flask import request
request.decoded_token = None
response = TestMCPServerConfig().post()
assert response.status_code == 401
def test_returns_400_missing_config(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
with app.test_request_context(
"/api/mcp_server/test", method="POST", json={},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 400
def test_unsupported_transport_returns_400(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
with app.test_request_context(
"/api/mcp_server/test", method="POST",
json={"config": {"transport_type": "websocket"}},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 400
def test_missing_url_returns_400(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
with app.test_request_context(
"/api/mcp_server/test", method="POST",
json={"config": {"transport_type": "http"}},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 400
def test_connection_success(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
fake_tool = MagicMock()
fake_tool.test_connection.return_value = {
"success": True, "message": "OK",
"tools_count": 3, "tools": ["a", "b", "c"],
}
with patch(
"docsgpt.api.user.tools.mcp.MCPTool",
return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/test", method="POST",
json={
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
"auth_type": "none",
},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 200
assert response.json["success"] is True
assert response.json["tools_count"] == 3
def test_connection_failure_returns_200_with_failure_message(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
fake_tool = MagicMock()
fake_tool.test_connection.return_value = {
"success": False, "message": "Cannot reach server",
}
with patch(
"docsgpt.api.user.tools.mcp.MCPTool",
return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/test", method="POST",
json={
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 200
assert response.json["success"] is False
def test_oauth_required_returns_200(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
fake_tool = MagicMock()
fake_tool.test_connection.return_value = {
"success": False,
"requires_oauth": True,
"auth_url": "https://auth/ex",
"task_id": "task-123",
"message": "OAuth required",
}
with patch(
"docsgpt.api.user.tools.mcp.MCPTool",
return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/test", method="POST",
json={
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
"auth_type": "oauth",
},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 200
assert response.json["requires_oauth"] is True
# The client follows the sign-in by this task's events.
assert response.json["task_id"] == "task-123"
def test_unexpected_exception_returns_500(self, app):
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
with patch(
"docsgpt.api.user.tools.mcp.MCPTool",
side_effect=RuntimeError("boom"),
), app.test_request_context(
"/api/mcp_server/test", method="POST",
json={
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = TestMCPServerConfig().post()
assert response.status_code == 500
class TestMCPServerSave:
def test_returns_401_unauthenticated(self, app):
from docsgpt.api.user.tools.mcp import MCPServerSave
with app.test_request_context(
"/api/mcp_server/save", method="POST",
json={"displayName": "n", "config": {}},
):
from flask import request
request.decoded_token = None
response = MCPServerSave().post()
assert response.status_code == 401
def test_returns_400_missing_fields(self, app):
from docsgpt.api.user.tools.mcp import MCPServerSave
with app.test_request_context(
"/api/mcp_server/save", method="POST", json={},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = MCPServerSave().post()
assert response.status_code == 400
def test_unsupported_transport_returns_400(self, app):
from docsgpt.api.user.tools.mcp import MCPServerSave
with app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"displayName": "Srv",
"config": {"transport_type": "bogus"},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = MCPServerSave().post()
assert response.status_code == 400
def test_missing_server_url_returns_400(self, app):
from docsgpt.api.user.tools.mcp import MCPServerSave
with app.test_request_context(
"/api/mcp_server/save", method="POST",
json={"displayName": "Srv", "config": {"transport_type": "http"}},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = MCPServerSave().post()
assert response.status_code == 400
def test_oauth_missing_task_id_returns_400(self, app):
from docsgpt.api.user.tools.mcp import MCPServerSave
# No OAuth task and no stored sign-in that works.
unauthorized = MagicMock()
unauthorized.discover_tools.side_effect = RuntimeError("401")
with patch("docsgpt.api.user.tools.mcp.MCPTool", return_value=unauthorized), app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"displayName": "Srv",
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
"auth_type": "oauth",
},
},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = MCPServerSave().post()
assert response.status_code == 400
def test_oauth_without_task_uses_the_stored_sign_in(self, app, pg_conn):
"""Signed in before: the server answers with the saved tokens, no new handshake."""
from docsgpt.api.user.tools.mcp import MCPServerSave
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
connection = ConnectorSessionsRepository(pg_conn).create(
"u-signed-in", "mcp:https://mcp.linear.app", connector_key="mcp:linear", auth_kind="mcp_oauth",
display_name="Linear", account_label="Linear", server_url="https://mcp.linear.app",
)
signed_in = MagicMock()
signed_in.get_actions_metadata.return_value = [{"name": "search"}]
with _patch_db(pg_conn), patch(
"docsgpt.api.user.tools.mcp.MCPTool", return_value=signed_in,
), patch(
"docsgpt.api.user.tools.mcp._mcp_connection", return_value=str(connection["id"]),
), app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"displayName": "Linear",
"config": {
"transport_type": "http",
"server_url": "https://mcp.linear.app/mcp",
"auth_type": "oauth",
},
},
):
from flask import request
request.decoded_token = {"sub": "u-signed-in"}
response = MCPServerSave().post()
assert response.status_code == 200
signed_in.discover_tools.assert_called_once()
def test_creates_mcp_tool_successfully(self, app, pg_conn):
from docsgpt.api.user.tools.mcp import MCPServerSave
user = "u-mcp-save"
fake_tool = MagicMock()
fake_tool.discover_tools.return_value = {"tools": ["t1"]}
fake_tool.get_actions_metadata.return_value = [{"name": "t1"}]
with _patch_db(pg_conn), patch(
"docsgpt.api.user.tools.mcp.MCPTool",
return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"displayName": "My MCP",
"config": {
"transport_type": "http",
"server_url": "https://example.com/mcp",
"auth_type": "none",
},
"status": True,
},
):
from flask import request
request.decoded_token = {"sub": user}
response = MCPServerSave().post()
assert response.status_code in (200, 201)
def test_saving_an_existing_server_keeps_fixed_values(self, app, pg_conn):
from docsgpt.api.user.tools.mcp import MCPServerSave
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
user = "u-mcp-resave"
repo = UserToolsRepository(pg_conn)
existing = repo.create(
user, "mcp_tool",
config={"server_url": "https://example.com/mcp", "auth_type": "none"},
display_name="My MCP",
actions=[{
"name": "search",
"active": True,
"parameters": {"properties": {
"q": {"type": "string", "filled_by_llm": True, "value": ""},
"team": {"type": "string", "filled_by_llm": False, "value": "ENG"},
}},
}],
)
fake_tool = MagicMock()
fake_tool.get_actions_metadata.return_value = [{
"name": "search",
"parameters": {"properties": {"q": {"type": "string"}, "team": {"type": "string"}}},
}]
with _patch_db(pg_conn), patch(
"docsgpt.api.user.tools.mcp.MCPTool", return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"id": str(existing["id"]),
"displayName": "My MCP",
"config": {"transport_type": "http", "server_url": "https://example.com/mcp", "auth_type": "none"},
},
):
from flask import request
request.decoded_token = {"sub": user}
response = MCPServerSave().post()
assert response.status_code == 200
team = repo.get_any(str(existing["id"]), user)["actions"][0]["parameters"]["properties"]["team"]
assert team["filled_by_llm"] is False and team["value"] == "ENG"
class TestSignInServerWithoutAConnection:
def test_an_oauth_server_is_not_saved_without_its_connection(self, app, pg_conn):
"""A client still cached from a removed connection can answer the
discovery; the tool must not then be saved as an unconnected server."""
from docsgpt.api.user.tools.mcp import MCPServerSave
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
user = "u-mcp-no-connection"
fake_tool = MagicMock()
fake_tool.get_actions_metadata.return_value = [{"name": "list_issues", "parameters": {"properties": {}}}]
with _patch_db(pg_conn), patch(
"docsgpt.api.user.tools.mcp.MCPTool", return_value=fake_tool,
), app.test_request_context(
"/api/mcp_server/save", method="POST",
json={
"displayName": "Linear",
"config": {"transport_type": "http", "server_url": "https://mcp.linear.app/mcp", "auth_type": "oauth"},
"status": True,
},
):
from flask import request
request.decoded_token = {"sub": user}
response = MCPServerSave().post()
assert response.status_code == 400
assert UserToolsRepository(pg_conn).list_for_user(user) == []
class TestMCPOAuthCallback:
def test_error_param_redirects_error(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
with app.test_request_context(
"/api/mcp_server/oauth_callback?error=access_denied"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
assert "status=error" in response.location
def test_missing_code_or_state_redirects_error(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
with app.test_request_context(
"/api/mcp_server/oauth_callback"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
def test_success_redirects_success(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
fake_redis = MagicMock()
fake_manager = MagicMock()
fake_manager.handle_oauth_callback.return_value = True
with patch(
"docsgpt.api.user.tools.mcp.get_redis_instance",
return_value=fake_redis,
), patch(
"docsgpt.api.user.tools.mcp.MCPOAuthManager",
return_value=fake_manager,
), app.test_request_context(
"/api/mcp_server/oauth_callback?code=c1&state=s1"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
assert "status=success" in response.location
def test_manager_failure_redirects_error(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
fake_redis = MagicMock()
fake_manager = MagicMock()
fake_manager.handle_oauth_callback.return_value = False
with patch(
"docsgpt.api.user.tools.mcp.get_redis_instance",
return_value=fake_redis,
), patch(
"docsgpt.api.user.tools.mcp.MCPOAuthManager",
return_value=fake_manager,
), app.test_request_context(
"/api/mcp_server/oauth_callback?code=c1&state=s1"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
assert "status=error" in response.location
def test_no_redis_redirects_error(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
with patch(
"docsgpt.api.user.tools.mcp.get_redis_instance",
return_value=None,
), app.test_request_context(
"/api/mcp_server/oauth_callback?code=c&state=s"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
assert "Redis" in response.location or "status=error" in response.location
def test_exception_redirects_error(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
with patch(
"docsgpt.api.user.tools.mcp.get_redis_instance",
side_effect=RuntimeError("boom"),
), app.test_request_context(
"/api/mcp_server/oauth_callback?code=c&state=s"
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
class TestMCPAuthStatus:
def test_returns_401_unauthenticated(self, app):
from docsgpt.api.user.tools.mcp import MCPAuthStatus
with app.test_request_context(
"/api/mcp_server/auth_status"
):
from flask import request
request.decoded_token = None
response = MCPAuthStatus().get()
assert response.status_code == 401
class TestMCPOAuthCallbackIssuer:
def test_passes_the_issuer_to_the_waiting_sign_in(self, app):
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
manager = MagicMock()
manager.handle_oauth_callback.return_value = True
with patch(
"docsgpt.api.user.tools.mcp.get_redis_instance", return_value=MagicMock(),
), patch(
"docsgpt.api.user.tools.mcp.MCPOAuthManager", return_value=manager,
), app.test_request_context(
"/api/mcp_server/callback?code=c&state=s&iss=https%3A%2F%2Fmcp.linear.app",
):
response = MCPOAuthCallback().get()
assert response.status_code == 302
manager.handle_oauth_callback.assert_called_once_with(
"s", "c", None, iss="https://mcp.linear.app",
)