move LSP serve method to main cli crate and fix shutdown handling (#143)

Author: joshuadavidthomas

crates/djls-server/src/lib.rs

@@ -3,17 +3,4 @@ mod queue;
 mod server;
 mod workspace;
 
-use crate::server::DjangoLanguageServer;
-use anyhow::Result;
-use tower_lsp_server::{LspService, Server};
-
-pub async fn serve() -> Result<()> {
-    let stdin = tokio::io::stdin();
-    let stdout = tokio::io::stdout();
-
-    let (service, socket) = LspService::build(DjangoLanguageServer::new).finish();
-
-    Server::new(stdin, stdout, socket).serve(service).await;
-
-    Ok(())
-}
+pub use crate::server::DjangoLanguageServer;

crates/djls-server/src/queue.rs

@@ -105,7 +105,7 @@ impl Queue {
                     }
                 }
             }
-            eprintln!("Queue worker task shutting down.");
+            eprintln!("Queue worker task shutting down");
         });
 
         Self {
@@ -179,7 +179,7 @@ impl Drop for QueueInner {
             // `.ok()` ignores the result, as the receiver might have already
             // terminated if the channel closed naturally or panicked.
             sender.send(()).ok();
-            eprintln!("Sent shutdown signal to queue worker.");
+            eprintln!("Sent shutdown signal to queue worker");
         }
     }
 }

crates/djls/src/cli.rs

@@ -1,8 +1,8 @@
 use crate::args::Args;
 use crate::commands::{Command, DjlsCommand};
+use crate::exit::Exit;
 use anyhow::Result;
 use clap::Parser;
-use std::process::ExitCode;
 
 /// Main CLI structure that defines the command-line interface
 #[derive(Parser)]
@@ -16,13 +16,24 @@ pub struct Cli {
     pub args: Args,
 }
 
-/// Parse CLI arguments and execute the chosen command
-pub async fn run(args: Vec<String>) -> Result<ExitCode> {
+/// Parse CLI arguments, execute the chosen command, and handle results
+pub fn run(args: Vec<String>) -> Result<()> {
     let cli = Cli::try_parse_from(args).unwrap_or_else(|e| {
         e.exit();
     });
 
-    match &cli.command {
-        DjlsCommand::Serve(cmd) => cmd.execute(&cli.args).await,
+    let result = match &cli.command {
+        DjlsCommand::Serve(cmd) => cmd.execute(&cli.args),
+    };
+
+    match result {
+        Ok(exit) => exit.process_exit(),
+        Err(e) => {
+            let mut msg = e.to_string();
+            if let Some(source) = e.source() {
+                msg += &format!(", caused by {}", source);
+            }
+            Exit::error().with_message(msg).process_exit()
+        }
     }
 }

crates/djls/src/commands.rs

@@ -1,12 +1,12 @@
 mod serve;
 
 use crate::args::Args;
+use crate::exit::Exit;
 use anyhow::Result;
 use clap::Subcommand;
-use std::process::ExitCode;
 
 pub trait Command {
-    async fn execute(&self, args: &Args) -> Result<ExitCode>;
+    fn execute(&self, args: &Args) -> Result<Exit>;
 }
 
 #[derive(Debug, Subcommand)]

crates/djls/src/commands/serve.rs

@@ -1,8 +1,10 @@
 use crate::args::Args;
 use crate::commands::Command;
+use crate::exit::Exit;
 use anyhow::Result;
 use clap::{Parser, ValueEnum};
-use std::process::ExitCode;
+use djls_server::DjangoLanguageServer;
+use tower_lsp_server::{LspService, Server};
 
 #[derive(Debug, Parser)]
 pub struct Serve {
@@ -17,8 +19,29 @@ enum ConnectionType {
 }
 
 impl Command for Serve {
-    async fn execute(&self, _args: &Args) -> Result<ExitCode> {
-        djls_server::serve().await?;
-        Ok(ExitCode::SUCCESS)
+    fn execute(&self, _args: &Args) -> Result<Exit> {
+        let runtime = tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+
+        let exit_status = runtime.block_on(async {
+            let stdin = tokio::io::stdin();
+            let stdout = tokio::io::stdout();
+
+            let (service, socket) = LspService::build(DjangoLanguageServer::new).finish();
+
+            Server::new(stdin, stdout, socket).serve(service).await;
+
+            // Exit here instead of returning control to the `Cli`, for ... reasons?
+            // If we don't exit here, ~~~ something ~~~ goes on with PyO3 (I assume)
+            // or the Python entrypoint wrapper to indefinitely hang the CLI and keep
+            // the process running
+            Exit::success()
+                .with_message("Server completed successfully")
+                .process_exit()
+        });
+
+        Ok(exit_status)
     }
 }

crates/djls/src/exit.rs

@@ -0,0 +1,100 @@
+use anyhow::Result;
+use std::error::Error;
+use std::fmt;
+
+type ExitMessage = Option<String>;
+
+#[derive(Debug)]
+pub enum ExitStatus {
+    Success,
+    Error,
+}
+
+impl ExitStatus {
+    pub fn as_raw(&self) -> i32 {
+        match self {
+            ExitStatus::Success => 0,
+            ExitStatus::Error => 1,
+        }
+    }
+
+    pub fn as_str(&self) -> &str {
+        match self {
+            ExitStatus::Success => "Command succeeded",
+            ExitStatus::Error => "Command error",
+        }
+    }
+}
+
+impl From<ExitStatus> for i32 {
+    fn from(status: ExitStatus) -> Self {
+        status.as_raw()
+    }
+}
+
+impl fmt::Display for ExitStatus {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let msg = self.as_str();
+        write!(f, "{}", msg)
+    }
+}
+
+#[derive(Debug)]
+pub struct Exit {
+    status: ExitStatus,
+    message: ExitMessage,
+}
+
+impl Exit {
+    fn new(status: ExitStatus) -> Self {
+        Self {
+            status,
+            message: None,
+        }
+    }
+
+    pub fn success() -> Self {
+        Self::new(ExitStatus::Success)
+    }
+
+    pub fn error() -> Self {
+        Self::new(ExitStatus::Error)
+    }
+
+    pub fn with_message<S: Into<String>>(mut self, message: S) -> Self {
+        self.message = Some(message.into());
+        self
+    }
+
+    pub fn process_exit(self) -> ! {
+        if let Some(message) = self.message {
+            println!("{}", message)
+        }
+        std::process::exit(self.status.as_raw())
+    }
+
+    #[allow(dead_code)]
+    pub fn ok(self) -> Result<()> {
+        match self.status {
+            ExitStatus::Success => Ok(()),
+            _ => Err(self.into()),
+        }
+    }
+
+    #[allow(dead_code)]
+    pub fn as_raw(&self) -> i32 {
+        self.status.as_raw()
+    }
+}
+
+impl fmt::Display for Exit {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let status_str = self.status.as_str();
+        match &self.message {
+            Some(msg) => write!(f, "{}: {}", status_str, msg),
+            None => write!(f, "{}", status_str),
+        }
+    }
+}
+
+impl Error for Exit {}

crates/djls/src/lib.rs

@@ -1,40 +1,29 @@
+/// PyO3 entrypoint for the Django Language Server CLI.
+///
+/// This module provides a Python interface using PyO3 to solve Python runtime
+/// interpreter linking issues. The PyO3 approach avoids complexities with
+/// static/dynamic linking when building binaries that interact with Python.
 mod args;
 mod cli;
 mod commands;
+mod exit;
 
 use pyo3::prelude::*;
 use std::env;
-use std::process::ExitCode;
 
 #[pyfunction]
+/// Entry point called by Python when the CLI is invoked.
+/// This function handles argument parsing from Python and routes to the Rust CLI logic.
 fn entrypoint(_py: Python) -> PyResult<()> {
     // Skip python interpreter and script path, add command name
     let args: Vec<String> = std::iter::once("djls".to_string())
         .chain(env::args().skip(2))
         .collect();
 
-    let runtime = tokio::runtime::Builder::new_multi_thread()
-        .enable_all()
-        .build()
-        .unwrap();
+    // Run the CLI with the adjusted args
+    cli::run(args).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
 
-    let result = runtime.block_on(cli::run(args));
-
-    match result {
-        Ok(code) => {
-            if code != ExitCode::SUCCESS {
-                std::process::exit(1);
-            }
-            Ok(())
-        }
-        Err(e) => {
-            eprintln!("Error: {}", e);
-            if let Some(source) = e.source() {
-                eprintln!("Caused by: {}", source);
-            }
-            std::process::exit(1);
-        }
-    }
+    Ok(())
 }
 
 #[pymodule]

crates/djls/src/main.rs

@@ -0,0 +1,19 @@
+/// Binary interface for local development only.
+///
+/// This binary exists for development and testing with `cargo run`.
+/// The production CLI is distributed through the PyO3 interface in lib.rs.
+mod args;
+mod cli;
+mod commands;
+mod exit;
+
+use anyhow::Result;
+
+/// Process CLI args and run the appropriate command.
+fn main() -> Result<()> {
+    // Get command line arguments
+    let args: Vec<String> = std::env::args().collect();
+
+    // Call the unified CLI run function
+    cli::run(args)
+}