50 lines
1.7 KiB
Rust
50 lines
1.7 KiB
Rust
use std::{collections::HashMap, sync::Arc};
|
|
|
|
use crate::core::downloads::models::DownloadManagerState;
|
|
use rmcp::{
|
|
model::{CallToolRequestParam, CallToolResult, InitializeRequestParam, Tool},
|
|
service::RunningService,
|
|
RoleClient, ServiceError,
|
|
};
|
|
use tokio::sync::{Mutex, oneshot};
|
|
use tokio::task::JoinHandle;
|
|
|
|
/// Server handle type for managing the proxy server lifecycle
|
|
pub type ServerHandle = JoinHandle<Result<(), Box<dyn std::error::Error + Send + Sync>>>;
|
|
|
|
pub enum RunningServiceEnum {
|
|
NoInit(RunningService<RoleClient, ()>),
|
|
WithInit(RunningService<RoleClient, InitializeRequestParam>),
|
|
}
|
|
pub type SharedMcpServers = Arc<Mutex<HashMap<String, RunningServiceEnum>>>;
|
|
|
|
#[derive(Default)]
|
|
pub struct AppState {
|
|
pub app_token: Option<String>,
|
|
pub mcp_servers: SharedMcpServers,
|
|
pub download_manager: Arc<Mutex<DownloadManagerState>>,
|
|
pub mcp_restart_counts: Arc<Mutex<HashMap<String, u32>>>,
|
|
pub mcp_active_servers: Arc<Mutex<HashMap<String, serde_json::Value>>>,
|
|
pub mcp_successfully_connected: Arc<Mutex<HashMap<String, bool>>>,
|
|
pub server_handle: Arc<Mutex<Option<ServerHandle>>>,
|
|
pub tool_call_cancellations: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
|
|
}
|
|
|
|
impl RunningServiceEnum {
|
|
pub async fn list_all_tools(&self) -> Result<Vec<Tool>, ServiceError> {
|
|
match self {
|
|
Self::NoInit(s) => s.list_all_tools().await,
|
|
Self::WithInit(s) => s.list_all_tools().await,
|
|
}
|
|
}
|
|
pub async fn call_tool(
|
|
&self,
|
|
params: CallToolRequestParam,
|
|
) -> Result<CallToolResult, ServiceError> {
|
|
match self {
|
|
Self::NoInit(s) => s.call_tool(params).await,
|
|
Self::WithInit(s) => s.call_tool(params).await,
|
|
}
|
|
}
|
|
}
|