Skip to content

Commit

Permalink
feat: add socket connection
Browse files Browse the repository at this point in the history
  • Loading branch information
McPatate committed Feb 7, 2024
1 parent 455b085 commit f827940
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 15 deletions.
17 changes: 9 additions & 8 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 5 additions & 4 deletions crates/llm-ls/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ edition = "2021"
name = "llm-ls"

[dependencies]
clap = { version = "4", features = ["derive"] }
home = "0.5"
ropey = { version = "1.6", default-features = false, features = [
"simd",
Expand All @@ -19,7 +20,7 @@ reqwest = { version = "0.11", default-features = false, features = [
serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "1"
tokenizers = { version = "0.14", default-features = false, features = ["onig"] }
tokenizers = { version = "0.15", default-features = false, features = ["onig"] }
tokio = { version = "1", features = [
"fs",
"io-std",
Expand All @@ -37,9 +38,9 @@ tree-sitter-c = "0.20"
tree-sitter-cpp = "0.20"
tree-sitter-c-sharp = "0.20"
tree-sitter-elixir = "0.1"
tree-sitter-erlang = "0.2"
tree-sitter-erlang = "0.4"
tree-sitter-go = "0.20"
tree-sitter-html = "0.19"
tree-sitter-html = "0.20"
tree-sitter-java = "0.20"
tree-sitter-javascript = "0.20"
tree-sitter-json = "0.20"
Expand All @@ -52,7 +53,7 @@ tree-sitter-r = "0.19"
tree-sitter-ruby = "0.20"
tree-sitter-rust = "0.20"
tree-sitter-scala = "0.20"
tree-sitter-swift = "0.3"
tree-sitter-swift = "0.4"
tree-sitter-typescript = "0.20"

[dependencies.uuid]
Expand Down
34 changes: 31 additions & 3 deletions crates/llm-ls/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use clap::Parser;
use ropey::Rope;
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
Expand All @@ -8,6 +9,7 @@ use std::sync::Arc;
use std::time::{Duration, Instant};
use tokenizers::Tokenizer;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::sync::RwLock;
use tower_lsp::jsonrpc::Result as LspResult;
use tower_lsp::lsp_types::*;
Expand Down Expand Up @@ -770,10 +772,22 @@ impl LanguageServer for LlmService {
}
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
/// Wether to use a tcp socket for data transfer
#[arg(long = "port")]
socket: Option<usize>,

/// Wether to use stdio transport for data transfer, ignored because it is the default
/// behaviour
#[arg(short, long, default_value_t = true)]
stdio: bool,
}

#[tokio::main]
async fn main() {
let stdin = tokio::io::stdin();
let stdout = tokio::io::stdout();
let args = Args::parse();

let home_dir = home::home_dir().ok_or(()).expect("failed to find home dir");
let cache_dir = home_dir.join(".cache/llm_ls");
Expand Down Expand Up @@ -822,5 +836,19 @@ async fn main() {
.custom_method("llm-ls/rejectCompletion", LlmService::reject_completion)
.finish();

Server::new(stdin, stdout, socket).serve(service).await;
if let Some(port) = args.socket {
let addr = format!("127.0.0.1:{port}");
let listener = TcpListener::bind(&addr)
.await
.unwrap_or_else(|_| panic!("failed to bind tcp listener to {addr}"));
let (stream, _) = listener
.accept()
.await
.unwrap_or_else(|_| panic!("failed to accept new connections on {addr}"));
let (read, write) = tokio::io::split(stream);
Server::new(read, write, socket).serve(service).await;
} else {
let (stdin, stdout) = (tokio::io::stdin(), tokio::io::stdout());
Server::new(stdin, stdout, socket).serve(service).await;
}
}

0 comments on commit f827940

Please sign in to comment.