mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 06:15:57 +00:00
Roles audit and revamp
This commit is contained in:
Pavel committed
1 parent
208d34abc7
commit
98afa5d93c
134 files changed
+13725
-1702
No files matched your search
+157
-90
@@ -7,7 +7,15 @@ 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.api.user.resource_access import AccessDenied, require
|
||||
from docsgpt.api.user.team_sharing import visible_with_access
|
||||
from docsgpt.api.user.tools.routes import (
|
||||
_CREDENTIALS_FOR_NEW_SERVER,
|
||||
_MCP_CREDENTIAL_AUTH_TYPES,
|
||||
_mcp_host_changed,
|
||||
denied_response,
|
||||
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
|
||||
@@ -75,12 +83,60 @@ def _validate_mcp_server_url(config: dict) -> None:
|
||||
raise ValueError(f"Invalid server URL: {exc}") from exc
|
||||
|
||||
|
||||
_ONLY_OWNER_RECONNECTS = "Only the owner can reconnect this account"
|
||||
|
||||
|
||||
def _existing_mcp_context(tool_id, user, config):
|
||||
"""Resolve the stored MCP tool a test/save refers to, and its credentials.
|
||||
|
||||
With no ``tool_id`` the caller acts on their own new server. With one,
|
||||
the caller needs ``edit_credentials`` on that tool and everything runs as
|
||||
its owner. Stored secrets are write-only, so an empty secret field reuses
|
||||
the stored one while the host is unchanged; a new host never inherits them.
|
||||
|
||||
Returns:
|
||||
``(existing_doc, owner_id, is_owner, moved, credentials)``, or a Flask
|
||||
response (404 / 400) to return as is.
|
||||
|
||||
Raises:
|
||||
AccessDenied: the caller can't see the tool (404) or can't change
|
||||
its credentials (403).
|
||||
"""
|
||||
auth_credentials = _extract_auth_credentials(config)
|
||||
if not tool_id:
|
||||
return None, user, True, False, auth_credentials
|
||||
with db_readonly() as conn:
|
||||
ra = require(conn, "tool", tool_id, user, "edit_credentials")
|
||||
existing_doc = UserToolsRepository(conn).get_any(ra.resource_id, ra.owner_id)
|
||||
if not existing_doc or existing_doc.get("name") != "mcp_tool":
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Tool not found or access denied"}), 404,
|
||||
)
|
||||
existing_config = existing_doc.get("config") or {}
|
||||
moved = _mcp_host_changed(config, existing_config)
|
||||
auth_type = config.get("auth_type", "none")
|
||||
new_secret_keys = set(auth_credentials) - {"api_key_header"}
|
||||
if moved and auth_type in _MCP_CREDENTIAL_AUTH_TYPES and not new_secret_keys:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": _CREDENTIALS_FOR_NEW_SERVER}), 400
|
||||
)
|
||||
credentials = dict(auth_credentials)
|
||||
existing_encrypted = None if moved else existing_config.get("encrypted_credentials")
|
||||
if existing_encrypted:
|
||||
credentials = {**decrypt_credentials(existing_encrypted, ra.owner_id), **auth_credentials}
|
||||
return existing_doc, ra.owner_id, ra.access == "owner", moved, credentials
|
||||
|
||||
|
||||
@tools_mcp_ns.route("/mcp_server/test")
|
||||
class TestMCPServerConfig(Resource):
|
||||
@api.expect(
|
||||
api.model(
|
||||
"MCPServerTestModel",
|
||||
{
|
||||
"id": fields.String(
|
||||
required=False,
|
||||
description="Stored tool to test with (empty secrets reuse its stored ones)",
|
||||
),
|
||||
"config": fields.Raw(
|
||||
required=True, description="MCP server configuration to test"
|
||||
),
|
||||
@@ -111,11 +167,20 @@ class TestMCPServerConfig(Resource):
|
||||
|
||||
_validate_mcp_server_url(config)
|
||||
|
||||
auth_credentials = _extract_auth_credentials(config)
|
||||
ctx = _existing_mcp_context(data.get("id"), user, config)
|
||||
if not isinstance(ctx, tuple):
|
||||
return ctx
|
||||
_existing_doc, owner_id, is_owner, _moved, auth_credentials = ctx
|
||||
if not is_owner and config.get("auth_type") == "oauth":
|
||||
# An OAuth flow would store tokens under the editor's account
|
||||
# (and its popup event goes to that account), not the owner's.
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": _ONLY_OWNER_RECONNECTS}), 403
|
||||
)
|
||||
test_config = config.copy()
|
||||
test_config["auth_credentials"] = auth_credentials
|
||||
|
||||
mcp_tool = MCPTool(config=test_config, user_id=user)
|
||||
mcp_tool = MCPTool(config=test_config, user_id=owner_id)
|
||||
result = mcp_tool.test_connection()
|
||||
|
||||
if result.get("requires_oauth"):
|
||||
@@ -148,6 +213,8 @@ class TestMCPServerConfig(Resource):
|
||||
"tools": result.get("tools", []),
|
||||
}
|
||||
return make_response(jsonify(safe_result), 200)
|
||||
except AccessDenied as e:
|
||||
return denied_response(e)
|
||||
except ValueError as e:
|
||||
current_app.logger.warning(f"Invalid MCP server test request: {e}")
|
||||
return make_response(
|
||||
@@ -207,13 +274,55 @@ class MCPServerSave(Resource):
|
||||
|
||||
_validate_mcp_server_url(config)
|
||||
|
||||
auth_credentials = _extract_auth_credentials(config)
|
||||
# An existing id is always an update of THAT row, written as its
|
||||
# owner; it never falls through to creating a copy for the caller.
|
||||
ctx = _existing_mcp_context(data.get("id"), user, config)
|
||||
if not isinstance(ctx, tuple):
|
||||
return ctx
|
||||
existing_doc, owner_id, is_owner, moved, merged_credentials = ctx
|
||||
existing_config = (existing_doc or {}).get("config") or {}
|
||||
auth_type = config.get("auth_type", "none")
|
||||
mcp_config = config.copy()
|
||||
mcp_config["auth_credentials"] = auth_credentials
|
||||
mcp_config["auth_credentials"] = merged_credentials
|
||||
keep_actions = False
|
||||
|
||||
if auth_type == "oauth":
|
||||
if not config.get("oauth_task_id"):
|
||||
if config.get("oauth_task_id"):
|
||||
if not is_owner:
|
||||
# The OAuth flow stores tokens under the account that
|
||||
# ran it; reconnecting as the owner is owner-only.
|
||||
return make_response(
|
||||
jsonify({
|
||||
"success": False,
|
||||
"message": _ONLY_OWNER_RECONNECTS,
|
||||
}),
|
||||
403,
|
||||
)
|
||||
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 (
|
||||
existing_doc is not None
|
||||
and not moved
|
||||
and existing_config.get("auth_type") == "oauth"
|
||||
):
|
||||
# Editing an already-connected server: keep its tools.
|
||||
actions_metadata = existing_doc.get("actions") or []
|
||||
keep_actions = True
|
||||
else:
|
||||
return make_response(
|
||||
jsonify(
|
||||
{
|
||||
@@ -223,24 +332,8 @@ class MCPServerSave(Resource):
|
||||
),
|
||||
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)
|
||||
elif auth_type == "none" or merged_credentials:
|
||||
mcp_tool = MCPTool(config=mcp_config, user_id=owner_id)
|
||||
mcp_tool.discover_tools()
|
||||
actions_metadata = mcp_tool.get_actions_metadata()
|
||||
else:
|
||||
@@ -249,30 +342,10 @@ class MCPServerSave(Resource):
|
||||
)
|
||||
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
|
||||
if merged_credentials:
|
||||
storage_config["encrypted_credentials"] = encrypt_credentials(
|
||||
auth_credentials, user
|
||||
merged_credentials, owner_id
|
||||
)
|
||||
elif existing_encrypted:
|
||||
storage_config["encrypted_credentials"] = existing_encrypted
|
||||
|
||||
for field in [
|
||||
"api_key",
|
||||
@@ -283,58 +356,54 @@ class MCPServerSave(Resource):
|
||||
"redirect_uri",
|
||||
]:
|
||||
storage_config.pop(field, None)
|
||||
transformed_actions = transform_actions(actions_metadata)
|
||||
# Kept actions already carry the owner's on/off and approval flags.
|
||||
transformed_actions = actions_metadata if keep_actions else transform_actions(actions_metadata)
|
||||
|
||||
display_name = data["displayName"]
|
||||
description = f"MCP Server: {storage_config.get('server_url', 'Unknown')}"
|
||||
status_bool = bool(data.get("status", True))
|
||||
fields_out = {
|
||||
"display_name": display_name,
|
||||
"custom_name": display_name,
|
||||
"description": description,
|
||||
"config": storage_config,
|
||||
"actions": transformed_actions,
|
||||
}
|
||||
updated_message = (
|
||||
f"MCP server updated successfully! Discovered {len(transformed_actions)} tools."
|
||||
)
|
||||
|
||||
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,
|
||||
},
|
||||
)
|
||||
if existing_doc is not None:
|
||||
# ``status`` is the owner's own chat switch; an editor's
|
||||
# save doesn't flip it.
|
||||
if is_owner:
|
||||
fields_out["status"] = status_bool
|
||||
repo.update(str(existing_doc["id"]), owner_id, fields_out)
|
||||
saved_id = str(existing_doc["id"])
|
||||
response_data = {
|
||||
"success": True,
|
||||
"id": saved_id,
|
||||
"message": f"MCP server updated successfully! Discovered {len(transformed_actions)} tools.",
|
||||
"message": updated_message,
|
||||
"tools_count": len(transformed_actions),
|
||||
}
|
||||
else:
|
||||
fields_out["status"] = status_bool
|
||||
# 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 (
|
||||
if 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,
|
||||
},
|
||||
)
|
||||
repo.update(str(existing_by_name["id"]), user, fields_out)
|
||||
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.",
|
||||
"message": updated_message,
|
||||
"tools_count": len(transformed_actions),
|
||||
}
|
||||
else:
|
||||
@@ -355,18 +424,9 @@ class MCPServerSave(Resource):
|
||||
"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 AccessDenied as e:
|
||||
return denied_response(e)
|
||||
except ValueError as e:
|
||||
current_app.logger.warning(f"Invalid MCP server save request: {e}")
|
||||
return make_response(
|
||||
@@ -455,6 +515,13 @@ class MCPAuthStatus(Resource):
|
||||
tools_repo = UserToolsRepository(conn)
|
||||
sessions_repo = ConnectorSessionsRepository(conn)
|
||||
all_tools = tools_repo.list_for_user(user)
|
||||
owned_ids = {str(t["id"]) for t in all_tools}
|
||||
# Team-shared MCP servers the caller can see run with the
|
||||
# owner's connection, so their status is the owner's.
|
||||
shared_ids = [
|
||||
tid for tid in visible_with_access(conn, user, "tool") if tid not in owned_ids
|
||||
]
|
||||
all_tools = all_tools + tools_repo.list_by_ids(shared_ids)
|
||||
mcp_tools = [t for t in all_tools if t.get("name") == "mcp_tool"]
|
||||
if not mcp_tools:
|
||||
return make_response(
|
||||
@@ -472,7 +539,7 @@ class MCPAuthStatus(Resource):
|
||||
if server_url:
|
||||
parsed = urlparse(server_url)
|
||||
base_url = f"{parsed.scheme}://{parsed.netloc}"
|
||||
oauth_server_urls[tool_id] = base_url
|
||||
oauth_server_urls[tool_id] = (tool.get("user_id") or user, base_url)
|
||||
else:
|
||||
statuses[tool_id] = "needs_auth"
|
||||
else:
|
||||
@@ -484,9 +551,9 @@ class MCPAuthStatus(Resource):
|
||||
# and the URL in ``server_url``; reuse the repo's
|
||||
# per-URL accessor rather than an ad-hoc $in query.
|
||||
url_has_tokens: dict = {}
|
||||
for base_url in set(oauth_server_urls.values()):
|
||||
for owner_id, base_url in set(oauth_server_urls.values()):
|
||||
session = sessions_repo.get_by_user_and_server_url(
|
||||
user, base_url,
|
||||
owner_id, base_url,
|
||||
)
|
||||
tokens = (
|
||||
(session or {}).get("session_data", {}) or {}
|
||||
@@ -494,13 +561,13 @@ class MCPAuthStatus(Resource):
|
||||
# MCP code also stashes tokens into token_info on
|
||||
# the row; consider either present as "connected".
|
||||
token_info = (session or {}).get("token_info") or {}
|
||||
url_has_tokens[base_url] = bool(
|
||||
url_has_tokens[(owner_id, base_url)] = bool(
|
||||
tokens.get("access_token")
|
||||
or token_info.get("access_token")
|
||||
)
|
||||
|
||||
for tool_id, base_url in oauth_server_urls.items():
|
||||
if url_has_tokens.get(base_url):
|
||||
for tool_id, key in oauth_server_urls.items():
|
||||
if url_has_tokens.get(key):
|
||||
statuses[tool_id] = "connected"
|
||||
else:
|
||||
statuses[tool_id] = "needs_auth"
|
||||
|
||||
Reference in new issue
Block a user