diff --git a/apps/backend/src/rhesis/backend/app/routers/services.py b/apps/backend/src/rhesis/backend/app/routers/services.py index ab57f165f..b9826d609 100644 --- a/apps/backend/src/rhesis/backend/app/routers/services.py +++ b/apps/backend/src/rhesis/backend/app/routers/services.py @@ -28,21 +28,28 @@ SearchMCPRequest, TestConfigRequest, TestConfigResponse, + TestMCPConnectionRequest, + TestMCPConnectionResponse, TextResponse, ) +from rhesis.backend.app.services.activities import RecentActivitiesService from rhesis.backend.app.services.gemini_client import ( create_chat_completion, get_chat_response, get_json_response, ) -from rhesis.backend.app.services.activities import RecentActivitiesService from rhesis.backend.app.services.generation import ( generate_multiturn_tests, generate_tests, ) from rhesis.backend.app.services.github import read_repo_contents from rhesis.backend.app.services.handlers import DocumentHandler -from rhesis.backend.app.services.mcp_service import extract_mcp, query_mcp, search_mcp +from rhesis.backend.app.services.mcp_service import ( + extract_mcp, + query_mcp, + search_mcp, + test_mcp_authentication, +) from rhesis.backend.app.services.storage_service import StorageService from rhesis.backend.app.services.test_config_generator import TestConfigGeneratorService from rhesis.backend.logging import logger @@ -456,7 +463,9 @@ async def extract_document_content(request: ExtractDocumentRequest) -> ExtractDo logger.error(f"Failed to extract document content: {str(e)}", exc_info=True) raise HTTPException( status_code=400, - detail="Failed to extract document content. Please check the file format and try again.", + detail=( + "Failed to extract document content. Please check the file format and try again." + ), ) @@ -666,6 +675,87 @@ async def query_mcp_server( ) +@router.post("/mcp/test-connection", response_model=TestMCPConnectionResponse) +async def test_mcp_connection( + request: TestMCPConnectionRequest, + db: Session = Depends(get_tenant_db_session), + tenant_context=Depends(get_tenant_context), + current_user: User = Depends(require_current_user_or_token), +): + """ + Test MCP connection authentication by calling a tool that requires authentication. + + This endpoint verifies that the MCP connection is properly authenticated by + instructing an agent to call a tool that requires authentication. If the + connection is not authenticated, a 401 unauthorized error will be detected. + + Args: + request: TestMCPConnectionRequest with either: + - tool_id: For testing existing tools + - provider_type_id + credentials + optional tool_metadata: + For testing non-existent tools + + Returns: + TestMCPConnectionResponse with: + - is_authenticated: str - "Yes" if authentication is valid, "No" otherwise + - message: str - Explanation of the authentication status + + Raises: + HTTPException: 400 error if validation fails, + 500 error if test fails due to connection issues + + Examples: + # Test existing tool + POST /mcp/test-connection + { + "tool_id": "tool-uuid-123" + } + + # Test non-existent tool (standard provider) + POST /mcp/test-connection + { + "provider_type_id": "provider-uuid-123", + "credentials": {"NOTION_TOKEN": "ntn_abc123..."} + } + + # Test non-existent tool (custom provider) + POST /mcp/test-connection + { + "provider_type_id": "provider-uuid-123", + "credentials": {"TOKEN": "token123"}, + "tool_metadata": { + "command": "bunx", + "args": ["--bun", "@notionhq/notion-mcp-server"], + "env": {"NOTION_TOKEN": "{{ TOKEN }}"} + } + } + """ + try: + organization_id, user_id = tenant_context + result = await test_mcp_authentication( + db=db, + user=current_user, + organization_id=organization_id, + tool_id=request.tool_id, + provider_type_id=request.provider_type_id, + credentials=request.credentials, + tool_metadata=request.tool_metadata, + user_id=user_id, + ) + return result + except ValueError as e: + logger.warning(f"Invalid request for MCP connection test: {str(e)}") + raise HTTPException(status_code=400, detail=str(e)) + except Exception as e: + logger.error(f"Failed to test MCP connection: {str(e)}", exc_info=True) + raise HTTPException( + status_code=500, + detail=( + "Failed to test MCP connection. Please verify the tool configuration and try again." + ), + ) + + @router.get("/recent-activities", response_model=RecentActivitiesResponse) def get_recent_activities( limit: int = Query(default=50, ge=1, le=200), diff --git a/apps/backend/src/rhesis/backend/app/schemas/services.py b/apps/backend/src/rhesis/backend/app/schemas/services.py index 46310710e..c74b8af6f 100644 --- a/apps/backend/src/rhesis/backend/app/schemas/services.py +++ b/apps/backend/src/rhesis/backend/app/schemas/services.py @@ -2,7 +2,7 @@ from enum import Enum from typing import Any, Dict, List, Optional -from pydantic import UUID4, BaseModel, Field +from pydantic import UUID4, BaseModel, Field, model_validator class GenerationConfig(BaseModel): @@ -299,6 +299,45 @@ class QueryMCPResponse(BaseModel): execution_history: List[ExecutionStep] +class TestMCPConnectionRequest(BaseModel): + """Request to test MCP connection authentication. + + Either tool_id (for existing tools) OR provider_type_id + credentials + (for non-existent tools) must be provided. + For custom providers, tool_metadata is required when using provider_type_id. + """ + + tool_id: Optional[str] = None + provider_type_id: Optional[UUID4] = None + credentials: Optional[Dict[str, str]] = None + tool_metadata: Optional[Dict[str, Any]] = None + + @model_validator(mode="after") + def validate_request(self): + """Ensure either tool_id OR (provider_type_id + credentials) is provided.""" + has_tool_id = self.tool_id is not None + has_params = self.provider_type_id is not None and self.credentials is not None + + if not has_tool_id and not has_params: + raise ValueError( + "Either 'tool_id' OR ('provider_type_id' + 'credentials') must be provided" + ) + + if has_tool_id and has_params: + raise ValueError( + "Cannot provide both 'tool_id' and parameter-based fields. Use one approach." + ) + + return self + + +class TestMCPConnectionResponse(BaseModel): + """Response from MCP connection authentication test.""" + + is_authenticated: str # "Yes" or "No" + message: str + + # Recent Activities Schemas class ActivityOperation(str, Enum): """Type of operation performed on an entity.""" @@ -339,3 +378,5 @@ class RecentActivitiesResponse(BaseModel): activities: List[ActivityItem] total: int # Total number of activity groups (not individual activities) + is_authenticated: str # "Yes" or "No" + message: str diff --git a/apps/backend/src/rhesis/backend/app/services/mcp_service.py b/apps/backend/src/rhesis/backend/app/services/mcp_service.py index d24c71e76..a5148c8d2 100644 --- a/apps/backend/src/rhesis/backend/app/services/mcp_service.py +++ b/apps/backend/src/rhesis/backend/app/services/mcp_service.py @@ -81,6 +81,64 @@ def _get_mcp_client_by_tool_id( return client +def _get_mcp_client_from_params( + provider_type_id: uuid.UUID, + credentials: Dict[str, str], + tool_metadata: Optional[Dict[str, Any]], + db: Session, + organization_id: str, + user_id: str = None, +): + """ + Get MCP client from parameters without requiring a tool in the database. + + Args: + provider_type_id: UUID of the provider type (TypeLookup) + credentials: Dictionary of credential key-value pairs + tool_metadata: Optional tool metadata (required for custom providers) + db: Database session + organization_id: Organization ID (for authorization check) + user_id: User ID (for authorization check) + + Returns: + MCPClient ready to use + + Raises: + ValueError: If provider not found, invalid configuration, or missing required fields + """ + # Fetch provider type from database + provider_type = crud.get_type_lookup(db, provider_type_id, organization_id, user_id) + + if not provider_type: + raise ValueError( + f"Provider type '{provider_type_id}' not found. Please verify the provider_type_id." + ) + + # Get provider name for the client + provider = provider_type.type_value + + # Check if provider uses custom provider (requires manual JSON config) or standard provider + if provider == "custom": + # Custom provider: requires tool_metadata with full JSON config + if not tool_metadata: + raise ValueError("Custom provider requires tool_metadata configuration") + + manager = MCPClientManager.from_tool_config( + tool_name=f"{provider}Api", + tool_config=tool_metadata, + credentials=credentials, + ) + else: + # Standard provider: SDK constructs config from YAML templates + manager = MCPClientManager.from_provider( + provider=provider, + credentials=credentials, + ) + + client = manager.create_client(f"{provider}Api") + return client + + async def search_mcp( query: str, tool_id: str, db: Session, user: User, organization_id: str, user_id: str = None ) -> List[Dict[str, str]]: @@ -254,3 +312,72 @@ async def query_mcp( raise ValueError(f"Query failed: {result.error}") return result.model_dump() + + +async def test_mcp_authentication( + db: Session, + user: User, + organization_id: str, + tool_id: Optional[str] = None, + provider_type_id: Optional[uuid.UUID] = None, + credentials: Optional[Dict[str, str]] = None, + tool_metadata: Optional[Dict[str, Any]] = None, + user_id: str = None, +) -> Dict[str, Any]: + """ + Test MCP connection authentication by calling a tool that requires authentication. + + Uses an AI agent to call a tool requiring authentication. The LLM determines + whether authentication is successful based on the tool call results. + + Args: + db: Database session + user: Current user (for retrieving default generation model) + organization_id: Organization ID for loading tools from database + tool_id: Optional ID of the configured tool instance (for existing tools) + provider_type_id: Optional UUID of the provider type (for non-existent tools) + credentials: Optional dictionary of credential key-value pairs (for non-existent tools) + tool_metadata: Optional tool metadata (for non-existent tools that use custom providers) + user_id: User ID for authorization check + + Returns: + Dict with: + - is_authenticated: str - "Yes" or "No" determined by the LLM based on tool call results + - message: str - Message from the LLM explaining the authentication status + + Raises: + ValueError: If authentication test fails due to connection issues + """ + model = get_user_generation_model(db, user) + + # Load MCP client from either tool_id or parameters + if tool_id is not None: + client = _get_mcp_client_by_tool_id(db, tool_id, organization_id, user_id) + else: + client = _get_mcp_client_from_params( + provider_type_id=provider_type_id, + credentials=credentials, + tool_metadata=tool_metadata, + db=db, + organization_id=organization_id, + user_id=user_id, + ) + + # Load the authentication test prompt + auth_prompt = jinja_env.get_template("mcp_test_auth_prompt.jinja2").render() + agent = MCPAgent( + model=model, + mcp_client=client, + system_prompt=auth_prompt, + max_iterations=5, # Keep it short for authentication test + verbose=False, + ) + + # Run agent with query to test authentication + query = "Call a tool that requires authentication to verify your credentials" + result = await agent.run_async(query) + + if not result.success: + raise ValueError(f"Authentication test failed: {result.error}") + + return json.loads(result.final_answer) diff --git a/apps/backend/src/rhesis/backend/app/templates/mcp_test_auth_prompt.jinja2 b/apps/backend/src/rhesis/backend/app/templates/mcp_test_auth_prompt.jinja2 new file mode 100644 index 000000000..bbcb28245 --- /dev/null +++ b/apps/backend/src/rhesis/backend/app/templates/mcp_test_auth_prompt.jinja2 @@ -0,0 +1,19 @@ +You are testing the authentication status of an MCP connection. + +Your task is to call a tool that requires authentication to verify that the connection is properly authenticated. + +Look at the available tools and select one that would require authentication, such as: +- Getting user information +- Listing protected resources +- Accessing user-specific data +- Any tool that would fail with a 401 unauthorized error if not authenticated + +Call the tool and then provide your final answer in this exact JSON format: +{ + "is_authenticated": "Yes" or "No", + "message": "Brief explanation of the authentication status" +} + +Use "Yes" if the tool call succeeded, or "No" if you received any authentication error (such as 401 unauthorized, authentication failed, etc.). + +Be concise and direct - your goal is to verify authentication status, not to perform a complex task. diff --git a/apps/frontend/src/app/(protected)/knowledge/components/MCPImportDialog.tsx b/apps/frontend/src/app/(protected)/knowledge/components/MCPImportDialog.tsx index 8ddc9d10b..aeb64c351 100644 --- a/apps/frontend/src/app/(protected)/knowledge/components/MCPImportDialog.tsx +++ b/apps/frontend/src/app/(protected)/knowledge/components/MCPImportDialog.tsx @@ -126,10 +126,6 @@ export default function MCPImportDialog({ ? err.message : 'Failed to search. Please try again.'; setError(errorMessage); - notifications.show('Search failed: ' + errorMessage, { - severity: 'error', - autoHideDuration: 6000, - }); } finally { setSearching(false); } diff --git a/apps/frontend/src/app/(protected)/mcp/components/MCPConnectionDialog.tsx b/apps/frontend/src/app/(protected)/mcp/components/MCPConnectionDialog.tsx index fac5a2aa7..65b899357 100644 --- a/apps/frontend/src/app/(protected)/mcp/components/MCPConnectionDialog.tsx +++ b/apps/frontend/src/app/(protected)/mcp/components/MCPConnectionDialog.tsx @@ -19,8 +19,12 @@ import VisibilityIcon from '@mui/icons-material/Visibility'; import VisibilityOffIcon from '@mui/icons-material/VisibilityOff'; import ExpandMoreIcon from '@mui/icons-material/ExpandMore'; import ExpandLessIcon from '@mui/icons-material/ExpandLess'; +import CheckCircleIcon from '@mui/icons-material/CheckCircle'; +import ErrorIcon from '@mui/icons-material/Error'; import dynamic from 'next/dynamic'; import { useTheme } from '@mui/material/styles'; +import { useSession } from 'next-auth/react'; +import { ApiClientFactory } from '@/utils/api-client/client-factory'; import { TypeLookup } from '@/utils/api-client/interfaces/type-lookup'; import { Tool, @@ -29,6 +33,7 @@ import { } from '@/utils/api-client/interfaces/tool'; import { UUID } from 'crypto'; import { MCP_PROVIDER_ICONS } from '@/config/mcp-providers'; +import { useNotifications } from '@/components/common/NotificationContext'; // Lazy load Monaco Editor const Editor = dynamic(() => import('@monaco-editor/react'), { @@ -97,6 +102,8 @@ export function MCPConnectionDialog({ onUpdate, }: MCPConnectionDialogProps) { const theme = useTheme(); + const { data: session } = useSession(); + const notifications = useNotifications(); const [name, setName] = useState(''); const [description, setDescription] = useState(''); const [authToken, setAuthToken] = useState(''); @@ -106,6 +113,12 @@ export function MCPConnectionDialog({ const [showAuthToken, setShowAuthToken] = useState(false); const [showAdvancedConfig, setShowAdvancedConfig] = useState(false); const [jsonError, setJsonError] = useState(null); + const [testingConnection, setTestingConnection] = useState(false); + const [testResult, setTestResult] = useState<{ + is_authenticated: string; + message: string; + } | null>(null); + const [connectionTested, setConnectionTested] = useState(false); const isEditMode = mode === 'edit'; @@ -154,6 +167,8 @@ export function MCPConnectionDialog({ setShowAuthToken(false); setLoading(false); setShowAdvancedConfig(!!tool.tool_metadata); + setTestResult(null); + setConnectionTested(true); // Skip test requirement in edit mode } else if (provider) { // Create mode: reset to defaults setName(''); @@ -165,10 +180,25 @@ export function MCPConnectionDialog({ setShowAuthToken(false); setLoading(false); setShowAdvancedConfig(isCustomProvider); + setTestResult(null); + setConnectionTested(false); } } }, [open, provider, tool, isEditMode, isCustomProvider]); + // Reset connection test status when critical fields change + useEffect(() => { + if (!isEditMode) { + // In create mode, always reset when fields change + setConnectionTested(false); + setTestResult(null); + } else if (authToken && authToken !== '************') { + // In edit mode, reset only if auth token was actually changed + setConnectionTested(false); + setTestResult(null); + } + }, [name, authToken, toolMetadata, provider, isEditMode]); + const validateToolMetadata = ( jsonString: string ): Record | null => { @@ -202,6 +232,108 @@ export function MCPConnectionDialog({ } }; + const handleTestConnection = async () => { + if (!session?.session_token) { + setError('Session not available. Please try again.'); + return; + } + + // In create mode, we need to validate required fields first + if (!isEditMode) { + if (!provider || !name || (requiresToken && !authToken)) { + setError('Please fill in all required fields before testing.'); + return; + } + if (isCustomProvider && !toolMetadata.trim()) { + setError('Tool metadata is required for custom providers.'); + return; + } + } + + // In edit mode, we need the tool ID + if (isEditMode && !tool?.id) { + setError('Tool ID not available. Please save the connection first.'); + return; + } + + setTestingConnection(true); + setError(null); + setTestResult(null); + + try { + const apiFactory = new ApiClientFactory(session.session_token); + const servicesClient = apiFactory.getServicesClient(); + + let testRequest: { + tool_id?: string; + provider_type_id?: string; + credentials?: Record; + tool_metadata?: Record; + }; + + if (isEditMode) { + // In edit mode, use existing tool ID + testRequest = { + tool_id: tool!.id, + }; + } else { + // In create mode, use direct parameters + if (!provider) { + setError('Provider not found. Please try again.'); + setTestingConnection(false); + return; + } + + const credentialKey = getCredentialKey(provider.type_value); + const credentials = requiresToken + ? { + [credentialKey]: authToken.trim(), + } + : {}; + + let parsedMetadata: Record | undefined = undefined; + if (isCustomProvider && toolMetadata.trim()) { + const validatedMetadata = validateToolMetadata(toolMetadata); + if (validatedMetadata === null) { + setError( + 'Please fix the JSON configuration errors before testing.' + ); + setTestingConnection(false); + return; + } + parsedMetadata = validatedMetadata; + } + + testRequest = { + provider_type_id: provider.id, + credentials, + tool_metadata: parsedMetadata, + }; + } + + // Test the connection + const result = await servicesClient.testMCPConnection(testRequest); + setTestResult(result); + + // Mark as tested if successful + if (result.is_authenticated === 'Yes') { + setConnectionTested(true); + } else { + setConnectionTested(false); + } + } catch (err) { + setError( + err instanceof Error + ? err.message + : 'Failed to test connection. Please try again.' + ); + setTestResult(null); + setConnectionTested(false); + } finally { + setTestingConnection(false); + } + }; + const handleSubmit = async (e: React.FormEvent) => { e.preventDefault(); @@ -263,7 +395,19 @@ export function MCPConnectionDialog({ setLoading(false); } } else { - // Create mode: validate required fields + // Create mode: require connection test + if (!connectionTested) { + setError('Please test the connection before saving the tool.'); + return; + } + + // Validate that test was successful + if (!testResult || testResult.is_authenticated !== 'Yes') { + setError('Connection test must be successful before saving.'); + return; + } + + // Validate required fields if (!provider || !name || (requiresToken && !authToken)) { setError('Please fill in all required fields.'); return; @@ -456,6 +600,49 @@ export function MCPConnectionDialog({ ) : null, }} /> + + + {testResult && ( + + ) : ( + + ) + } + sx={{ mt: 2 }} + > + + {testResult.is_authenticated === 'Yes' + ? 'Connection Successful' + : 'Connection Failed'} + + + {testResult.message} + + + )} + )} @@ -572,6 +759,27 @@ export function MCPConnectionDialog({ + {/* Connection Test Required Message */} + {!isEditMode && !connectionTested && !testResult && ( + + + Please test the connection before saving the tool configuration. + + + )} + {isEditMode && + requiresToken && + authToken !== '************' && + !connectionTested && + !testResult && ( + + + Please test the connection with the new credentials before + updating. + + + )} + @@ -585,6 +793,11 @@ export function MCPConnectionDialog({ !name || (!isEditMode && requiresToken && !authToken) || (!isEditMode && isCustomProvider && !toolMetadata.trim()) || + (!isEditMode && !connectionTested) || + (isEditMode && + requiresToken && + authToken !== '************' && + !connectionTested) || // Require test if token changed in edit mode loading || !!jsonError } diff --git a/apps/frontend/src/utils/api-client/services-client.ts b/apps/frontend/src/utils/api-client/services-client.ts index 3466aca75..ad966bee0 100644 --- a/apps/frontend/src/utils/api-client/services-client.ts +++ b/apps/frontend/src/utils/api-client/services-client.ts @@ -116,6 +116,18 @@ export interface MCPExtractResponse { content: string; } +export interface TestMCPConnectionRequest { + tool_id?: string; + provider_type_id?: string; + credentials?: Record; + tool_metadata?: Record; +} + +export interface TestMCPConnectionResponse { + is_authenticated: string; // "Yes" or "No" + message: string; +} + export class ServicesClient extends BaseApiClient { async getGitHubContents(repo_url: string): Promise { return this.fetch( @@ -284,6 +296,26 @@ export class ServicesClient extends BaseApiClient { ); } + /** + * Test MCP connection authentication + * @param request - Either tool_id (for existing tools) OR provider_type_id + credentials + optional tool_metadata (for non-existent tools) + * @returns Test result with authentication status and message + */ + async testMCPConnection( + request: TestMCPConnectionRequest + ): Promise { + return this.fetch( + `${API_ENDPOINTS.services}/mcp/test-connection`, + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(request), + } + ); + } + /** * Get recent activities across all trackable entities * @param limit - Maximum number of activities to return (default 50, max 200)