538 lines
19 KiB
Python
538 lines
19 KiB
Python
"""
|
|
Copyright 2024-2026 ReplyPier
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, status
|
|
from typing import List
|
|
from app.core.logger import get_logger
|
|
from app.database import get_db
|
|
from app.models.user import User
|
|
from app.core.auth import require_permissions
|
|
from app.repositories.mcp_tool import MCPToolRepository
|
|
from app.repositories.agent import AgentRepository
|
|
|
|
# Enterprise feature check
|
|
try:
|
|
from app.enterprise.repositories.plan import PlanRepository
|
|
from app.enterprise.services.feature_access import require_accessible_subscription
|
|
HAS_ENTERPRISE = True
|
|
except ImportError:
|
|
HAS_ENTERPRISE = False
|
|
from app.models.schemas.mcp_tool import (
|
|
MCPToolCreate, MCPToolUpdate, MCPToolResponse,
|
|
MCPToolToAgentCreate, MCPToolToAgentResponse,
|
|
AgentMCPToolsResponse, MCPToolTestResponse, MCPToolReferencesResponse,
|
|
MCPTransportTypeEnum, unmask_secret_values
|
|
)
|
|
from app.core.outbound_policy import require_public_target
|
|
from app.core.params import DbIntId
|
|
from app.tools.mcp_manager import MCP_URL_LABEL, TEST_CONNECT_BUDGET_SECONDS, MCPToolsManager
|
|
from app.core.config import settings
|
|
from sqlalchemy.orm import Session
|
|
from uuid import UUID
|
|
|
|
def require_stdio_enabled() -> None:
|
|
"""Refuse a STDIO connector where running admin-typed commands is off."""
|
|
if not settings.MCP_STDIO_ENABLED:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Local (STDIO) MCP servers are not available here. Connect a remote server over HTTP or SSE.",
|
|
)
|
|
|
|
|
|
router = APIRouter()
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def check_mcp_feature_access(current_user: User, db: Session):
|
|
"""Check if user has access to MCP tools feature"""
|
|
if not HAS_ENTERPRISE:
|
|
return # Allow access in non-enterprise mode
|
|
|
|
# Accessible = active/trial/past-due-in-period OR cancelled-but-still-in-
|
|
# paid-period; raises 403 when the org has no accessible plan.
|
|
subscription = require_accessible_subscription(db, current_user.organization_id)
|
|
plan_repo = PlanRepository(db)
|
|
|
|
# Check if MCP tools feature is available in the plan
|
|
if not plan_repo.check_feature_availability(str(subscription.plan_id), 'mcp_tools'):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="MCP Tools feature is not available in your current plan. Please upgrade to access this feature."
|
|
)
|
|
|
|
|
|
@router.post("", response_model=MCPToolResponse)
|
|
async def create_mcp_tool(
|
|
mcp_tool_data: MCPToolCreate,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Create a new MCP tool"""
|
|
try:
|
|
# Check MCP tools feature access
|
|
check_mcp_feature_access(current_user, db)
|
|
if mcp_tool_data.transport_type == MCPTransportTypeEnum.STDIO:
|
|
require_stdio_enabled()
|
|
else:
|
|
require_public_target(mcp_tool_data.url, MCP_URL_LABEL)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
|
|
# Set organization_id from current user
|
|
create_data = mcp_tool_data.model_dump()
|
|
create_data['organization_id'] = current_user.organization_id
|
|
|
|
# Check if MCP tool with same name already exists
|
|
existing = mcp_tool_repo.get_by_name(
|
|
mcp_tool_data.name,
|
|
current_user.organization_id
|
|
)
|
|
if existing:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="MCP tool with this name already exists"
|
|
)
|
|
|
|
mcp_tool = mcp_tool_repo.create_mcp_tool(**create_data)
|
|
return mcp_tool
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"MCP tool creation error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("", response_model=List[MCPToolResponse])
|
|
async def get_organization_mcp_tools(
|
|
enabled_only: bool = True,
|
|
current_user: User = Depends(require_permissions("view_all", "manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Get all MCP tools for the current user's organization"""
|
|
try:
|
|
# Check MCP tools feature access
|
|
check_mcp_feature_access(current_user, db)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tools = mcp_tool_repo.get_org_mcp_tools(
|
|
current_user.organization_id,
|
|
enabled_only=enabled_only
|
|
)
|
|
return mcp_tools
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error fetching MCP tools: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/{mcp_tool_id}", response_model=MCPToolResponse)
|
|
async def get_mcp_tool(
|
|
mcp_tool_id: DbIntId,
|
|
current_user: User = Depends(require_permissions("view_all", "manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Get a specific MCP tool"""
|
|
try:
|
|
check_mcp_feature_access(current_user, db)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
# Check if user has access to this MCP tool's organization
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this MCP tool"
|
|
)
|
|
|
|
return mcp_tool
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error fetching MCP tool: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/{mcp_tool_id}/test", response_model=MCPToolTestResponse)
|
|
async def test_mcp_tool(
|
|
mcp_tool_id: DbIntId,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Connect to a configured MCP tool once and report whether it comes up
|
|
and which functions it exposes — surfaces environment problems (e.g. a
|
|
missing npx) that otherwise only show as a silent 0-tool run."""
|
|
try:
|
|
check_mcp_feature_access(current_user, db)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this MCP tool"
|
|
)
|
|
|
|
# A request is waiting on this, so cap it no matter how generous the
|
|
# tool's own timeout is — a hung server must not hold the worker and
|
|
# its DB session for minutes.
|
|
manager = MCPToolsManager(connect_budget=TEST_CONNECT_BUDGET_SECONDS)
|
|
return await manager.test_tool_config(mcp_tool)
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"MCP tool test error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/{mcp_tool_id}/references", response_model=MCPToolReferencesResponse)
|
|
async def get_mcp_tool_references(
|
|
mcp_tool_id: DbIntId,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""What currently points at this connector, so a delete confirmation can
|
|
name what it is about to break. Deliberately ungated, like the delete it
|
|
precedes."""
|
|
try:
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this MCP tool"
|
|
)
|
|
|
|
return MCPToolReferencesResponse(
|
|
agents=mcp_tool_repo.get_mcp_tool_agent_names(mcp_tool_id),
|
|
used_in_investigations=mcp_tool_repo.is_used_in_investigations(
|
|
mcp_tool_id, current_user.organization_id
|
|
),
|
|
)
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error fetching MCP tool references: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.put("/{mcp_tool_id}", response_model=MCPToolResponse)
|
|
async def update_mcp_tool(
|
|
mcp_tool_id: DbIntId,
|
|
update_data: MCPToolUpdate,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Update an existing MCP tool"""
|
|
try:
|
|
check_mcp_feature_access(current_user, db)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
# Check if user has access to this MCP tool's organization
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to update this MCP tool"
|
|
)
|
|
|
|
# Update MCP tool with provided fields
|
|
update_dict = update_data.model_dump(exclude_unset=True)
|
|
|
|
# Secrets are masked on the way out, so a form that round-trips them
|
|
# untouched must leave the stored credential alone.
|
|
if 'env_vars' in update_dict:
|
|
update_dict['env_vars'] = unmask_secret_values(
|
|
update_dict['env_vars'], mcp_tool.env_vars
|
|
)
|
|
if 'headers' in update_dict:
|
|
update_dict['headers'] = unmask_secret_values(
|
|
update_dict['headers'], mcp_tool.headers
|
|
)
|
|
|
|
# transport_type is editable, and MCPToolUpdate carries none of the
|
|
# create-time validators, so check the state the update would leave
|
|
# behind — otherwise a connector can be saved with a transport it has
|
|
# no command or URL for and fails only mid-investigation.
|
|
transport = update_dict.get('transport_type', mcp_tool.transport_type)
|
|
transport = getattr(transport, 'value', transport)
|
|
if transport == MCPTransportTypeEnum.STDIO.value:
|
|
require_stdio_enabled()
|
|
if not update_dict.get('command', mcp_tool.command):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Command is required for STDIO transport"
|
|
)
|
|
elif not update_dict.get('url', mcp_tool.url):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="URL is required for SSE/HTTP transport"
|
|
)
|
|
else:
|
|
require_public_target(update_dict.get('url', mcp_tool.url), MCP_URL_LABEL)
|
|
|
|
# Check for name conflicts if name is being updated
|
|
if 'name' in update_dict:
|
|
existing = mcp_tool_repo.get_by_name(
|
|
update_dict['name'],
|
|
current_user.organization_id
|
|
)
|
|
if existing and existing.id != mcp_tool_id:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="MCP tool with this name already exists"
|
|
)
|
|
|
|
updated_mcp_tool = mcp_tool_repo.update_mcp_tool(mcp_tool_id, **update_dict)
|
|
return updated_mcp_tool
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"MCP tool update error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.delete("/{mcp_tool_id}")
|
|
async def delete_mcp_tool(
|
|
mcp_tool_id: DbIntId,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Delete an MCP tool.
|
|
|
|
Deliberately not behind check_mcp_feature_access: an org whose plan no
|
|
longer carries mcp_tools must still be able to remove its connectors,
|
|
otherwise a downgrade would strand their credentials in env_vars for good.
|
|
"""
|
|
try:
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
# Check if user has access to this MCP tool's organization
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to delete this MCP tool"
|
|
)
|
|
|
|
success = mcp_tool_repo.delete_mcp_tool(mcp_tool_id)
|
|
if not success:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="Failed to delete MCP tool"
|
|
)
|
|
|
|
return {"message": "MCP tool deleted successfully"}
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"MCP tool deletion error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# Agent-MCPTool association endpoints
|
|
@router.post("/agent-association", response_model=MCPToolToAgentResponse)
|
|
async def add_mcp_tool_to_agent(
|
|
association_data: MCPToolToAgentCreate,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Add MCP tool to agent"""
|
|
try:
|
|
# Check MCP tools feature access
|
|
check_mcp_feature_access(current_user, db)
|
|
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
agent_repo = AgentRepository(db)
|
|
|
|
# Verify MCP tool exists and belongs to user's organization
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(association_data.mcp_tool_id)
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this MCP tool"
|
|
)
|
|
|
|
# Verify agent exists and belongs to user's organization
|
|
agent = agent_repo.get_agent(association_data.agent_id)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="Agent not found"
|
|
)
|
|
|
|
if agent.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this agent"
|
|
)
|
|
|
|
# Create association
|
|
association = mcp_tool_repo.add_mcp_tool_to_agent(
|
|
association_data.mcp_tool_id,
|
|
association_data.agent_id
|
|
)
|
|
|
|
return association
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error adding MCP tool to agent: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.delete("/agent-association/{mcp_tool_id}/{agent_id}")
|
|
async def remove_mcp_tool_from_agent(
|
|
mcp_tool_id: DbIntId,
|
|
agent_id: UUID,
|
|
current_user: User = Depends(require_permissions("manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Remove MCP tool from agent"""
|
|
try:
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
agent_repo = AgentRepository(db)
|
|
|
|
# Verify MCP tool exists and belongs to user's organization
|
|
mcp_tool = mcp_tool_repo.get_mcp_tool(mcp_tool_id)
|
|
if not mcp_tool:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="MCP tool not found"
|
|
)
|
|
|
|
if mcp_tool.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this MCP tool"
|
|
)
|
|
|
|
# Verify agent exists and belongs to user's organization
|
|
agent = agent_repo.get_agent(agent_id)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="Agent not found"
|
|
)
|
|
|
|
if agent.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this agent"
|
|
)
|
|
|
|
# Remove association
|
|
success = mcp_tool_repo.remove_mcp_tool_from_agent(mcp_tool_id, agent_id)
|
|
if not success:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="Association not found"
|
|
)
|
|
|
|
return {"message": "MCP tool removed from agent successfully"}
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error removing MCP tool from agent: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/agent/{agent_id}", response_model=AgentMCPToolsResponse)
|
|
async def get_agent_mcp_tools(
|
|
agent_id: UUID,
|
|
current_user: User = Depends(require_permissions("view_all", "manage_agents")),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""Get all MCP tools associated with an agent"""
|
|
try:
|
|
mcp_tool_repo = MCPToolRepository(db)
|
|
agent_repo = AgentRepository(db)
|
|
|
|
# Verify agent exists and belongs to user's organization
|
|
agent = agent_repo.get_agent(agent_id)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="Agent not found"
|
|
)
|
|
|
|
if agent.organization_id != current_user.organization_id:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Not authorized to access this agent"
|
|
)
|
|
|
|
# Get agent's MCP tools
|
|
mcp_tools = mcp_tool_repo.get_agent_mcp_tools(agent_id)
|
|
|
|
return AgentMCPToolsResponse(
|
|
id=agent.id,
|
|
name=agent.name,
|
|
mcp_tools=mcp_tools
|
|
)
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error fetching agent MCP tools: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e)) |