main.rs

  1use std::ops::ControlFlow;
  2use std::path::PathBuf;
  3
  4use async_lsp::concurrency::ConcurrencyLayer;
  5use async_lsp::lsp_types::notification::{Progress, PublishDiagnostics, ShowMessage};
  6use async_lsp::lsp_types::{
  7    ClientCapabilities, DidOpenTextDocumentParams, InitializeParams, InitializedParams, Url,
  8    WindowClientCapabilities, WorkspaceFolder,
  9};
 10use async_lsp::panic::CatchUnwindLayer;
 11use async_lsp::router::Router;
 12use async_lsp::tracing::TracingLayer;
 13use async_lsp::{Error, ErrorCode, LanguageServer};
 14use async_process::{Command as ProcessCommand, Stdio};
 15use clap::Parser;
 16use futures::channel::oneshot;
 17use rmcp::ServerHandler as MCPServerHandler;
 18use rmcp::handler::server::tool::ToolRouter;
 19use rmcp::handler::server::wrapper::Parameters;
 20use rmcp::model::CallToolResult;
 21use rmcp::{ErrorData as MCPError, schemars};
 22use rmcp::{tool, tool_handler, tool_router};
 23use tower::ServiceBuilder;
 24
 25struct ClientState {
 26    indexed_tx: Option<oneshot::Sender<()>>,
 27}
 28
 29mod ControlFlowState {
 30    pub struct Stop;
 31}
 32
 33#[derive(Clone)]
 34struct MCPServer {
 35    tool_router: ToolRouter<Self>,
 36}
 37
 38#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
 39struct ReadToolArgs {
 40    path: PathBuf,
 41    range_start: usize,
 42    range_end: usize,
 43}
 44
 45struct ReadToolOutput {}
 46
 47#[tool_router]
 48impl MCPServer {
 49    /// Description take from Amp's description of their built-in `read` tool
 50    #[tool(description = "Read a file or list a directory from the file system")]
 51    async fn read(
 52        &self,
 53        Parameters(args): Parameters<ReadToolArgs>,
 54    ) -> Result<CallToolResult, MCPError> {
 55        // TODO: Read the file in the file system
 56        // TODO: Push "didOpenFile" to LSP
 57        // TODO: Get diagnostics from LSP
 58        // TODO: Construct and return result
 59        Ok(CallToolResult {
 60            content: vec![],
 61            structured_content: None,
 62            is_error: None,
 63            meta: None,
 64        })
 65    }
 66}
 67
 68#[tool_handler]
 69impl MCPServerHandler for MCPServer {}
 70
 71#[derive(clap::Parser)]
 72#[command(version, about)]
 73struct CommandLineArgs {
 74    #[arg(short, long, value_name = "FILE")]
 75    config: PathBuf,
 76}
 77
 78#[derive(serde::Deserialize, Debug)]
 79struct Command {
 80    command: String,
 81    args: Vec<String>,
 82}
 83
 84#[derive(serde::Deserialize, Debug)]
 85struct Config {
 86    lsp_server_command: Command,
 87    project_root: PathBuf,
 88}
 89
 90#[tokio::main]
 91async fn main() -> Result<(), Box<dyn std::error::Error>> {
 92    let args = CommandLineArgs::parse();
 93
 94    let config_content = tokio::fs::read_to_string(&args.config).await?;
 95    let config: Config = toml::from_str(&config_content)?;
 96
 97    let mut child = ProcessCommand::new(&config.lsp_server_command.command)
 98        .args(&config.lsp_server_command.args)
 99        .current_dir(&config.project_root)
100        .stdin(Stdio::piped())
101        .stdout(Stdio::piped())
102        .stderr(Stdio::inherit())
103        .kill_on_drop(true)
104        .spawn()
105        .expect(format!("Failed to start lsp: {:?}", &config.lsp_server_command).as_ref());
106
107    let stdin = child.stdin.take().unwrap();
108    let stdout = child.stdout.take().unwrap();
109
110    let (mainloop, mut server) = async_lsp::MainLoop::new_client(|server| {
111        let mut router = Router::new(ClientState { indexed_tx: None });
112
113        router
114            .notification::<Progress>(|this, params| {
115                tracing::info!("{:?} {:?}", params.token, params.value);
116                ControlFlow::Continue(())
117            })
118            .notification::<PublishDiagnostics>(|this, params| {
119                tracing::info!(
120                    "{:?} {:?} {:?}",
121                    params.uri,
122                    params.version,
123                    params.diagnostics
124                );
125                ControlFlow::Continue(())
126            })
127            .notification::<ShowMessage>(|this, params| {
128                tracing::info!("Message {:?}: {}", params.typ, params.message);
129                ControlFlow::Continue(())
130            })
131            .event(|_, _: ControlFlowState::Stop| ControlFlow::Break(Ok(())));
132
133        ServiceBuilder::new().service(router)
134    }); // Initialize.
135
136    let mainloop_fut = tokio::spawn(async move {
137        mainloop.run_buffered(stdout, stdin).await.unwrap();
138    });
139
140    let init_ret = server
141        .initialize(InitializeParams {
142            workspace_folders: Some(vec![WorkspaceFolder {
143                uri: Url::from_file_path(&config.project_root).unwrap(),
144                name: "root".into(),
145            }]),
146            capabilities: ClientCapabilities {
147                window: Some(WindowClientCapabilities {
148                    work_done_progress: Some(true),
149                    ..WindowClientCapabilities::default()
150                }),
151                ..ClientCapabilities::default()
152            },
153            ..InitializeParams::default()
154        })
155        .await
156        .unwrap();
157    tracing::info!("Initialized: {init_ret:?}");
158    server.initialized(InitializedParams {}).unwrap();
159
160    mainloop_fut.await.unwrap();
161
162    Ok(())
163}