[MCP Gateway] - Add custom cost configuration for each MCP tool (#12499)

* add endpoints to test MCP connection * fix route names * add testMCPToolsListRequest * fix MCP connection test * add tool_name_to_cost_per_query * fixes tool_name_to_cost_per_query * fix networking * fix test MCP connection * use POST for test tools endpoints * ui fixes * fix config * fixes for cost config * fixes * decent connection status * fix Created At * fix MCP table * Potential fix for code scanning alert no. 2928: Information exposure through an exception Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * MCPServerCostConfig --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
2026-08-17 04:26:20 +00:00 · 2025-07-10 17:00:39 -07:00
parent 0691ff8d13
commit 5fdbb218a8
10 changed files with 817 additions and 28 deletions
@@ -149,3 +149,72 @@ if MCP_AVAILABLE:
            proxy_config=proxy_config,
        )
        return await call_mcp_tool(**data)
+    
+    ########################################################
+    # MCP Connection testing routes
+    # /health -> Test if we can connect to the MCP server
+    # /health/tools/list -> List tools from MCP server
+    # For these routes users will dynamically pass the MCP connection params, they don't need to be on the MCP registry
+    ########################################################
+    from litellm.proxy._experimental.mcp_server.server import MCPServer
+    from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+        NewMCPServerRequest,
+    )
+    @router.post("/test/connection")
+    async def test_connection(
+        request: NewMCPServerRequest,
+    ):
+        """
+        Test if we can connect to the provided MCP server before adding it
+        """
+        try:
+            client = global_mcp_server_manager._create_mcp_client(
+                server=MCPServer(
+                    server_id=request.server_id or "",
+                    name=request.alias or "",
+                    url=request.url,
+                    transport=request.transport,
+                    spec_version=request.spec_version,
+                    auth_type=request.auth_type,
+                    mcp_info=request.mcp_info,
+                ),
+                mcp_auth_header=None,
+            )
+
+            await client.connect()
+        except Exception as e:
+            verbose_logger.error(f"Error in test_connection: {e}", exc_info=True)
+            return {"status": "error", "message": "An internal error has occurred."}
+        return {"status": "ok"}
+        
+    
+    @router.post("/test/tools/list")
+    async def test_tools_list(
+        request: NewMCPServerRequest,
+        user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    ):
+        """
+        Preview tools available from MCP server before adding it
+        """
+        try:
+            client = global_mcp_server_manager._create_mcp_client(
+                server=MCPServer(
+                    server_id=request.server_id or "",
+                    name=request.alias or "",
+                    url=request.url,
+                    transport=request.transport,
+                    spec_version=request.spec_version,
+                    auth_type=request.auth_type,
+                    mcp_info=request.mcp_info,
+                ),
+                mcp_auth_header=None,
+            )
+            list_tools_result = await client.list_tools()
+        except Exception as e:
+            verbose_logger.error(f"Error in test_tools_list: {e}", exc_info=True)
+            return {"status": "error", "message": "An internal error has occurred."}
+        return {
+            "tools": list_tools_result,
+            "error": None,
+            "message": "Successfully retrieved tools"
+        }
@@ -6,9 +6,6 @@ model_list:
    litellm_params:
      model: openai/*

-general_settings:
-  custom_ui_sso_sign_in_handler: "custom_hooks.custom_ui_sso_hook.custom_ui_sso_sign_in_handler"
-
 guardrails:
  - guardrail_name: "bedrock-pre-guard"
    litellm_params:
@@ -0,0 +1,254 @@
+import React from 'react';
+import { Typography, Space, Button, Divider, message } from 'antd';
+import { WarningOutlined, InfoCircleOutlined, CopyOutlined } from '@ant-design/icons';
+import { testMCPConnectionRequest } from "./networking";
+
+const { Text } = Typography;
+
+interface MCPConnectionTestProps {
+  formValues: Record<string, any>;
+  accessToken: string;
+  serverName?: string;
+  onClose?: () => void;
+  onTestComplete?: () => void;
+}
+
+const MCPConnectionTest: React.FC<MCPConnectionTestProps> = ({ 
+  formValues, 
+  accessToken, 
+  serverName = "this MCP server", 
+  onClose,
+  onTestComplete
+}) => {
+  const [connectionError, setConnectionError] = React.useState<Error | string | null>(null);
+  const [rawRequest, setRawRequest] = React.useState<any>(null);
+  const [rawResponse, setRawResponse] = React.useState<any>(null);
+  const [isLoading, setIsLoading] = React.useState<boolean>(true);
+  const [connectionSuccess, setConnectionSuccess] = React.useState<boolean>(false);
+  const [showDetails, setShowDetails] = React.useState<boolean>(false);
+
+  const testMCPConnection = async () => {
+    setIsLoading(true);
+    setShowDetails(false);
+    setConnectionError(null);
+    setRawRequest(null);
+    setRawResponse(null);
+    setConnectionSuccess(false);
+    
+    // Add a small delay to ensure form values are fully populated
+    await new Promise(resolve => setTimeout(resolve, 100));
+    
+    try {
+      console.log("Testing MCP connection with form values:", formValues);
+      
+      // Prepare the MCP server config from form values
+      const mcpServerConfig = {
+        server_id: formValues.server_id || "",
+        alias: formValues.alias || "",
+        url: formValues.url,
+        transport: formValues.transport,
+        spec_version: formValues.spec_version,
+        auth_type: formValues.auth_type,
+        mcp_info: formValues.mcp_info,
+      };
+
+      setRawRequest(mcpServerConfig);
+
+      // Test connection
+      const connectionResponse = await testMCPConnectionRequest(accessToken, mcpServerConfig);
+      console.log("Connection test response:", connectionResponse);
+      
+      if (connectionResponse.status === "ok") {
+        setConnectionError(null);
+        setConnectionSuccess(true);
+      } else {
+        const errorMessage = connectionResponse.message || "Unknown connection error";
+        setConnectionError(errorMessage);
+        setRawResponse(connectionResponse);
+      }
+    } catch (error) {
+      console.error("MCP connection test error:", error);
+      setConnectionError(error instanceof Error ? error.message : String(error));
+    } finally {
+      setIsLoading(false);
+      if (onTestComplete) onTestComplete();
+    }
+  };
+
+  React.useEffect(() => {
+    // Run the test once when component mounts
+    // Add a small timeout to ensure form values are ready
+    const timer = setTimeout(() => {
+      testMCPConnection();
+    }, 200);
+    
+    return () => clearTimeout(timer);
+  }, []); // Empty dependency array means this runs once on mount
+
+  const getCleanErrorMessage = (errorMsg: string) => {
+    if (!errorMsg) return "Unknown error";
+    
+    const mainError = errorMsg.split('stack trace:')[0].trim();
+    
+    const cleanedError = mainError.replace(/^(.*?)Error: /, '');
+    
+    return cleanedError;
+  };
+
+  const connectionErrorMessage = typeof connectionError === 'string' 
+    ? getCleanErrorMessage(connectionError) 
+    : connectionError?.message ? getCleanErrorMessage(connectionError.message) : "Unknown error";
+
+  const formatMCPRequest = (mcpConfig: Record<string, any>) => {
+    return JSON.stringify(mcpConfig, null, 2);
+  };
+
+  const isOverallSuccess = connectionSuccess && !connectionError;
+
+  return (
+    <div style={{ padding: '24px', borderRadius: '8px', backgroundColor: '#fff' }}>
+      {isLoading ? (
+        <div style={{ textAlign: 'center', padding: '32px 20px' }}>
+          <div className="loading-spinner" style={{ marginBottom: '16px' }}>
+            {/* Simple CSS spinner */}
+            <div style={{ 
+              border: '3px solid #f3f3f3',
+              borderTop: '3px solid #1890ff',
+              borderRadius: '50%',
+              width: '30px',
+              height: '30px',
+              animation: 'spin 1s linear infinite',
+              margin: '0 auto'
+            }} />
+          </div>
+          <Text style={{ fontSize: '16px' }}>Testing connection to {serverName}...</Text>
+          <style jsx>{`
+            @keyframes spin {
+              0% { transform: rotate(0deg); }
+              100% { transform: rotate(360deg); }
+            }
+          `}</style>
+        </div>
+      ) : isOverallSuccess ? (
+        <div>
+          <div style={{ display: 'flex', alignItems: 'center', justifyContent: 'center', padding: '32px 20px' }}>
+            <div style={{ color: '#52c41a', fontSize: '24px', display: 'flex', alignItems: 'center' }}>
+              <svg viewBox="64 64 896 896" focusable="false" data-icon="check-circle" width="1em" height="1em" fill="currentColor" aria-hidden="true">
+                <path d="M512 64C264.6 64 64 264.6 64 512s200.6 448 448 448 448-200.6 448-448S759.4 64 512 64zm193.5 301.7l-210.6 292a31.8 31.8 0 01-51.7 0L318.5 484.9c-3.8-5.3 0-12.7 6.5-12.7h46.9c10.2 0 19.9 4.9 25.9 13.3l71.2 98.8 157.2-218c6-8.3 15.6-13.3 25.9-13.3H699c6.5 0 10.3 7.4 6.5 12.7z"></path>
+              </svg>
+            </div>
+            <Text type="success" style={{ fontSize: '18px', fontWeight: 500, marginLeft: '10px' }}>
+              Connection to {serverName} successful!
+            </Text>
+          </div>
+          
+
+        </div>
+      ) : (
+        <>
+          <div>
+            <div style={{ display: 'flex', alignItems: 'center', marginBottom: '20px' }}>
+              <WarningOutlined style={{ color: '#ff4d4f', fontSize: '24px', marginRight: '12px' }} />
+              <Text type="danger" style={{ fontSize: '18px', fontWeight: 500 }}>Connection to {serverName} failed</Text>
+            </div>
+            
+            <div style={{ 
+              backgroundColor: '#fff2f0', 
+              border: '1px solid #ffccc7', 
+              borderRadius: '8px', 
+              padding: '16px', 
+              marginBottom: '20px',
+              boxShadow: '0 1px 2px rgba(0, 0, 0, 0.03)'
+            }}>
+              <Text strong style={{ display: 'block', marginBottom: '8px' }}>Error: </Text>
+              <Text type="danger" style={{ fontSize: '14px', lineHeight: '1.5' }}>{connectionErrorMessage}</Text>
+              
+              {connectionError && (
+                <div style={{ marginTop: '12px' }}>
+                  <Button 
+                    type="link" 
+                    onClick={() => setShowDetails(!showDetails)}
+                    style={{ paddingLeft: 0, height: 'auto' }}
+                  >
+                    {showDetails ? 'Hide Details' : 'Show Details'}
+                  </Button>
+                </div>
+              )}
+            </div>
+            
+            {showDetails && (
+              <div style={{ marginBottom: '20px' }}>
+                <Text strong style={{ display: 'block', marginBottom: '8px', fontSize: '15px' }}>Troubleshooting Details</Text>
+                <pre style={{ 
+                  backgroundColor: '#f5f5f5', 
+                  padding: '16px', 
+                  borderRadius: '8px',
+                  fontSize: '13px',
+                  maxHeight: '200px',
+                  overflow: 'auto',
+                  border: '1px solid #e8e8e8',
+                  lineHeight: '1.5'
+                }}>
+                  {typeof connectionError === 'string' ? connectionError : JSON.stringify(connectionError, null, 2)}
+                </pre>
+              </div>
+            )}
+            
+            <div>
+              <Text strong style={{ display: 'block', marginBottom: '8px', fontSize: '15px' }}>MCP Server Configuration</Text>
+              <pre style={{ 
+                backgroundColor: '#f5f5f5', 
+                padding: '16px', 
+                borderRadius: '8px',
+                fontSize: '13px',
+                maxHeight: '250px',
+                overflow: 'auto',
+                border: '1px solid #e8e8e8',
+                lineHeight: '1.5'
+              }}>
+                {formatMCPRequest(rawRequest || {})}
+              </pre>
+              <Button 
+                style={{ marginTop: '8px' }}
+                icon={<CopyOutlined />} 
+                onClick={() => {
+                  navigator.clipboard.writeText(formatMCPRequest(rawRequest || {}));
+                  message.success('Copied to clipboard');
+                }}
+              >
+                Copy Configuration
+              </Button>
+            </div>
+          </div>
+        </>
+      )}
+      <Divider style={{ margin: '24px 0 16px' }} />
+      <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
+        <Button 
+          type="link" 
+          href="https://docs.litellm.ai/docs/proxy/mcp_server" 
+          target="_blank"
+          icon={<InfoCircleOutlined />}
+        >
+          View MCP Documentation
+        </Button>
+        <Space>
+          <Button 
+            onClick={testMCPConnection}
+            loading={isLoading}
+          >
+            Test Again
+          </Button>
+          <Button 
+            type="primary"
+            onClick={onClose}
+          >
+            Close
+          </Button>
+        </Space>
+      </div>
+    </div>
+  );
+};
+
+export default MCPConnectionTest; 
@@ -12,6 +12,7 @@ import { Button, TextInput } from "@tremor/react";
 import { createMCPServer } from "../networking";
 import { MCPServer, MCPServerCostInfo } from "./types";
 import MCPServerCostConfig from "./mcp_server_cost_config";
+import MCPConnectionStatus from "./mcp_connection_status";
 import { isAdminRole } from "@/utils/roles";


@@ -32,6 +33,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
  const [form] = Form.useForm();
  const [isLoading, setIsLoading] = useState(false);
  const [costConfig, setCostConfig] = useState<MCPServerCostInfo>({});
+  const [formValues, setFormValues] = useState<Record<string, any>>({});
+  const [tools, setTools] = useState<any[]>([]);

  const handleCreate = async (formValues: Record<string, any>) => {
    setIsLoading(true);
@@ -57,6 +60,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
        message.success("MCP Server created successfully");
        form.resetFields();
        setCostConfig({});
+        setTools([]);
        setModalVisible(false);
        onCreateSuccess(response);
      }
@@ -73,9 +77,12 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
  const handleCancel = () => {
    form.resetFields();
    setCostConfig({});
+    setTools([]);
    setModalVisible(false);
  };

+
+
  // rendering
  if (!isAdminRole(userRole)) {
    return null;
@@ -121,6 +128,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
          <Form
            form={form}
            onFinish={handleCreate}
+            onValuesChange={(_, allValues) => setFormValues(allValues)}
            layout="vertical"
            className="space-y-6"
          >
@@ -194,8 +202,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
                    className="rounded-lg"
                    size="large"
                  >
-                    <Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
                    <Select.Option value="http">HTTP</Select.Option>
+                    <Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
                  </Select>
                </Form.Item>

@@ -246,12 +254,21 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
              </Form.Item>
            </div>

-            {/* Cost Configuration Section */}
+            {/* Connection Status Section */}
            <div className="mt-8 pt-6 border-t border-gray-200">
+              <MCPConnectionStatus
+                accessToken={accessToken}
+                formValues={formValues}
+                onToolsLoaded={setTools}
+              />
+            </div>
+
+            {/* Cost Configuration Section */}
+            <div className="mt-6">
              <MCPServerCostConfig
                value={costConfig}
                onChange={setCostConfig}
-                accessToken={accessToken}
+                tools={tools}
                disabled={false}
              />
            </div>
@@ -273,6 +290,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
          </Form>
        </div>
      </Modal>
+
+
    </div>
  );
 };
@@ -0,0 +1,230 @@
+import React, { useState, useEffect } from "react";
+import { Button, message, Spin, Alert, Collapse, Badge } from "antd";
+import { CheckCircleOutlined, ExclamationCircleOutlined, ReloadOutlined, ToolOutlined, InfoCircleOutlined } from "@ant-design/icons";
+import { Card, Title, Text } from "@tremor/react";
+import { testMCPToolsListRequest } from "../networking";
+
+const { Panel } = Collapse;
+
+interface MCPConnectionStatusProps {
+  accessToken: string | null;
+  formValues: Record<string, any>;
+  onToolsLoaded?: (tools: any[]) => void;
+}
+
+const MCPConnectionStatus: React.FC<MCPConnectionStatusProps> = ({
+  accessToken,
+  formValues,
+  onToolsLoaded
+}) => {
+  const [tools, setTools] = useState<any[]>([]);
+  const [isLoadingTools, setIsLoadingTools] = useState(false);
+  const [toolsError, setToolsError] = useState<string | null>(null);
+  const [hasShownSuccessMessage, setHasShownSuccessMessage] = useState(false);
+
+  // Check if we have the minimum required fields to fetch tools
+  const canFetchTools = formValues.url && formValues.transport && formValues.auth_type && accessToken;
+
+  const fetchTools = async () => {
+    if (!accessToken || !formValues.url) {
+      return;
+    }
+
+    setIsLoadingTools(true);
+    setToolsError(null);
+    
+    try {
+      // Prepare the MCP server config from form values
+      const mcpServerConfig = {
+        server_id: formValues.server_id || "",
+        alias: formValues.alias || "",
+        url: formValues.url,
+        transport: formValues.transport,
+        spec_version: formValues.spec_version,
+        auth_type: formValues.auth_type,
+        mcp_info: formValues.mcp_info,
+      };
+
+      const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig);
+      
+      if (toolsResponse.tools && !toolsResponse.error) {
+        setTools(toolsResponse.tools);
+        setToolsError(null);
+        onToolsLoaded?.(toolsResponse.tools);
+        if (toolsResponse.tools.length > 0 && !hasShownSuccessMessage) {
+          setHasShownSuccessMessage(true);
+        }
+      } else {
+        const errorMessage = toolsResponse.message || "Failed to retrieve tools list";
+        setToolsError(errorMessage);
+        setTools([]);
+        onToolsLoaded?.([]);
+        setHasShownSuccessMessage(false);
+      }
+    } catch (error) {
+      console.error("Tools fetch error:", error);
+      setToolsError(error instanceof Error ? error.message : String(error));
+      setTools([]);
+      onToolsLoaded?.([]);
+      setHasShownSuccessMessage(false);
+    } finally {
+      setIsLoadingTools(false);
+    }
+  };
+
+  // Auto-fetch tools when form values change and required fields are available
+  useEffect(() => {
+    if (canFetchTools) {
+      fetchTools();
+    } else {
+      // Clear tools if required fields are missing
+      setTools([]);
+      setToolsError(null);
+      setHasShownSuccessMessage(false);
+      onToolsLoaded?.([]);
+    }
+  }, [formValues.url, formValues.transport, formValues.auth_type, formValues.spec_version, accessToken]);
+
+  // Don't show anything if required fields aren't filled
+  if (!canFetchTools && !formValues.url) {
+    return null;
+  }
+
+  return (
+    <Card>
+      <div className="space-y-4">
+        <div className="flex items-center gap-2">
+          <CheckCircleOutlined className="text-blue-600" />
+          <Title>Connection Status</Title>
+        </div>
+
+        {!canFetchTools && formValues.url && (
+          <div className="text-center py-6 text-gray-400 border rounded-lg border-dashed">
+            <ToolOutlined className="text-2xl mb-2" />
+            <Text>Complete required fields to test connection</Text>
+            <br />
+            <Text className="text-sm">
+              Fill in URL, Transport, and Authentication to test MCP server connection
+            </Text>
+          </div>
+        )}
+
+        {canFetchTools && (
+          <div>
+            <div className="flex items-center justify-between mb-4">
+              <div>
+                <Text className="text-gray-700 font-medium">
+                  {isLoadingTools 
+                    ? "Testing connection to MCP server..." 
+                    : tools.length > 0 
+                      ? "Connection successful"
+                      : toolsError
+                        ? "Connection failed"
+                        : "Ready to test connection"}
+                </Text>
+                <br />
+                <Text className="text-gray-500 text-sm">
+                  Server: {formValues.url}
+                </Text>
+              </div>
+              
+              {isLoadingTools && (
+                <div className="flex items-center text-blue-600">
+                  <Spin size="small" className="mr-2" />
+                  <Text className="text-blue-600">Connecting...</Text>
+                </div>
+              )}
+              
+              {!isLoadingTools && !toolsError && tools.length > 0 && (
+                <div className="flex items-center text-green-600">
+                  <CheckCircleOutlined className="mr-1" />
+                  <Text className="text-green-600 font-medium">Connected</Text>
+                </div>
+              )}
+              
+              {toolsError && (
+                <div className="flex items-center text-red-600">
+                  <ExclamationCircleOutlined className="mr-1" />
+                  <Text className="text-red-600 font-medium">Failed</Text>
+                </div>
+              )}
+            </div>
+
+            {isLoadingTools && (
+              <div className="flex items-center justify-center py-6">
+                <Spin size="large" />
+                <Text className="ml-3">Testing connection and loading tools...</Text>
+              </div>
+            )}
+
+            {toolsError && (
+              <Alert
+                message="Connection Failed"
+                description={toolsError}
+                type="error"
+                showIcon
+                action={
+                  <Button
+                    icon={<ReloadOutlined />}
+                    onClick={fetchTools}
+                    size="small"
+                  >
+                    Retry
+                  </Button>
+                }
+              />
+            )}
+
+            {!isLoadingTools && tools.length > 0 && (
+              <Collapse
+                items={[
+                  {
+                    key: '1',
+                    label: (
+                      <div className="flex items-center">
+                        <ToolOutlined className="mr-2 text-green-500" />
+                        <span className="font-medium">Available Tools</span>
+                        <Badge 
+                          count={tools.length} 
+                          style={{ 
+                            backgroundColor: '#52c41a', 
+                            marginLeft: '8px' 
+                          }} 
+                        />
+                      </div>
+                    ),
+                    children: (
+                      <div className="space-y-2 max-h-48 overflow-y-auto">
+                        {tools.map((tool, index) => (
+                          <div key={index} className="p-3 bg-gray-50 rounded-lg">
+                            <Text className="font-medium text-gray-900">{tool.name}</Text>
+                            {tool.description && (
+                              <Text className="text-gray-500 text-sm block mt-1">
+                                {tool.description}
+                              </Text>
+                            )}
+                          </div>
+                        ))}
+                      </div>
+                    ),
+                  },
+                ]}
+              />
+            )}
+
+            {!isLoadingTools && tools.length === 0 && !toolsError && (
+              <div className="text-center py-6 text-gray-500 border rounded-lg border-dashed">
+                <CheckCircleOutlined className="text-2xl mb-2 text-green-500" />
+                <Text className="text-green-600 font-medium">Connection successful!</Text>
+                <br />
+                <Text className="text-gray-500">No tools found for this MCP server</Text>
+              </div>
+            )}
+          </div>
+        )}
+      </div>
+    </Card>
+  );
+};
+
+export default MCPConnectionStatus; 
@@ -56,9 +56,35 @@ export const mcpServerColumns = (
      </span>
    ),
  },
+  {
+    header: "Created At",
+    accessorKey: "created_at",
+    sortingFn: "datetime",
+    cell: ({ row }) => {
+      const server = row.original;
+      return (
+        <span className="text-xs">
+          {server.created_at ? new Date(server.created_at).toLocaleDateString() : "-"}
+        </span>
+      );
+    },
+  },
+  {
+    header: "Updated At",
+    accessorKey: "updated_at",
+    sortingFn: "datetime",
+    cell: ({ row }) => {
+      const server = row.original;
+      return (
+        <span className="text-xs">
+          {server.updated_at ? new Date(server.updated_at).toLocaleDateString() : "-"}
+        </span>
+      );
+    },
+  },
  {
    id: "actions",
-    header: "Info",
+    header: "Actions",
    cell: ({ row }) => (
      <div className="flex items-center gap-2">
        <Icon
@@ -1,24 +1,20 @@
 import React from "react";
-import { Tooltip, InputNumber } from "antd";
-import { InfoCircleOutlined, DollarOutlined } from "@ant-design/icons";
+import { Tooltip, InputNumber, Collapse, Badge } from "antd";
+import { InfoCircleOutlined, DollarOutlined, ToolOutlined } from "@ant-design/icons";
 import { Card, Title, Text } from "@tremor/react";
 import { MCPServerCostInfo } from "./types";

 interface MCPServerCostConfigProps {
  value?: MCPServerCostInfo;
  onChange?: (value: MCPServerCostInfo) => void;
-  serverId?: string;
-  serverUrl?: string;
-  accessToken: string | null;
+  tools?: any[]; // Receive tools from connection component
  disabled?: boolean;
 }

 const MCPServerCostConfig: React.FC<MCPServerCostConfigProps> = ({
  value = {},
  onChange,
-  serverId,
-  serverUrl,
-  accessToken,
+  tools = [],
  disabled = false
 }) => {
  const handleDefaultCostChange = (defaultCost: number | null) => {
@@ -29,13 +25,24 @@ const MCPServerCostConfig: React.FC<MCPServerCostConfigProps> = ({
    onChange?.(updated);
  };

+  const handleToolCostChange = (toolName: string, cost: number | null) => {
+    const updated = {
+      ...value,
+      tool_name_to_cost_per_query: {
+        ...value.tool_name_to_cost_per_query,
+        [toolName]: cost
+      }
+    };
+    onChange?.(updated);
+  };
+
  return (
    <Card>
      <div className="space-y-6">
        <div className="flex items-center gap-2 mb-4">
          <DollarOutlined className="text-green-600" />
          <Title>Cost Configuration</Title>
-          <Tooltip title="Configure costs for this MCP server's tool calls. These costs will be tracked when the server's tools are used.">
+          <Tooltip title="Configure costs for this MCP server's tool calls. Set a default rate and per-tool overrides.">
            <InfoCircleOutlined className="text-gray-400" />
          </Tooltip>
        </div>
@@ -63,15 +70,84 @@ const MCPServerCostConfig: React.FC<MCPServerCostConfigProps> = ({
              Set a default cost for all tool calls to this server
            </Text>
          </div>
+
+          {tools.length > 0 && (
+            <div className="space-y-4">
+              <label className="block text-sm font-medium text-gray-700">
+                Tool-Specific Costs ($)
+                <Tooltip title="Override the default cost for specific tools. Leave blank to use the default rate.">
+                  <InfoCircleOutlined className="ml-1 text-gray-400" />
+                </Tooltip>
+              </label>
+              <Collapse
+                items={[
+                  {
+                    key: '1',
+                    label: (
+                      <div className="flex items-center">
+                        <ToolOutlined className="mr-2 text-blue-500" />
+                        <span className="font-medium">Available Tools</span>
+                        <Badge 
+                          count={tools.length} 
+                          style={{ 
+                            backgroundColor: '#52c41a', 
+                            marginLeft: '8px' 
+                          }} 
+                        />
+                      </div>
+                    ),
+                    children: (
+                      <div className="space-y-3 max-h-64 overflow-y-auto">
+                        {tools.map((tool, index) => (
+                          <div key={index} className="flex items-center justify-between p-3 bg-gray-50 rounded-lg">
+                            <div className="flex-1">
+                              <Text className="font-medium text-gray-900">{tool.name}</Text>
+                              {tool.description && (
+                                <Text className="text-gray-500 text-sm block mt-1">
+                                  {tool.description}
+                                </Text>
+                              )}
+                            </div>
+                            <div className="ml-4">
+                              <InputNumber
+                                min={0}
+                                step={0.0001}
+                                precision={4}
+                                placeholder="Use default"
+                                value={value.tool_name_to_cost_per_query?.[tool.name]}
+                                onChange={(cost) => handleToolCostChange(tool.name, cost)}
+                                disabled={disabled}
+                                style={{ width: '120px' }}
+                                addonBefore="$"
+                              />
+                            </div>
+                          </div>
+                        ))}
+                      </div>
+                    ),
+                  },
+                ]}
+              />
+            </div>
+          )}
        </div>

-        {value.default_cost_per_query && (
+        {(value.default_cost_per_query || (value.tool_name_to_cost_per_query && Object.keys(value.tool_name_to_cost_per_query).length > 0)) && (
          <div className="mt-6 p-4 bg-blue-50 border border-blue-200 rounded-lg">
            <Text className="text-blue-800 font-medium">Cost Summary:</Text>
            <div className="mt-2 space-y-1">
-              <Text className="text-blue-700">
-                • Default cost: ${value.default_cost_per_query.toFixed(4)} per query
-              </Text>
+              {value.default_cost_per_query && (
+                <Text className="text-blue-700">
+                  • Default cost: ${value.default_cost_per_query.toFixed(4)} per query
+                </Text>
+              )}
+              {value.tool_name_to_cost_per_query && Object.entries(value.tool_name_to_cost_per_query).map(([toolName, cost]) => 
+                cost !== null && cost !== undefined && (
+                  <Text key={toolName} className="text-blue-700">
+                    • {toolName}: ${cost.toFixed(4)} per query
+                  </Text>
+                )
+              )}
            </div>
          </div>
        )}
@@ -92,14 +92,12 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({ mcpServer, accessToken, o
        
        <TabPanel>
          <div className="space-y-6">
-            <MCPServerCostConfig
-              value={costConfig}
-              onChange={setCostConfig}
-              serverId={mcpServer.server_id}
-              serverUrl={mcpServer.url}
-              accessToken={accessToken}
-              disabled={false}
-            />
+              <MCPServerCostConfig
+                value={costConfig}
+                onChange={setCostConfig}
+                tools={[]}
+                disabled={false}
+              />
            
            <div className="flex justify-end gap-2">
              <AntdButton onClick={onCancel}>Cancel</AntdButton>
@@ -54,6 +54,7 @@ export interface InputSchemaProperty {
  // Define MCPServerCostInfo for cost tracking
  export interface MCPServerCostInfo {
    default_cost_per_query?: number | null;
+    tool_name_to_cost_per_query?: Record<string, number | null>;
  }

  // Define MCP provider info
@@ -5999,3 +5999,122 @@ export const mcpToolsCall = async (accessToken: string) => {

  return await response.json();
 };
+
+export const testMCPConnectionRequest = async (
+  accessToken: string,
+  mcpServerConfig: Record<string, any>
+) => {
+  try {
+    console.log(
+      "Testing MCP connection with config:",
+      JSON.stringify(mcpServerConfig)
+    );
+
+    // Construct the URL for POST request
+    const url = proxyBaseUrl
+      ? `${proxyBaseUrl}/mcp-rest/test/connection`
+      : `/mcp-rest/test/connection`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+      },
+      body: JSON.stringify(mcpServerConfig),
+    });
+
+    // Check for non-JSON responses first
+    const contentType = response.headers.get("content-type");
+    if (!contentType || !contentType.includes("application/json")) {
+      const text = await response.text();
+      console.error("Received non-JSON response:", text);
+      throw new Error(
+        `Received non-JSON response (${response.status}: ${response.statusText}). Check network tab for details.`
+      );
+    }
+
+    const data = await response.json();
+
+    if (!response.ok || data.status === "error") {
+      // Return the error response instead of throwing an error
+      // This allows the caller to handle the error format properly
+      if (data.status === "error") {
+        return data; // Return the full error response
+      } else {
+        return {
+          status: "error",
+          message:
+            data.error?.message ||
+            `MCP connection test failed: ${response.status} ${response.statusText}`,
+        };
+      }
+    }
+
+    return data;
+  } catch (error) {
+    console.error("MCP connection test error:", error);
+    // For network errors or other exceptions, still throw
+    throw error;
+  }
+};
+
+export const testMCPToolsListRequest = async (
+  accessToken: string,
+  mcpServerConfig: Record<string, any>
+) => {
+  try {
+    console.log(
+      "Testing MCP tools list with config:",
+      JSON.stringify(mcpServerConfig)
+    );
+
+    // Construct the URL for POST request
+    const url = proxyBaseUrl
+      ? `${proxyBaseUrl}/mcp-rest/test/tools/list`
+      : `/mcp-rest/test/tools/list`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+      },
+      body: JSON.stringify(mcpServerConfig),
+    });
+
+    // Check for non-JSON responses first
+    const contentType = response.headers.get("content-type");
+    if (!contentType || !contentType.includes("application/json")) {
+      const text = await response.text();
+      console.error("Received non-JSON response:", text);
+      throw new Error(
+        `Received non-JSON response (${response.status}: ${response.statusText}). Check network tab for details.`
+      );
+    }
+
+    const data = await response.json();
+
+    if (!response.ok || data.error) {
+      // Return the error response instead of throwing an error
+      // This allows the caller to handle the error format properly
+      if (data.error) {
+        return data; // Return the full error response
+      } else {
+        return {
+          tools: [],
+          error: "request_failed",
+          message:
+            data.message ||
+            `MCP tools list failed: ${response.status} ${response.statusText}`,
+        };
+      }
+    }
+
+    return data;
+  } catch (error) {
+    console.error("MCP tools list test error:", error);
+    // For network errors or other exceptions, still throw
+    throw error;
+  }
+};