Skip to main content

lean_ctx/engine/
mod.rs

1use std::path::PathBuf;
2use std::sync::atomic::{AtomicI64, Ordering};
3
4use anyhow::{anyhow, Context, Result};
5use rmcp::{
6    model::{
7        CallToolRequest, CallToolRequestParams, CallToolResult, ClientJsonRpcMessage,
8        ClientRequest, JsonRpcRequest, NumberOrString, ServerJsonRpcMessage, ServerResult,
9    },
10    service::serve_directly,
11    service::RoleServer,
12    transport::OneshotTransport,
13};
14use serde_json::{Map, Value};
15
16use crate::tools::LeanCtxServer;
17
18pub struct ContextEngine {
19    server: LeanCtxServer,
20    next_id: AtomicI64,
21}
22
23impl ContextEngine {
24    pub fn new() -> Self {
25        Self {
26            server: LeanCtxServer::new(),
27            next_id: AtomicI64::new(1),
28        }
29    }
30
31    pub fn with_project_root(project_root: impl Into<PathBuf>) -> Self {
32        Self {
33            server: LeanCtxServer::new_with_project_root(Some(
34                project_root.into().to_string_lossy().to_string(),
35            )),
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        crate::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_result(
61        &self,
62        name: &str,
63        arguments: Option<Value>,
64    ) -> Result<CallToolResult> {
65        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
66        let req_id = NumberOrString::Number(id);
67
68        let args_obj: Map<String, Value> = match arguments {
69            None => Map::new(),
70            Some(Value::Object(m)) => m,
71            Some(other) => {
72                return Err(anyhow!(
73                    "tool arguments must be a JSON object (got {})",
74                    other
75                ))
76            }
77        };
78
79        let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
80        let call: CallToolRequest = CallToolRequest::new(params);
81        let client_req = ClientRequest::CallToolRequest(call);
82        let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));
83
84        let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
85        let service = serve_directly(self.server.clone(), transport, None);
86        tokio::spawn(async move {
87            let _ = service.waiting().await;
88        });
89
90        let Some(server_msg) = rx.recv().await else {
91            return Err(anyhow!("no response from tool call"));
92        };
93
94        match server_msg {
95            ServerJsonRpcMessage::Response(r) => match r.result {
96                ServerResult::CallToolResult(result) => Ok(result),
97                other => Err(anyhow!("unexpected server result: {:?}", other)),
98            },
99            ServerJsonRpcMessage::Error(e) => Err(anyhow!("{e:?}")).context("tool call error"),
100            ServerJsonRpcMessage::Notification(_) => Err(anyhow!("unexpected notification")),
101            ServerJsonRpcMessage::Request(_) => Err(anyhow!("unexpected request")),
102        }
103    }
104
105    pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
106        let result = self.call_tool_result(name, arguments).await?;
107        let mut out = String::new();
108        for c in result.content {
109            if let Some(t) = c.as_text() {
110                out.push_str(&t.text);
111            }
112        }
113        if out.is_empty() {
114            if let Some(v) = result.structured_content {
115                out = v.to_string();
116            }
117        }
118        Ok(out)
119    }
120}
121
122impl Default for ContextEngine {
123    fn default() -> Self {
124        Self::new()
125    }
126}