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 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
88 let Some(server_msg) = rx.recv().await else {
89 return Err(anyhow!("no response from tool call"));
90 };
91
92 match server_msg {
93 ServerJsonRpcMessage::Response(r) => match r.result {
94 ServerResult::CallToolResult(result) => Ok(result),
95 other => Err(anyhow!("unexpected server result: {other:?}")),
96 },
97 ServerJsonRpcMessage::Error(e) => Err(anyhow!("{e:?}")).context("tool call error"),
98 ServerJsonRpcMessage::Notification(_) => Err(anyhow!("unexpected notification")),
99 ServerJsonRpcMessage::Request(_) => Err(anyhow!("unexpected request")),
100 }
101 }
102
103 pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
104 let result = self.call_tool_result(name, arguments).await?;
105 let mut out = String::new();
106 for c in result.content {
107 if let Some(t) = c.as_text() {
108 out.push_str(&t.text);
109 }
110 }
111 if out.is_empty() {
112 if let Some(v) = result.structured_content {
113 out = v.to_string();
114 }
115 }
116 Ok(out)
117 }
118}
119
120impl Default for ContextEngine {
121 fn default() -> Self {
122 Self::new()
123 }
124}