feat: add MCP server connection status (#5018)

This commit is contained in:
Louis 2025-05-19 15:17:19 +07:00 committed by GitHub
parent 90da49f873
commit 2ae7417e10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 88 additions and 104 deletions

View File

@ -2,7 +2,7 @@ use std::{collections::HashMap, env, sync::Arc};
use rmcp::{service::RunningService, transport::TokioChildProcess, RoleClient, ServiceExt};
use serde_json::Value;
use tauri::{AppHandle, Emitter, State};
use tauri::{AppHandle, Emitter, Runtime, State};
use tokio::{process::Command, sync::Mutex};
use super::{cmd::get_jan_data_folder_path, state::AppState};
@ -16,15 +16,17 @@ use super::{cmd::get_jan_data_folder_path, state::AppState};
/// # Returns
/// * `Ok(())` if servers were initialized successfully
/// * `Err(String)` if there was an error reading config or starting servers
pub async fn run_mcp_commands(
app_path: String,
pub async fn run_mcp_commands<R: Runtime>(
app: &AppHandle<R>,
servers_state: Arc<Mutex<HashMap<String, RunningService<RoleClient, ()>>>>,
) -> Result<(), String> {
let app_path = get_jan_data_folder_path(app.clone());
let app_path_str = app_path.to_str().unwrap().to_string();
log::info!(
"Load MCP configs from {}",
app_path.clone() + "/mcp_config.json"
app_path_str.clone() + "/mcp_config.json"
);
let config_content = std::fs::read_to_string(app_path.clone() + "/mcp_config.json")
let config_content = std::fs::read_to_string(app_path_str.clone() + "/mcp_config.json")
.map_err(|e| format!("Failed to read config file: {}", e))?;
let mcp_servers: serde_json::Value = serde_json::from_str(&config_content)
@ -68,12 +70,28 @@ pub async fn run_mcp_commands(
}
});
let service =
().serve(TokioChildProcess::new(cmd).map_err(|e| e.to_string())?)
.await
.map_err(|e| e.to_string())?;
let process = TokioChildProcess::new(cmd);
match process {
Ok(p) => {
let service = ().serve(p).await;
servers_state.lock().await.insert(name.clone(), service);
match service {
Ok(running_service) => {
servers_state
.lock()
.await
.insert(name.clone(), running_service);
log::info!("Server {name} started successfully.");
}
Err(e) => {
log::error!("Failed to start server {name}: {e}");
}
}
}
Err(e) => {
log::error!("Failed to run command {name}: {e}");
}
}
}
}
}
@ -84,6 +102,19 @@ pub async fn run_mcp_commands(
// Initialize
let _server_info = service.peer_info();
log::info!("Connected to server: {_server_info:#?}");
// Emit event to the frontend
let event = format!("mcp-connected");
let server_info: &rmcp::model::InitializeResult = service.peer_info();
let name = server_info.server_info.name.clone();
let version = server_info.server_info.version.clone();
let payload = serde_json::json!({
"name": name,
"version": version,
});
// service.peer_info().server_info.name
app.emit(&event, payload)
.map_err(|e| format!("Failed to emit event: {}", e))?;
log::info!("Emitted event: {event}");
}
Ok(())
}
@ -110,14 +141,12 @@ fn extract_active_status(config: &Value) -> Option<bool> {
#[tauri::command]
pub async fn restart_mcp_servers(app: AppHandle, state: State<'_, AppState>) -> Result<(), String> {
let app_path = get_jan_data_folder_path(app.clone());
let app_path_str = app_path.to_str().unwrap().to_string();
let servers = state.mcp_servers.clone();
// Stop the servers
stop_mcp_servers(state.mcp_servers.clone()).await?;
// Restart the servers
run_mcp_commands(app_path_str, servers).await?;
run_mcp_commands(&app, servers).await?;
app.emit("mcp-update", "MCP servers updated")
.map_err(|e| format!("Failed to emit event: {}", e))
@ -136,6 +165,15 @@ pub async fn stop_mcp_servers(
drop(servers_map); // Release the lock after stopping
Ok(())
}
#[tauri::command]
pub async fn get_connected_servers(
_app: AppHandle,
state: State<'_, AppState>,
) -> Result<Vec<String>, String> {
let servers = state.mcp_servers.clone();
let servers_map = servers.lock().await;
Ok(servers_map.keys().cloned().collect())
}
#[cfg(test)]
mod tests {
@ -144,21 +182,22 @@ mod tests {
use std::fs::File;
use std::io::Write;
use std::sync::Arc;
use tauri::test::mock_app;
use tokio::sync::Mutex;
#[tokio::test]
async fn test_run_mcp_commands() {
let app = mock_app();
// Create a mock mcp_config.json file
let config_path = "mcp_config.json";
let mut file = File::create(config_path).expect("Failed to create config file");
let mut file: File = File::create(config_path).expect("Failed to create config file");
file.write_all(b"{\"mcpServers\":{}}")
.expect("Failed to write to config file");
// Call the run_mcp_commands function
let app_path = ".".to_string();
let servers_state: Arc<Mutex<HashMap<String, RunningService<RoleClient, ()>>>> =
Arc::new(Mutex::new(HashMap::new()));
let result = run_mcp_commands(app_path, servers_state).await;
let result = run_mcp_commands(app.handle(), servers_state).await;
// Assert that the function returns Ok(())
assert!(result.is_ok());

View File

@ -190,14 +190,11 @@ fn extract_extension_manifest<R: Read>(
}
pub fn setup_mcp(app: &App) {
let app_path = get_jan_data_folder_path(app.handle().clone());
let state = app.state::<AppState>().inner();
let app_path_str = app_path.to_str().unwrap().to_string();
let servers = state.mcp_servers.clone();
let app_handle = app.handle().clone();
let app_handle: tauri::AppHandle = app.handle().clone();
tauri::async_runtime::spawn(async move {
if let Err(e) = run_mcp_commands(app_path_str, servers).await {
if let Err(e) = run_mcp_commands(&app_handle, servers).await {
log::error!("Failed to run mcp commands: {}", e);
}
app_handle

View File

@ -50,6 +50,7 @@ pub fn run() {
core::cmd::get_tools,
core::cmd::call_tool,
core::mcp::restart_mcp_servers,
core::mcp::get_connected_servers,
// Threads
core::threads::list_threads,
core::threads::create_thread,

View File

@ -25,21 +25,6 @@ type MCPServerStoreState = {
addServer: (key: string, config: MCPServerConfig) => void
editServer: (key: string, config: MCPServerConfig) => void
deleteServer: (key: string) => void
fetchMCPServers: () => Promise<void>
}
// Mock data for MCP servers
export const mockMCPServers: MCPServers = {
puppeteer: {
command: 'npx',
args: ['-y', '@tokenizin/mcp-npx-fetch'],
env: {},
},
inspector: {
command: 'npx',
args: ['-y', '@modelcontextprotocol/inspector'],
env: {},
},
}
export const useMCPServers = create<MCPServerStoreState>()(
@ -90,66 +75,6 @@ export const useMCPServers = create<MCPServerStoreState>()(
deletedServerKeys: [...state.deletedServerKeys, key],
}
}),
// Fetch MCP servers
fetchMCPServers: async () => {
set({ loading: true })
// Simulate API call with mock data
const response = await new Promise<MCPServers>((resolve) =>
setTimeout(() => resolve(mockMCPServers), 0)
)
set((state) => {
// Filter out deleted servers from the response
const filteredResponse = { ...response }
state.deletedServerKeys.forEach((key) => {
delete filteredResponse[key]
})
const localKeys = Object.keys(state.mcpServers)
const responseKeys = Object.keys(filteredResponse)
// Check if the keys are the same
const hasSameKeys =
localKeys.length === responseKeys.length &&
localKeys.every((key) => responseKeys.includes(key))
// Check if values are the same for each key
const hasSameValues =
hasSameKeys &&
localKeys.every((key) => {
const current = state.mcpServers[key]
const resp = filteredResponse[key]
return (
current.command === resp.command &&
JSON.stringify(current.args) === JSON.stringify(resp.args) &&
JSON.stringify(current.env) === JSON.stringify(resp.env)
)
})
// If everything is the same, don't update
if (hasSameValues) {
return { loading: false }
}
// Add only new servers, preserving existing ones
const existingKeys = new Set(localKeys)
const newServers: MCPServers = {}
responseKeys.forEach((key) => {
if (!existingKeys.has(key)) {
newServers[key] = filteredResponse[key]
}
})
return {
mcpServers: { ...newServers, ...state.mcpServers },
loading: false,
}
})
},
}),
{
name: localStoregeKey.settingMCPSevers, // Using existing key for now

View File

@ -177,7 +177,10 @@ export const stopModel = async (
* @param tools
* @returns
*/
export const normalizeTools = (tools: MCPTool[]): ChatCompletionTool[] => {
export const normalizeTools = (
tools: MCPTool[]
): ChatCompletionTool[] | undefined => {
if (tools.length === 0) return undefined
return tools.map((tool) => ({
type: 'function',
function: {

View File

@ -25,6 +25,7 @@ export const AppRoutes = [
'saveMcpConfigs',
'getMcpConfigs',
'restartMcpServers',
'getConnectedServers',
]
// Define API routes based on different route types
export const Routes = [...CoreRoutes, ...APIRoutes, ...AppRoutes].map((r) => ({

View File

@ -16,6 +16,8 @@ import AddEditMCPServer from '@/containers/dialogs/AddEditMCPServer'
import DeleteMCPServerConfirm from '@/containers/dialogs/DeleteMCPServerConfirm'
import EditJsonMCPserver from '@/containers/dialogs/EditJsonMCPserver'
import { Switch } from '@/components/ui/switch'
import { twMerge } from 'tailwind-merge'
import { getConnectedServers } from '@/services/mcp'
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export const Route = createFileRoute(route.settings.mcp_servers as any)({
@ -23,12 +25,7 @@ export const Route = createFileRoute(route.settings.mcp_servers as any)({
})
function MCPServers() {
const { fetchMCPServers, mcpServers, addServer, editServer, deleteServer } =
useMCPServers()
useEffect(() => {
fetchMCPServers()
}, [fetchMCPServers])
const { mcpServers, addServer, editServer, deleteServer } = useMCPServers()
const [open, setOpen] = useState(false)
const [editingKey, setEditingKey] = useState<string | null>(null)
@ -46,6 +43,7 @@ function MCPServers() {
const [jsonEditorData, setJsonEditorData] = useState<
MCPServerConfig | Record<string, MCPServerConfig> | undefined
>(undefined)
const [connectedServers, setConnectedServers] = useState<string[]>([])
const handleOpenDialog = (serverKey?: string) => {
if (serverKey) {
@ -137,6 +135,10 @@ function MCPServers() {
}
}
useEffect(() => {
getConnectedServers().then(setConnectedServers)
}, [setConnectedServers])
return (
<div className="flex flex-col h-full">
<HeaderPage>
@ -198,7 +200,14 @@ function MCPServers() {
align="start"
title={
<div className="flex items-center gap-x-2">
<div className="size-2 rounded-full bg-accent" />
<div
className={twMerge(
'size-2 rounded-full',
connectedServers.includes(key)
? 'bg-accent'
: 'bg-main-view-fg/50'
)}
/>
{/* condition here when server is running or not */}
{/* <div className="size-2 rounded-full bg-main-view-fg/20" /> */}
<h1 className="text-main-view-fg text-base capitalize">

View File

@ -18,6 +18,15 @@ export const getTools = (): Promise<MCPTool[]> => {
return window.core?.api?.getTools()
}
/**
* @description This function gets connected MCP servers.
* @returns {Promise<string[]>} The MCP names
* @returns
*/
export const getConnectedServers = (): Promise<string[]> => {
return window.core?.api?.getConnectedServers()
}
/**
* @description This function invoke an MCP tool
* @param tool