lean_ctx/engine/
mod.rs

1use std::path::PathBuf;
2use std::sync::atomic::{AtomicI64, Ordering};
3
4use anyhow::{Context, Result, anyhow};
5use rmcp::{
6    model::{
7        CallToolRequest, CallToolRequestParams, CallToolResult, ClientJsonRpcMessage,
8        ClientRequest, JsonRpcRequest, NumberOrString, ServerJsonRpcMessage, ServerResult,
9    },
10    service::RoleServer,
11    service::serve_directly,
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        let root = project_root.into().to_string_lossy().to_string();
33        Self {
34            server: LeanCtxServer::new_with_project_root(Some(&root)),
35            next_id: AtomicI64::new(1),
36        }
37    }
38
39    pub fn from_server(server: LeanCtxServer) -> Self {
40        Self {
41            server,
42            next_id: AtomicI64::new(1),
43        }
44    }
45
46    pub fn server(&self) -> &LeanCtxServer {
47        &self.server
48    }
49
50    pub fn manifest(&self) -> Value {
51        crate::core::mcp_manifest::manifest_value()
52    }
53
54    pub async fn call_tool_value(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
55        let result = self.call_tool_result(name, arguments).await?;
56        serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
57    }
58
59    pub async fn call_tool_result(
60        &self,
61        name: &str,
62        arguments: Option<Value>,
63    ) -> Result<CallToolResult> {
64        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
65        let req_id = NumberOrString::Number(id);
66
67        let args_obj: Map<String, Value> = match arguments {
68            None => Map::new(),
69            Some(Value::Object(m)) => m,
70            Some(other) => {
71                return Err(anyhow!(
72                    "tool arguments must be a JSON object (got {other})"
73                ));
74            }
75        };
76
77        let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
78        let call: CallToolRequest = CallToolRequest::new(params);
79        let client_req = ClientRequest::CallToolRequest(call);
80        let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));
81
82        let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
83        let service = serve_directly(self.server.clone(), transport, None);
84        tokio::spawn(async move {
85            let _ = service.waiting().await;
86        });
87        let deadline = tokio::time::Instant::now() + std::time::Duration::from_mins(2);
88        loop {
89            let server_msg = match tokio::time::timeout_at(deadline, rx.recv()).await {
90                Ok(Some(msg)) => msg,
91                Ok(None) => return Err(anyhow!("no response from tool call")),
92                Err(_) => return Err(anyhow!("tool call timed out after 120s")),
93            };
94
95            match server_msg {
96                ServerJsonRpcMessage::Notification(_) => {}
97                ServerJsonRpcMessage::Response(r) => {
98                    return match r.result {
99                        ServerResult::CallToolResult(result) => Ok(result),
100                        other => Err(anyhow!("unexpected server result: {other:?}")),
101                    };
102                }
103                ServerJsonRpcMessage::Error(e) => {
104                    return Err(anyhow!("{e:?}")).context("tool call error");
105                }
106                ServerJsonRpcMessage::Request(_) => {
107                    return Err(anyhow!("unexpected request"));
108                }
109            }
110        }
111    }
112
113    pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
114        let result = self.call_tool_result(name, arguments).await?;
115        let mut out = String::new();
116        for c in result.content {
117            if let Some(t) = c.as_text() {
118                out.push_str(&t.text);
119            }
120        }
121        if out.is_empty()
122            && let Some(v) = result.structured_content
123        {
124            out = v.to_string();
125        }
126        Ok(out)
127    }
128}
129
130impl Default for ContextEngine {
131    fn default() -> Self {
132        Self::new()
133    }
134}
lean_ctx/engine/mod.rs

lean_ctx/engine/
mod.rs
