use std::sync::Arc;
use crate::config::Config;
use crate::interceptors::completions::CompletionCompressor;
use crate::interceptors::diagnostics::DiagnosticsCompressor;
use crate::interceptors::hover::HoverCompressor;
use crate::interceptors::locations::LocationCompressor;
use crate::interceptors::symbols::DocumentSymbolCompressor;
use crate::interceptors::workspace_diagnostics::WorkspaceDiagnosticCompressor;
use crate::interceptors::workspace_symbols::WorkspaceSymbolCompressor;
use crate::interceptors::{Direction, Interceptor, InterceptorChain};
use crate::mcp::LspSession;
use serde_json::Value;
use tokio::sync::RwLock;
pub struct AgentHandle {
session: LspSession,
language: String,
interceptor_chain: Option<InterceptorChain>,
}
impl std::fmt::Debug for AgentHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentHandle")
.field("language", &self.language)
.field("compression", &self.interceptor_chain.is_some())
.field("session", &"LspSession { .. }")
.finish()
}
}
impl AgentHandle {
pub fn builder() -> AgentBuilder {
AgentBuilder::default()
}
#[allow(dead_code)]
pub(crate) fn new(session: LspSession, language: String, compression: bool) -> Self {
let interceptor_chain = if compression {
Some(build_interceptor_chain())
} else {
None
};
Self {
session,
language,
interceptor_chain,
}
}
async fn open_file(&mut self, uri: &str) -> Result<String, anyhow::Error> {
let path = uri
.strip_prefix("file://")
.ok_or_else(|| anyhow::anyhow!("URI must start with file://"))?;
let content = tokio::fs::read_to_string(path).await?;
self.session
.send_notification(
"textDocument/didOpen",
serde_json::json!({
"textDocument": {
"uri": uri,
"languageId": self.language,
"version": 1,
"text": content,
}
}),
)
.await?;
Ok(content)
}
async fn process_through_chain(
&self,
method: &str,
params: Value,
) -> Result<Value, anyhow::Error> {
match &self.interceptor_chain {
Some(chain) => match chain
.process(method, params.clone(), Direction::ServerToClient)
.await
{
Ok(Some(p)) => Ok(p),
Ok(None) => Ok(Value::Null),
Err(e) => {
tracing::warn!(
error = %e,
method = %method,
"Interceptor chain failed, returning original"
);
Ok(params)
}
},
None => Ok(params),
}
}
pub async fn get_diagnostics(&mut self, uri: &str) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let params = self
.session
.wait_for_notification("textDocument/publishDiagnostics")
.await?;
let processed = self
.process_through_chain("textDocument/publishDiagnostics", params)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_completions(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/completion",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/completion", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_symbols(&mut self, uri: &str) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/documentSymbol",
serde_json::json!({
"textDocument": { "uri": uri },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/documentSymbol", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_hover(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/hover",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/hover", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_references(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/references",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
"context": { "includeDeclaration": true },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/references", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_definition(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/definition",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/definition", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_implementation(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/implementation",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/implementation", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_type_definition(
&mut self,
uri: &str,
line: u32,
character: u32,
) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"textDocument/typeDefinition",
serde_json::json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
}),
)
.await?;
let processed = self
.process_through_chain("textDocument/typeDefinition", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_workspace_symbols(&mut self, query: &str) -> Result<String, anyhow::Error> {
let result = self
.session
.send_request(
"workspace/symbol",
serde_json::json!({
"query": query,
}),
)
.await?;
let processed = self
.process_through_chain("workspace/symbol", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub async fn get_workspace_diagnostics(&mut self, uri: &str) -> Result<String, anyhow::Error> {
self.open_file(uri).await?;
let result = self
.session
.send_request(
"workspace/diagnostic",
serde_json::json!({
"previousResultId": null,
"textDocument": { "uri": uri },
}),
)
.await?;
let processed = self
.process_through_chain("workspace/diagnostic", result)
.await?;
Ok(serde_json::to_string_pretty(&processed)?)
}
pub fn inflate(compressed_json: &str) -> Result<String, anyhow::Error> {
let compressed: Value = serde_json::from_str(compressed_json)?;
let expanded = crate::codec::compact::decompress(&compressed)?;
Ok(serde_json::to_string_pretty(&expanded)?)
}
pub fn compress(raw_json: &str) -> Result<String, anyhow::Error> {
let raw: Value = serde_json::from_str(raw_json)?;
let compressed = crate::codec::compact::compress(&raw)?;
Ok(serde_json::to_string_pretty(&compressed)?)
}
pub async fn shutdown(mut self) -> Result<(), anyhow::Error> {
let _ = self
.session
.send_request("shutdown", serde_json::json!({}))
.await;
self.session
.send_notification("exit", serde_json::json!({}))
.await?;
Ok(())
}
}
#[derive(Default)]
pub struct AgentBuilder {
backend: Option<String>,
language: Option<String>,
compression: bool,
}
impl AgentBuilder {
pub fn backend(mut self, cmd: impl Into<String>) -> Self {
self.backend = Some(cmd.into());
self
}
pub fn language(mut self, lang: impl Into<String>) -> Self {
self.language = Some(lang.into());
self
}
pub fn enable_compression(mut self, enabled: bool) -> Self {
self.compression = enabled;
self
}
pub async fn start(self) -> Result<AgentHandle, anyhow::Error> {
let backend = self
.backend
.ok_or_else(|| anyhow::anyhow!("backend is required"))?;
let language = self
.language
.ok_or_else(|| anyhow::anyhow!("language is required"))?;
let mut session = LspSession::spawn(&backend)?;
session.initialize().await?;
tracing::info!(%backend, %language, "Agent session started");
Ok(AgentHandle::new(session, language, self.compression))
}
}
fn build_interceptor_chain() -> InterceptorChain {
let config = Arc::new(RwLock::new(Config {
backend_cmd: String::new(),
capping: crate::CappingConfig::default(),
enable_diag_compress: true,
enable_completion_compress: true,
enable_hover_compress: true,
enable_document_symbol_compress: true,
enable_location_compress: true,
enable_workspace_symbol_compress: true,
enable_workspace_diag_compress: true,
output_format: crate::OutputFormat::Json,
log_level: "info".into(),
metrics: crate::MetricsConfig::default(),
}));
let interceptors: Vec<Box<dyn Interceptor>> = vec![
Box::new(DiagnosticsCompressor::default()),
Box::new(CompletionCompressor::default()),
Box::new(HoverCompressor::default()),
Box::new(DocumentSymbolCompressor),
Box::new(LocationCompressor),
Box::new(WorkspaceSymbolCompressor),
Box::new(WorkspaceDiagnosticCompressor),
];
InterceptorChain::new(interceptors, config)
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicU16, Ordering};
use crate::codec::json_rpc::LspMessage;
use crate::mcp::LspSession;
use crate::transport::mock::MockTransport;
use serde_json::json;
use super::*;
static TEST_COUNTER: AtomicU16 = AtomicU16::new(0);
fn temp_file(content: &str) -> (String, String) {
let id = TEST_COUNTER.fetch_add(1, Ordering::SeqCst);
let path = format!("/tmp/lspz-test-{id}.rs");
std::fs::write(&path, content).unwrap();
(format!("file://{path}"), path)
}
fn mock_handle(responses: Vec<LspMessage>) -> AgentHandle {
let mock = MockTransport::new();
for msg in responses {
mock.push_message(&msg).unwrap();
}
let session = LspSession::with_transport(Box::new(mock));
AgentHandle::new(session, "rust".into(), false)
}
fn mock_handle_compressed(responses: Vec<LspMessage>) -> AgentHandle {
let mock = MockTransport::new();
for msg in responses {
mock.push_message(&msg).unwrap();
}
let session = LspSession::with_transport(Box::new(mock));
AgentHandle::new(session, "rust".into(), true)
}
#[tokio::test]
async fn test_builder_missing_backend() {
let err = AgentHandle::builder()
.language("rust")
.start()
.await
.unwrap_err();
assert!(err.to_string().contains("backend"), "{err}");
}
#[tokio::test]
async fn test_builder_missing_language() {
let err = AgentHandle::builder()
.backend("rust-analyzer")
.start()
.await
.unwrap_err();
assert!(err.to_string().contains("language"), "{err}");
}
#[tokio::test]
async fn test_get_diagnostics_no_compression() {
let diag_notif = LspMessage::Notification {
method: "textDocument/publishDiagnostics".into(),
params: json!({
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "test error",
}],
}),
};
let mut agent = mock_handle(vec![diag_notif]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_diagnostics(&uri).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["diagnostics"][0]["message"], "test error");
}
#[tokio::test]
async fn test_get_diagnostics_with_compression() {
let diag_notif = LspMessage::Notification {
method: "textDocument/publishDiagnostics".into(),
params: json!({
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "test error",
}],
}),
};
let mut agent = mock_handle_compressed(vec![diag_notif]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_diagnostics(&uri).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert!(parsed.get("r").is_some() || parsed.get("diagnostics").is_some());
}
#[tokio::test]
async fn test_uri_must_be_file() {
let mut agent = mock_handle(vec![]);
let err = agent
.get_diagnostics("http://example.com/test.rs")
.await
.unwrap_err();
assert!(
err.to_string().contains("URI must start with file://"),
"{err}"
);
}
#[tokio::test]
async fn test_get_completions() {
let comp_resp = LspMessage::Response {
id: 1,
result: Some(json!({
"items": [
{ "label": "fn", "kind": 14 },
{ "label": "for", "kind": 14 },
],
})),
error: None,
};
let mut agent = mock_handle(vec![comp_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_completions(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["items"][0]["label"], "fn");
}
#[tokio::test]
async fn test_get_symbols() {
let sym_resp = LspMessage::Response {
id: 1,
result: Some(json!([
{ "name": "main", "kind": 12 },
])),
error: None,
};
let mut agent = mock_handle(vec![sym_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_symbols(&uri).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed[0]["name"], "main");
}
#[tokio::test]
async fn test_get_hover() {
let hover_resp = LspMessage::Response {
id: 1,
result: Some(json!({
"contents": {
"kind": "markdown",
"value": "**fn main** — Entry point",
},
})),
error: None,
};
let mut agent = mock_handle(vec![hover_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_hover(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["contents"]["value"], "**fn main** — Entry point");
}
#[tokio::test]
async fn test_get_references() {
let ref_resp = LspMessage::Response {
id: 1,
result: Some(json!([
{
"uri": "file:///lib.rs",
"range": { "start": { "line": 5, "character": 0 }, "end": { "line": 5, "character": 1 } },
}
])),
error: None,
};
let mut agent = mock_handle(vec![ref_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_references(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed[0]["uri"], "file:///lib.rs");
}
#[tokio::test]
async fn test_get_definition() {
let def_resp = LspMessage::Response {
id: 1,
result: Some(json!({
"uri": "file:///src/lib.rs",
"range": { "start": { "line": 1, "character": 0 }, "end": { "line": 1, "character": 10 } },
})),
error: None,
};
let mut agent = mock_handle(vec![def_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_definition(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["uri"], "file:///src/lib.rs");
}
#[tokio::test]
async fn test_get_implementation() {
let impl_resp = LspMessage::Response {
id: 1,
result: Some(json!([
{
"uri": "file:///src/impl.rs",
"range": { "start": { "line": 10, "character": 0 }, "end": { "line": 10, "character": 5 } },
}
])),
error: None,
};
let mut agent = mock_handle(vec![impl_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_implementation(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed[0]["uri"], "file:///src/impl.rs");
}
#[tokio::test]
async fn test_get_type_definition() {
let td_resp = LspMessage::Response {
id: 1,
result: Some(json!({
"uri": "file:///src/types.rs",
"range": { "start": { "line": 3, "character": 0 }, "end": { "line": 3, "character": 8 } },
})),
error: None,
};
let mut agent = mock_handle(vec![td_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_type_definition(&uri, 0, 0).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["uri"], "file:///src/types.rs");
}
#[tokio::test]
async fn test_get_workspace_symbols() {
let ws_resp = LspMessage::Response {
id: 1,
result: Some(json!([
{ "name": "main", "kind": 12, "location": {
"uri": "file:///src/main.rs",
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 0, "character": 1 } },
}},
])),
error: None,
};
let mut agent = mock_handle(vec![ws_resp]);
let result = agent.get_workspace_symbols("main").await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed[0]["name"], "main");
}
#[tokio::test]
async fn test_get_workspace_diagnostics() {
let wd_resp = LspMessage::Response {
id: 1,
result: Some(json!({
"kind": "full",
"items": [{
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "ws diag",
}],
}],
"resultId": "abc",
})),
error: None,
};
let mut agent = mock_handle(vec![wd_resp]);
let (uri, _path) = temp_file("fn main() {}");
let result = agent.get_workspace_diagnostics(&uri).await.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert!(parsed.get("items").is_some() || parsed.get("resultId").is_some());
}
#[test]
fn test_compress() {
let raw = json!({
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "test",
}],
});
let compact = AgentHandle::compress(&raw.to_string()).unwrap();
let parsed: Value = serde_json::from_str(&compact).unwrap();
assert!(parsed["diagnostics"][0].get("r").is_some() || parsed.get("r").is_some());
}
#[test]
fn test_inflate() {
let raw = json!({
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "test error",
}],
});
let compact = AgentHandle::compress(&raw.to_string()).unwrap();
let expanded = AgentHandle::inflate(&compact).unwrap();
let parsed: Value = serde_json::from_str(&expanded).unwrap();
assert_eq!(parsed["diagnostics"][0]["message"], "test error");
}
#[test]
fn test_inflate_compress_roundtrip() {
let raw = json!({
"uri": "file:///test.rs",
"diagnostics": [{
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 1, "character": 0 } },
"severity": 1,
"message": "roundtrip test",
}],
});
let raw_str = raw.to_string();
let compact = AgentHandle::compress(&raw_str).unwrap();
let expanded = AgentHandle::inflate(&compact).unwrap();
let expanded_val: Value = serde_json::from_str(&expanded).unwrap();
assert_eq!(expanded_val["diagnostics"][0]["message"], "roundtrip test");
}
#[tokio::test]
async fn test_shutdown() {
let shutdown_resp = LspMessage::Response {
id: 1,
result: Some(json!(null)),
error: None,
};
let agent = mock_handle(vec![shutdown_resp]);
agent.shutdown().await.unwrap();
}
}