lean_ctx_server/engine/
mod.rs1use std::path::PathBuf;
3use std::sync::atomic::{AtomicI64, Ordering};
4
5use anyhow::{Context, Result, anyhow};
6use rmcp::{
7 model::{
8 CallToolRequest, CallToolRequestParams, CallToolResult, ClientJsonRpcMessage,
9 ClientRequest, JsonRpcRequest, NumberOrString, ServerJsonRpcMessage, ServerResult,
10 },
11 service::RoleServer,
12 service::serve_directly,
13 transport::OneshotTransport,
14};
15use serde_json::{Map, Value};
16
17use crate::tools::LeanCtxServer;
18
19pub struct ContextEngine {
20 server: LeanCtxServer,
21 next_id: AtomicI64,
22}
23
24impl ContextEngine {
25 pub fn new() -> Self {
26 Self {
27 server: LeanCtxServer::new(),
28 next_id: AtomicI64::new(1),
29 }
30 }
31
32 pub fn with_project_root(project_root: impl Into<PathBuf>) -> Self {
33 let root = project_root.into().to_string_lossy().to_string();
34 Self {
35 server: LeanCtxServer::new_with_project_root(Some(&root)),
36 next_id: AtomicI64::new(1),
37 }
38 }
39
40 pub fn from_server(server: LeanCtxServer) -> Self {
41 Self {
42 server,
43 next_id: AtomicI64::new(1),
44 }
45 }
46
47 pub fn server(&self) -> &LeanCtxServer {
48 &self.server
49 }
50
51 pub fn manifest(&self) -> Value {
52 lean_ctx::core::mcp_manifest::manifest_value()
53 }
54
55 pub async fn call_tool_value(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
56 let result = self.call_tool_result(name, arguments).await?;
57 serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
58 }
59
60 pub async fn call_tool_value_with_rights(
64 &self,
65 name: &str,
66 arguments: Option<Value>,
67 rights: lean_ctx::core::session_rights::State,
68 ) -> Result<Value> {
69 let result = self.call_tool_bound(name, arguments, Some(rights)).await?;
70 serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
71 }
72
73 pub async fn call_tool_result(
74 &self,
75 name: &str,
76 arguments: Option<Value>,
77 ) -> Result<CallToolResult> {
78 self.call_tool_bound(name, arguments, None).await
79 }
80
81 async fn call_tool_bound(
82 &self,
83 name: &str,
84 arguments: Option<Value>,
85 rights: Option<lean_ctx::core::session_rights::State>,
86 ) -> Result<CallToolResult> {
87 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
88 let req_id = NumberOrString::Number(id);
89
90 let args_obj: Map<String, Value> = match arguments {
91 None => Map::new(),
92 Some(Value::Object(m)) => m,
93 Some(other) => {
94 return Err(anyhow!(
95 "tool arguments must be a JSON object (got {other})"
96 ));
97 }
98 };
99
100 let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
101 let mut call: CallToolRequest = CallToolRequest::new(params);
102 if let Some(rights) = rights {
103 call.extensions
105 .insert(lean_ctx::core::session_rights::RequestRights(rights));
106 }
107 let client_req = ClientRequest::CallToolRequest(call);
108 let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));
109
110 let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
111 let service = serve_directly(self.server.clone(), transport, None);
112 tokio::spawn(async move {
113 let _ = service.waiting().await;
114 });
115 let deadline = tokio::time::Instant::now() + std::time::Duration::from_mins(2);
116 loop {
117 let server_msg = match tokio::time::timeout_at(deadline, rx.recv()).await {
118 Ok(Some(msg)) => msg,
119 Ok(None) => return Err(anyhow!("no response from tool call")),
120 Err(_) => return Err(anyhow!("tool call timed out after 120s")),
121 };
122
123 match server_msg {
124 ServerJsonRpcMessage::Notification(_) => {}
125 ServerJsonRpcMessage::Response(r) => {
126 return match r.result {
127 ServerResult::CallToolResult(result) => Ok(result),
128 other => Err(anyhow!("unexpected server result: {other:?}")),
129 };
130 }
131 ServerJsonRpcMessage::Error(e) => {
132 return Err(anyhow!("{e:?}")).context("tool call error");
133 }
134 ServerJsonRpcMessage::Request(_) => {
135 return Err(anyhow!("unexpected request"));
136 }
137 }
138 }
139 }
140
141 pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
142 let result = self.call_tool_result(name, arguments).await?;
143 let mut out = String::new();
144 for c in result.content {
145 if let Some(t) = c.as_text() {
146 out.push_str(&t.text);
147 }
148 }
149 if out.is_empty()
150 && let Some(v) = result.structured_content
151 {
152 out = v.to_string();
153 }
154 Ok(out)
155 }
156}
157
158impl Default for ContextEngine {
159 fn default() -> Self {
160 Self::new()
161 }
162}
163
164pub mod operators;
166
167#[cfg(test)]
168mod canary_sweep_tests;
169#[cfg(test)]
170mod transport_tests;