mod common;
use std::borrow::Cow;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio::sync::mpsc;
use lspf::types::request::Request;
use lspf::{
Context, RawMessage, RequestId, Server, Transport, TransportError, TransportReader,
TransportWriter,
};
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct GreetParams {
name: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct GreetResult {
greeting: String,
}
enum Greet {}
impl Request for Greet {
type Params = GreetParams;
type Result = GreetResult;
const METHOD: &'static str = "custom/greet";
}
struct AppState {
handled: Arc<AtomicUsize>,
prefix: String,
}
async fn greet(
state: Arc<AppState>,
_ctx: Context,
params: GreetParams,
_ct: lspf::CancellationToken,
) -> Result<GreetResult, lspf::LspError> {
state.handled.fetch_add(1, Ordering::SeqCst);
Ok(GreetResult {
greeting: format!("{}, {}!", state.prefix, params.name),
})
}
struct ChannelTransport {
in_rx: mpsc::UnboundedReceiver<RawMessage>,
out_tx: mpsc::UnboundedSender<RawMessage>,
}
struct ChannelReader {
in_rx: mpsc::UnboundedReceiver<RawMessage>,
}
struct ChannelWriter {
out_tx: mpsc::UnboundedSender<RawMessage>,
}
impl Transport for ChannelTransport {
type Reader = ChannelReader;
type Writer = ChannelWriter;
fn split(self) -> (Self::Reader, Self::Writer) {
(
ChannelReader { in_rx: self.in_rx },
ChannelWriter {
out_tx: self.out_tx,
},
)
}
}
impl TransportReader for ChannelReader {
async fn recv(&mut self) -> Result<RawMessage, TransportError> {
self.in_rx.recv().await.ok_or(TransportError::Closed)
}
}
impl TransportWriter for ChannelWriter {
async fn send(&mut self, msg: RawMessage) -> Result<(), TransportError> {
self.out_tx.send(msg).map_err(|_| TransportError::Closed)
}
async fn shutdown(self) -> Result<(), TransportError> {
Ok(())
}
}
fn request(id: i32, method: &'static str, params: serde_json::Value) -> RawMessage {
RawMessage::Request {
id: RequestId::Number(id),
method: Cow::Borrowed(method),
params: Bytes::from(serde_json::to_vec(¶ms).unwrap()),
}
}
fn initialize_request(id: i32) -> RawMessage {
request(
id,
"initialize",
json!({ "processId": null, "rootUri": null, "capabilities": {} }),
)
}
fn notification(method: &'static str) -> RawMessage {
RawMessage::Notification {
method: Cow::Borrowed(method),
params: Bytes::from_static(b"null"),
}
}
async fn drive(messages: Vec<RawMessage>) -> (Vec<RawMessage>, usize) {
let (in_tx, in_rx) = mpsc::unbounded_channel::<RawMessage>();
let (out_tx, mut out_rx) = mpsc::unbounded_channel::<RawMessage>();
let transport = ChannelTransport { in_rx, out_tx };
let handled = Arc::new(AtomicUsize::new(0));
let server = Server::builder(AppState {
handled: Arc::clone(&handled),
prefix: "Hello".to_string(),
})
.request::<Greet, _, _>(greet)
.build()
.expect("server builds");
let handle = tokio::spawn(async move { server.serve(transport).await });
let mut outbox = Vec::new();
for msg in messages {
let response_id = msg.id().cloned();
in_tx.send(msg).unwrap();
if let Some(response_id) = response_id {
let response = tokio::time::timeout(Duration::from_secs(2), out_rx.recv())
.await
.expect("response arrived within 2s")
.expect("writer remained open");
assert_eq!(response.id(), Some(&response_id));
outbox.push(response);
}
}
drop(in_tx);
tokio::time::timeout(Duration::from_secs(2), handle)
.await
.expect("serve returned within 2s")
.expect("server task did not panic")
.expect("serve ended cleanly");
outbox.extend(std::iter::from_fn(|| out_rx.try_recv().ok()));
(outbox, handled.load(Ordering::SeqCst))
}
fn response(outbox: &[RawMessage], id: i32) -> Option<&RawMessage> {
outbox.iter().find(
|m| matches!(m, RawMessage::Response { id: rid, .. } if *rid == RequestId::Number(id)),
)
}
fn ok_result(outbox: &[RawMessage], id: i32) -> Option<serde_json::Value> {
match response(outbox, id)? {
RawMessage::Response {
result: Ok(bytes), ..
} => Some(serde_json::from_slice(bytes).unwrap()),
_ => None,
}
}
fn error_code(outbox: &[RawMessage], id: i32) -> Option<i32> {
match response(outbox, id)? {
RawMessage::Response { result: Err(e), .. } => Some(e.code),
_ => None,
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn initialize_custom_request_shutdown_round_trip() {
let (outbox, handled) = drive(vec![
initialize_request(1),
request(2, "custom/greet", json!({ "name": "Ada" })),
request(3, "shutdown", json!(null)),
notification("exit"),
])
.await;
let init: lspf::types::InitializeResult =
serde_json::from_value(ok_result(&outbox, 1).expect("initialize response")).unwrap();
assert_eq!(
init.capabilities,
lspf::types::ServerCapabilities {
position_encoding: Some(lspf::types::PositionEncodingKind::UTF16),
text_document_sync: Some(lspf::types::TextDocumentSyncCapability::Kind(
lspf::types::TextDocumentSyncKind::INCREMENTAL,
)),
workspace: Some(common::workspace_capabilities()),
..lspf::types::ServerCapabilities::default()
},
"custom requests must not add ServerCapabilities beyond protocol-owned fields"
);
let greeting: GreetResult =
serde_json::from_value(ok_result(&outbox, 2).expect("greet response")).unwrap();
assert_eq!(
greeting,
GreetResult {
greeting: "Hello, Ada!".to_string()
}
);
assert_eq!(ok_result(&outbox, 3), Some(serde_json::Value::Null));
assert_eq!(handled, 1, "the typed handler ran exactly once");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn invalid_params_return_invalid_params_and_a_later_request_succeeds() {
let (outbox, handled) = drive(vec![
initialize_request(1),
request(2, "custom/greet", json!({ "name": 42 })),
request(3, "custom/greet", json!({ "name": "Bob" })),
notification("exit"),
])
.await;
assert_eq!(
error_code(&outbox, 2),
Some(-32602),
"malformed params must return InvalidParams"
);
let greeting: GreetResult =
serde_json::from_value(ok_result(&outbox, 3).expect("later greet response")).unwrap();
assert_eq!(greeting.greeting, "Hello, Bob!");
assert_eq!(
handled, 1,
"the handler ran only for the valid request, never for malformed params"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn request_before_initialize_returns_server_not_initialized() {
let (outbox, handled) = drive(vec![
request(1, "custom/greet", json!({ "name": "Ada" })),
notification("exit"),
])
.await;
assert_eq!(
error_code(&outbox, 1),
Some(-32002),
"a request before initialize must return ServerNotInitialized"
);
assert_eq!(handled, 0, "the handler must not run before initialize");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_second_initialize_is_rejected() {
let (outbox, _handled) = drive(vec![
initialize_request(1),
initialize_request(2),
notification("exit"),
])
.await;
assert!(
ok_result(&outbox, 1).is_some(),
"the first initialize succeeds"
);
assert_eq!(
error_code(&outbox, 2),
Some(-32600),
"a second initialize must be refused"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn request_after_shutdown_returns_invalid_request() {
let (outbox, handled) = drive(vec![
initialize_request(1),
request(2, "shutdown", json!(null)),
request(3, "custom/greet", json!({ "name": "Ada" })),
notification("exit"),
])
.await;
assert_eq!(
error_code(&outbox, 3),
Some(-32600),
"a request after shutdown must return InvalidRequest"
);
assert_eq!(handled, 0, "the handler must not run after shutdown");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unknown_request_returns_method_not_found() {
let (outbox, _handled) = drive(vec![
initialize_request(1),
request(2, "custom/unregistered", json!({})),
notification("exit"),
])
.await;
assert_eq!(
error_code(&outbox, 2),
Some(-32601),
"an unregistered method must return MethodNotFound"
);
}