nebu_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::NebuCtxServer;
17
18pub struct ContextEngine {
19    server: NebuCtxServer,
20    next_id: AtomicI64,
21}
22
23impl ContextEngine {
24    pub fn new() -> Self {
25        Self {
26            server: NebuCtxServer::new(),
27            next_id: AtomicI64::new(1),
28        }
29    }
30
31    pub fn with_project_root(project_root: impl Into<PathBuf>) -> Self {
32        let project_root = crate::core::pathutil::safe_canonicalize_or_self(&project_root.into())
33            .to_string_lossy()
34            .to_string();
35        Self {
36            server: NebuCtxServer::new_with_project_root(Some(project_root)),
37            next_id: AtomicI64::new(1),
38        }
39    }
40
41    pub fn from_server(server: NebuCtxServer) -> Self {
42        Self {
43            server,
44            next_id: AtomicI64::new(1),
45        }
46    }
47
48    pub fn server(&self) -> &NebuCtxServer {
49        &self.server
50    }
51
52    pub fn manifest(&self) -> Value {
53        crate::core::mcp_manifest::manifest_value()
54    }
55
56    pub async fn call_tool_value(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
57        let result = self.call_tool_result(name, arguments).await?;
58        serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
59    }
60
61    pub async fn call_tool_result(
62        &self,
63        name: &str,
64        arguments: Option<Value>,
65    ) -> Result<CallToolResult> {
66        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
67        let req_id = NumberOrString::Number(id);
68
69        let args_obj: Map<String, Value> = match arguments {
70            None => Map::new(),
71            Some(Value::Object(m)) => m,
72            Some(other) => {
73                return Err(anyhow!(
74                    "tool arguments must be a JSON object (got {})",
75                    other
76                ))
77            }
78        };
79
80        let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
81        let call: CallToolRequest = CallToolRequest::new(params);
82        let client_req = ClientRequest::CallToolRequest(call);
83        let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));
84
85        let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
86        let service = serve_directly(self.server.clone(), transport, None);
87        tokio::spawn(async move {
88            let _ = service.waiting().await;
89        });
90
91        let Some(server_msg) = rx.recv().await else {
92            return Err(anyhow!("no response from tool call"));
93        };
94
95        match server_msg {
96            ServerJsonRpcMessage::Response(r) => match r.result {
97                ServerResult::CallToolResult(result) => Ok(result),
98                other => Err(anyhow!("unexpected server result: {:?}", other)),
99            },
100            ServerJsonRpcMessage::Error(e) => Err(anyhow!("{e:?}")).context("tool call error"),
101            ServerJsonRpcMessage::Notification(_) => Err(anyhow!("unexpected notification")),
102            ServerJsonRpcMessage::Request(_) => Err(anyhow!("unexpected request")),
103        }
104    }
105
106    pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
107        let result = self.call_tool_result(name, arguments).await?;
108        let mut out = String::new();
109        for c in result.content {
110            if let Some(t) = c.as_text() {
111                out.push_str(&t.text);
112            }
113        }
114        if out.is_empty() {
115            if let Some(v) = result.structured_content {
116                out = v.to_string();
117            }
118        }
119        Ok(out)
120    }
121}
122
123impl Default for ContextEngine {
124    fn default() -> Self {
125        Self::new()
126    }
127}
nebu_ctx/engine/mod.rs

nebu_ctx/engine/
mod.rs
