Skip to main content

lean_ctx_server/engine/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2use 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    /// [`Self::call_tool_value`] under a caller's session rights (process-v1
61    /// ยง9): the shared daemon runs a forwarded call with the rights of the
62    /// session that sent it, not with its own.
63    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            // In-process only: extensions never travel over the wire.
104            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
164/// Stable Rust-native boundary over the existing Local operators.
165pub mod operators;
166
167#[cfg(test)]
168mod canary_sweep_tests;
169#[cfg(test)]
170mod transport_tests;