chore: import upstream snapshot with attribution

application/api/user/tools/mcp.py

@@ -0,0 +1,516 @@
+"""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 application.agents.tools.mcp_tool import MCPOAuthManager, MCPTool
+from application.api import api
+from application.api.user.tools.routes import transform_actions
+from application.cache import get_redis_instance
+from application.core.url_validation import SSRFError, validate_url
+from application.security.encryption import decrypt_credentials, encrypt_credentials
+from application.storage.db.repositories.connector_sessions import (
+    ConnectorSessionsRepository,
+)
+from application.storage.db.repositories.user_tools import UserToolsRepository
+from application.storage.db.session import db_readonly, db_session
+from application.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 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["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
+
+
+@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)
+
+            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)
+
+            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)
+            transformed_actions = 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))
+
+            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,
+                        },
+                    )
+                    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,
+                            },
+                        )
+                        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,
+                        )
+                        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
+                    )
+
+                oauth_server_urls: dict = {}
+                statuses: dict = {}
+                for tool in mcp_tools:
+                    tool_id = str(tool["id"])
+                    config = tool.get("config") or {}
+                    auth_type = config.get("auth_type", "none")
+                    if auth_type == "oauth":
+                        server_url = config.get("server_url", "")
+                        if server_url:
+                            parsed = urlparse(server_url)
+                            base_url = f"{parsed.scheme}://{parsed.netloc}"
+                            oauth_server_urls[tool_id] = base_url
+                        else:
+                            statuses[tool_id] = "needs_auth"
+                    else:
+                        statuses[tool_id] = "configured"
+
+                if oauth_server_urls:
+                    # Look up a session per distinct base URL. MCP sessions
+                    # are stored with ``provider = "mcp:<server_url>"``
+                    # 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()):
+                        session = sessions_repo.get_by_user_and_server_url(
+                            user, base_url,
+                        )
+                        tokens = (
+                            (session or {}).get("session_data", {}) or {}
+                        ).get("tokens", {}) or {}
+                        # 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(
+                            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):
+                            statuses[tool_id] = "connected"
+                        else:
+                            statuses[tool_id] = "needs_auth"
+
+            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,
+            )