Compare commits
22
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7dc0c329ad | ||
|
|
ae0ea9ab22 | ||
|
|
bf7bf48163 | ||
|
|
d3fce30ca5 | ||
|
|
58ad2e76a3 | ||
|
|
4c4150a042 | ||
|
|
42d3638cf8 | ||
|
|
1af21b3aec | ||
|
|
8934857635 | ||
|
|
1d5e4bd00d | ||
|
|
84ad64156c | ||
|
|
075bc50645 | ||
|
|
891bd481b9 | ||
|
|
1f40e6d32b | ||
|
|
8e10950fc0 | ||
|
|
f8d98a77fd | ||
|
|
9476bcd649 | ||
|
|
6b37799cc5 | ||
|
|
3c31aeec1f | ||
|
|
5da422d23a | ||
|
|
73ab5037d8 | ||
|
|
8afbf97b11 |
No files matched your search
Binary file not shown.
@@ -0,0 +1,14 @@
|
|||||||
|
import re
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "r", encoding="utf-8") as f:
|
||||||
|
text = f.read()
|
||||||
|
|
||||||
|
pattern = r'"(list_milestones|list_pinned_files|read_handoff_memos)" => \{\s*let req = parse_tool!\(args, id, ([a-zA-Z]+)\);\s*let mut [a-zA-Z_]+ = self\.state\.([a-zA-Z_]+)\.read\(\);\s*if let Some\(ns\) = req\.namespace \{\s*[a-zA-Z_]+\.retain\(\|.\| \w+\.namespace == ns\);\s*\}\s*let data = serde_json::to_string\(&[a-zA-Z_]+\)\.unwrap_or_default\(\);\s*Ok\(data\.to_string\(\)\)\s*\}'
|
||||||
|
|
||||||
|
def repl(m):
|
||||||
|
return f'"{m.group(1)}" => handle_list_with_namespace!(self, {m.group(3)}, {m.group(2)}, args, id),'
|
||||||
|
|
||||||
|
new_text = re.sub(pattern, repl, text)
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "w", encoding="utf-8") as f:
|
||||||
|
f.write(new_text)
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
import re
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "r", encoding="utf-8") as f:
|
||||||
|
text = f.read()
|
||||||
|
|
||||||
|
pattern = r'"([a-zA-Z_]+)" => \{\s*(?:let req = parse_tool!\(args, id, ([a-zA-Z]+)\);\s*)?let data = serde_json::to_string\(&self\.state\.([a-zA-Z_]+)\.read\(\)\)\s*\.unwrap_or_else\(\|_\| "\[\]"\.to_string\(\)\);\s*Ok\(data\.to_string\(\)\)\s*\}'
|
||||||
|
|
||||||
|
def repl(m):
|
||||||
|
return f'"{m.group(1)}" => {{\n let data = serde_json::to_string(&self.state.{m.group(3)}.read()).unwrap_or_else(|_| "[]".to_string());\n Ok(data)\n}},'
|
||||||
|
|
||||||
|
new_text = re.sub(pattern, repl, text)
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "w", encoding="utf-8") as f:
|
||||||
|
f.write(new_text)
|
||||||
+2
-6
@@ -15,10 +15,6 @@ fn main() {
|
|||||||
.and_then(|out| String::from_utf8(out.stdout).ok())
|
.and_then(|out| String::from_utf8(out.stdout).ok())
|
||||||
.unwrap_or_else(|| "unknown".to_string());
|
.unwrap_or_else(|| "unknown".to_string());
|
||||||
|
|
||||||
let version = format!(
|
let version = format!("{} ({})", git_date.trim(), git_hash.trim());
|
||||||
"{} ({})",
|
println!("cargo:rustc-env=APP_VERSION={version}");
|
||||||
git_date.trim(),
|
|
||||||
git_hash.trim()
|
|
||||||
);
|
|
||||||
println!("cargo:rustc-env=APP_VERSION={}", version);
|
|
||||||
}
|
}
|
||||||
@@ -1,3 +1,5 @@
|
|||||||
|
#![cfg(unix)]
|
||||||
|
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::io::{BufRead, BufReader, Read, Write};
|
use std::io::{BufRead, BufReader, Read, Write};
|
||||||
use std::process::{Command, Stdio};
|
use std::process::{Command, Stdio};
|
||||||
|
|||||||
@@ -22,9 +22,8 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
let lower_line = line.to_lowercase();
|
if line.len() >= 15 && line[..15].eq_ignore_ascii_case("content-length:") {
|
||||||
if let Some(len_str) = lower_line.strip_prefix("content-length:") {
|
length = line[15..].trim().parse().unwrap_or(0);
|
||||||
length = len_str.trim().parse().unwrap_or(0);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -39,3 +38,4 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
|
|||||||
|
|
||||||
String::from_utf8(buffer).ok()
|
String::from_utf8(buffer).ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
+142
-96
@@ -20,10 +20,8 @@ pub struct JsonRpcResponse {
|
|||||||
pub error: Option<Value>,
|
pub error: Option<Value>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
pub async fn send_response(response: JsonRpcResponse) {
|
pub async fn send_response(response: JsonRpcResponse) {
|
||||||
let msg = serde_json::to_string(&response).unwrap();
|
let msg = serde_json::to_string(&response).unwrap_or_else(|_| "{}".to_string());
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Sending JSON-RPC response (id: {:?}): {}",
|
"Sending JSON-RPC response (id: {:?}): {}",
|
||||||
response.id,
|
response.id,
|
||||||
@@ -57,7 +55,7 @@ pub async fn send_error(id: Value, code: i32, message: &str) {
|
|||||||
async fn get_socket_path() -> Result<String, String> {
|
async fn get_socket_path() -> Result<String, String> {
|
||||||
let profile =
|
let profile =
|
||||||
std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
|
std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
|
||||||
let path = format!("{}\\.gemini\\active_nvim.txt", profile);
|
let path = format!("{profile}\\.gemini\\active_nvim.txt");
|
||||||
|
|
||||||
if let Ok(content) = tokio::fs::read_to_string(&path).await {
|
if let Ok(content) = tokio::fs::read_to_string(&path).await {
|
||||||
let p = content.trim().to_string();
|
let p = content.trim().to_string();
|
||||||
@@ -65,10 +63,9 @@ async fn get_socket_path() -> Result<String, String> {
|
|||||||
if p.starts_with(r"\\.\pipe\") {
|
if p.starts_with(r"\\.\pipe\") {
|
||||||
return Ok(p);
|
return Ok(p);
|
||||||
} else if p.starts_with("nvim.") {
|
} else if p.starts_with("nvim.") {
|
||||||
return Ok(format!(r"\\.\pipe\{}", p));
|
return Ok(format!(r"\\.\pipe\{p}"));
|
||||||
} else {
|
|
||||||
return Ok(p);
|
|
||||||
}
|
}
|
||||||
|
return Ok(p);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -78,7 +75,7 @@ async fn get_socket_path() -> Result<String, String> {
|
|||||||
let name = entry.file_name();
|
let name = entry.file_name();
|
||||||
let name_str = name.to_string_lossy();
|
let name_str = name.to_string_lossy();
|
||||||
if name_str.starts_with("nvim.") {
|
if name_str.starts_with("nvim.") {
|
||||||
return Ok(format!(r"\\.\pipe\{}", name_str));
|
return Ok(format!(r"\\.\pipe\{name_str}"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -111,11 +108,10 @@ async fn get_socket_path() -> Result<String, String> {
|
|||||||
}
|
}
|
||||||
Err("Could not find Neovim socket".to_string())
|
Err("Could not find Neovim socket".to_string())
|
||||||
}
|
}
|
||||||
use std::sync::LazyLock;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::sync::{mpsc, oneshot};
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::LazyLock;
|
||||||
|
use tokio::sync::{mpsc, oneshot};
|
||||||
|
|
||||||
pub struct NvimRequest {
|
pub struct NvimRequest {
|
||||||
pub msgid_str: String,
|
pub msgid_str: String,
|
||||||
@@ -123,11 +119,14 @@ pub struct NvimRequest {
|
|||||||
pub reply: oneshot::Sender<Result<rmpv::Value, String>>,
|
pub reply: oneshot::Sender<Result<rmpv::Value, String>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
static NVIM_CONN: LazyLock<Arc<std::sync::Mutex<Option<mpsc::Sender<NvimRequest>>>>> = LazyLock::new(|| Arc::new(std::sync::Mutex::new(None)));
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
static NEXT_MSGID: AtomicU64 = AtomicU64::new(1);
|
||||||
|
static NVIM_CONN: LazyLock<Arc<std::sync::Mutex<Option<mpsc::Sender<NvimRequest>>>>> =
|
||||||
|
LazyLock::new(|| Arc::new(std::sync::Mutex::new(None)));
|
||||||
|
|
||||||
async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
||||||
{
|
{
|
||||||
let conn_lock = NVIM_CONN.lock().unwrap();
|
let conn_lock = NVIM_CONN.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
if let Some(sender) = conn_lock.as_ref() {
|
if let Some(sender) = conn_lock.as_ref() {
|
||||||
if !sender.is_closed() {
|
if !sender.is_closed() {
|
||||||
return Ok(sender.clone());
|
return Ok(sender.clone());
|
||||||
@@ -141,18 +140,23 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
#[cfg(windows)]
|
#[cfg(windows)]
|
||||||
let stream = {
|
let stream = {
|
||||||
use tokio::net::windows::named_pipe::ClientOptions;
|
use tokio::net::windows::named_pipe::ClientOptions;
|
||||||
ClientOptions::new().open(&socket_path).map_err(|e| e.to_string())?
|
ClientOptions::new()
|
||||||
|
.open(&socket_path)
|
||||||
|
.map_err(|e| e.to_string())?
|
||||||
};
|
};
|
||||||
|
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
let stream = {
|
let stream = {
|
||||||
use tokio::net::UnixStream;
|
use tokio::net::UnixStream;
|
||||||
UnixStream::connect(socket_path).await.map_err(|e| e.to_string())?
|
UnixStream::connect(socket_path)
|
||||||
|
.await
|
||||||
|
.map_err(|e| e.to_string())?
|
||||||
};
|
};
|
||||||
|
|
||||||
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
||||||
let (tx, mut rx) = mpsc::channel::<NvimRequest>(32);
|
let (tx, mut rx) = mpsc::channel::<NvimRequest>(32);
|
||||||
type PendingRequestsMap = Arc<std::sync::Mutex<HashMap<String, oneshot::Sender<Result<rmpv::Value, String>>>>>;
|
type PendingRequestsMap =
|
||||||
|
Arc<std::sync::Mutex<HashMap<String, oneshot::Sender<Result<rmpv::Value, String>>>>>;
|
||||||
let pending_requests: PendingRequestsMap = Arc::new(std::sync::Mutex::new(HashMap::new()));
|
let pending_requests: PendingRequestsMap = Arc::new(std::sync::Mutex::new(HashMap::new()));
|
||||||
|
|
||||||
// Write task
|
// Write task
|
||||||
@@ -165,7 +169,10 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
pending_clone.lock().unwrap().insert(req.msgid_str.clone(), req.reply);
|
pending_clone
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.insert(req.msgid_str.clone(), req.reply);
|
||||||
|
|
||||||
if write_half.write_all(&buf).await.is_err() {
|
if write_half.write_all(&buf).await.is_err() {
|
||||||
tracing::error!("Failed to write to Neovim socket");
|
tracing::error!("Failed to write to Neovim socket");
|
||||||
@@ -190,9 +197,11 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
if let rmpv::Value::Array(ref arr) = val {
|
if let rmpv::Value::Array(ref arr) = val {
|
||||||
if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) {
|
if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) {
|
||||||
let msgid = &arr[1];
|
let msgid = &arr[1];
|
||||||
let msgid_str = format!("{:?}", msgid);
|
let msgid_str = format!("{msgid:?}");
|
||||||
|
|
||||||
if let Some(reply_sender) = pending_clone2.lock().unwrap().remove(&msgid_str) {
|
if let Some(reply_sender) =
|
||||||
|
pending_clone2.lock().unwrap_or_else(|e| e.into_inner()).remove(&msgid_str)
|
||||||
|
{
|
||||||
let _ = reply_sender.send(Ok(val));
|
let _ = reply_sender.send(Ok(val));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -204,12 +213,16 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
}
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
Err(rmpv::decode::Error::InvalidMarkerRead(e)) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
Err(rmpv::decode::Error::InvalidMarkerRead(e))
|
||||||
|
if e.kind() == std::io::ErrorKind::UnexpectedEof =>
|
||||||
|
{
|
||||||
resp_buf.drain(..offset);
|
resp_buf.drain(..offset);
|
||||||
offset = 0;
|
offset = 0;
|
||||||
|
|
||||||
let read_future = read_half.read(&mut chunk);
|
let read_future = read_half.read(&mut chunk);
|
||||||
match tokio::time::timeout(tokio::time::Duration::from_secs(60), read_future).await {
|
match tokio::time::timeout(tokio::time::Duration::from_secs(60), read_future)
|
||||||
|
.await
|
||||||
|
{
|
||||||
Ok(Ok(n)) if n > 0 => {
|
Ok(Ok(n)) if n > 0 => {
|
||||||
resp_buf.extend_from_slice(&chunk[..n]);
|
resp_buf.extend_from_slice(&chunk[..n]);
|
||||||
}
|
}
|
||||||
@@ -227,7 +240,7 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Cleanup pending requests on disconnect
|
// Cleanup pending requests on disconnect
|
||||||
let mut pending = pending_clone2.lock().unwrap();
|
let mut pending = pending_clone2.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
for (_, sender) in pending.drain() {
|
for (_, sender) in pending.drain() {
|
||||||
let _ = sender.send(Err("Connection closed".to_string()));
|
let _ = sender.send(Err("Connection closed".to_string()));
|
||||||
}
|
}
|
||||||
@@ -242,11 +255,14 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
|||||||
if Arc::strong_count(&pending_clone3) <= 1 {
|
if Arc::strong_count(&pending_clone3) <= 1 {
|
||||||
break; // Socket closed and other tasks finished, no need to keep cleaning up
|
break; // Socket closed and other tasks finished, no need to keep cleaning up
|
||||||
}
|
}
|
||||||
pending_clone3.lock().unwrap().retain(|_, sender| !sender.is_closed());
|
pending_clone3
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.retain(|_, sender| !sender.is_closed());
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
let mut conn_lock = NVIM_CONN.lock().unwrap();
|
let mut conn_lock = NVIM_CONN.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
if let Some(existing_sender) = conn_lock.as_ref() {
|
if let Some(existing_sender) = conn_lock.as_ref() {
|
||||||
if !existing_sender.is_closed() {
|
if !existing_sender.is_closed() {
|
||||||
// Another task established the connection while we were waiting
|
// Another task established the connection while we were waiting
|
||||||
@@ -268,7 +284,7 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
|
|||||||
rmpv::Value::Nil
|
rmpv::Value::Nil
|
||||||
};
|
};
|
||||||
|
|
||||||
let msgid_str = format!("{:?}", msgid);
|
let msgid_str = format!("{msgid:?}");
|
||||||
let tx = get_nvim_connection().await?;
|
let tx = get_nvim_connection().await?;
|
||||||
let (reply_tx, reply_rx) = oneshot::channel();
|
let (reply_tx, reply_rx) = oneshot::channel();
|
||||||
|
|
||||||
@@ -276,7 +292,9 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
|
|||||||
msgid_str,
|
msgid_str,
|
||||||
req,
|
req,
|
||||||
reply: reply_tx,
|
reply: reply_tx,
|
||||||
}).await.map_err(|_| "Failed to send request to Neovim connection manager")?;
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|_| "Failed to send request to Neovim connection manager")?;
|
||||||
|
|
||||||
match tokio::time::timeout(tokio::time::Duration::from_secs(10), reply_rx).await {
|
match tokio::time::timeout(tokio::time::Duration::from_secs(10), reply_rx).await {
|
||||||
Ok(Ok(res)) => res,
|
Ok(Ok(res)) => res,
|
||||||
@@ -287,9 +305,10 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
|
|||||||
|
|
||||||
async fn send_nvim_command(cmd: &str) -> Result<(), String> {
|
async fn send_nvim_command(cmd: &str) -> Result<(), String> {
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(1.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_command".into()),
|
RmpValue::String("nvim_command".into()),
|
||||||
RmpValue::Array(vec![RmpValue::String(cmd.into())]),
|
RmpValue::Array(vec![RmpValue::String(cmd.into())]),
|
||||||
]);
|
]);
|
||||||
@@ -306,9 +325,10 @@ async fn send_nvim_command(cmd: &str) -> Result<(), String> {
|
|||||||
|
|
||||||
async fn get_nvim_active_buffer() -> Result<String, String> {
|
async fn get_nvim_active_buffer() -> Result<String, String> {
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(2.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_buf_get_lines".into()),
|
RmpValue::String("nvim_buf_get_lines".into()),
|
||||||
RmpValue::Array(vec![
|
RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
@@ -341,9 +361,10 @@ async fn get_nvim_active_buffer() -> Result<String, String> {
|
|||||||
|
|
||||||
async fn get_nvim_cursor() -> Result<String, String> {
|
async fn get_nvim_cursor() -> Result<String, String> {
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(3.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_win_get_cursor".into()),
|
RmpValue::String("nvim_win_get_cursor".into()),
|
||||||
RmpValue::Array(vec![RmpValue::Integer(0.into())]),
|
RmpValue::Array(vec![RmpValue::Integer(0.into())]),
|
||||||
]);
|
]);
|
||||||
@@ -356,7 +377,7 @@ async fn get_nvim_cursor() -> Result<String, String> {
|
|||||||
if let RmpValue::Array(pos) = &arr[3] {
|
if let RmpValue::Array(pos) = &arr[3] {
|
||||||
if pos.len() == 2 {
|
if pos.len() == 2 {
|
||||||
if let (RmpValue::Integer(row), RmpValue::Integer(col)) = (&pos[0], &pos[1]) {
|
if let (RmpValue::Integer(row), RmpValue::Integer(col)) = (&pos[0], &pos[1]) {
|
||||||
return Ok(format!("Line: {}, Column: {}", row, col));
|
return Ok(format!("Line: {row}, Column: {col}"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -377,9 +398,10 @@ async fn get_nvim_visual_selection() -> Result<String, String> {
|
|||||||
"#;
|
"#;
|
||||||
|
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(4.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_exec_lua".into()),
|
RmpValue::String("nvim_exec_lua".into()),
|
||||||
RmpValue::Array(vec![
|
RmpValue::Array(vec![
|
||||||
RmpValue::String(lua_script.into()),
|
RmpValue::String(lua_script.into()),
|
||||||
@@ -402,25 +424,25 @@ async fn get_nvim_visual_selection() -> Result<String, String> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> {
|
async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> {
|
||||||
let escaped_message = message.replace("\\", "\\\\").replace("\"", "\\\"");
|
let escaped_message = message.replace('\\', "\\\\").replace('"', "\\\"");
|
||||||
let lua_script = format!(
|
let lua_script = format!(
|
||||||
r#"
|
r#"
|
||||||
local ns = vim.api.nvim_create_namespace("gemini_diagnostics")
|
local ns = vim.api.nvim_create_namespace("gemini_diagnostics")
|
||||||
local diagnostics = {{{{
|
local diagnostics = {{{{
|
||||||
lnum = {} - 1,
|
lnum = {line} - 1,
|
||||||
col = 0,
|
col = 0,
|
||||||
severity = vim.diagnostic.severity.WARN,
|
severity = vim.diagnostic.severity.WARN,
|
||||||
message = "{}",
|
message = "{escaped_message}",
|
||||||
}}}}
|
}}}}
|
||||||
vim.diagnostic.set(ns, 0, diagnostics, {{}})
|
vim.diagnostic.set(ns, 0, diagnostics, {{}})
|
||||||
"#,
|
"#
|
||||||
line, escaped_message
|
|
||||||
);
|
);
|
||||||
|
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(5.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_exec_lua".into()),
|
RmpValue::String("nvim_exec_lua".into()),
|
||||||
RmpValue::Array(vec![
|
RmpValue::Array(vec![
|
||||||
RmpValue::String(lua_script.into()),
|
RmpValue::String(lua_script.into()),
|
||||||
@@ -470,7 +492,7 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
|
|||||||
let key_str = if let rmpv::Value::String(s) = k {
|
let key_str = if let rmpv::Value::String(s) = k {
|
||||||
s.as_str().unwrap_or("").to_string()
|
s.as_str().unwrap_or("").to_string()
|
||||||
} else {
|
} else {
|
||||||
format!("{:?}", k)
|
format!("{k:?}")
|
||||||
};
|
};
|
||||||
obj.insert(key_str, rmpv_to_json(v));
|
obj.insert(key_str, rmpv_to_json(v));
|
||||||
}
|
}
|
||||||
@@ -482,9 +504,10 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
|
|||||||
|
|
||||||
async fn execute_nvim_lua(code: &str) -> Result<String, String> {
|
async fn execute_nvim_lua(code: &str) -> Result<String, String> {
|
||||||
use rmpv::Value as RmpValue;
|
use rmpv::Value as RmpValue;
|
||||||
|
let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);
|
||||||
let req = RmpValue::Array(vec![
|
let req = RmpValue::Array(vec![
|
||||||
RmpValue::Integer(0.into()),
|
RmpValue::Integer(0.into()),
|
||||||
RmpValue::Integer(6.into()), // msgid
|
RmpValue::Integer(msgid.into()), // msgid
|
||||||
RmpValue::String("nvim_exec_lua".into()),
|
RmpValue::String("nvim_exec_lua".into()),
|
||||||
RmpValue::Array(vec![RmpValue::String(code.into()), RmpValue::Array(vec![])]),
|
RmpValue::Array(vec![RmpValue::String(code.into()), RmpValue::Array(vec![])]),
|
||||||
]);
|
]);
|
||||||
@@ -497,7 +520,7 @@ async fn execute_nvim_lua(code: &str) -> Result<String, String> {
|
|||||||
if arr.len() > 3 {
|
if arr.len() > 3 {
|
||||||
return Ok(serde_json::to_string_pretty(&rmpv_to_json(&arr[3])).unwrap_or_default());
|
return Ok(serde_json::to_string_pretty(&rmpv_to_json(&arr[3])).unwrap_or_default());
|
||||||
}
|
}
|
||||||
return Ok("".to_string());
|
return Ok(String::new());
|
||||||
}
|
}
|
||||||
Err("Invalid response".to_string())
|
Err("Invalid response".to_string())
|
||||||
}
|
}
|
||||||
@@ -532,21 +555,22 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
tracing::info!("{} MCP server started", app_name);
|
tracing::info!("{} MCP server started", app_name);
|
||||||
let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
|
let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
|
||||||
loop {
|
loop {
|
||||||
let raw_msg = match mcp_stdio::read_mcp_message(&mut stdin).await {
|
let raw_msg = if let Some(m) = mcp_stdio::read_mcp_message(&mut stdin).await { m } else {
|
||||||
Some(m) => m,
|
|
||||||
None => {
|
|
||||||
tracing::info!("Stdin closed, exiting loop");
|
tracing::info!("Stdin closed, exiting loop");
|
||||||
break;
|
break;
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let msg = match serde_json::from_str::<JsonRpcRequest>(&raw_msg) {
|
let msg = match serde_json::from_str::<JsonRpcRequest>(&raw_msg) {
|
||||||
Ok(m) => {
|
Ok(m) => {
|
||||||
tracing::info!("Received message method: {}", m.method);
|
tracing::info!("Received message method: {}", m.method);
|
||||||
m
|
m
|
||||||
},
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::error!("Failed to parse JSON-RPC request from JSONL: {}. Payload: {}", e, raw_msg);
|
tracing::error!(
|
||||||
|
"Failed to parse JSON-RPC request from JSONL: {}. Payload: {}",
|
||||||
|
e,
|
||||||
|
raw_msg
|
||||||
|
);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -559,13 +583,17 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
let id_clone = id.clone();
|
let id_clone = id.clone();
|
||||||
let start_time = std::time::Instant::now();
|
let start_time = std::time::Instant::now();
|
||||||
let method_clone = if msg.method == "tools/call" {
|
let method_clone = if msg.method == "tools/call" {
|
||||||
let tool_name = msg.params.as_ref().and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown");
|
let tool_name = msg
|
||||||
format!("ToolCall[{}]", tool_name)
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|p| p.get("name"))
|
||||||
|
.and_then(|n| n.as_str())
|
||||||
|
.unwrap_or("unknown");
|
||||||
|
format!("ToolCall[{tool_name}]")
|
||||||
} else {
|
} else {
|
||||||
msg.method.clone()
|
msg.method.clone()
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
||||||
match msg.method.as_str() {
|
match msg.method.as_str() {
|
||||||
|
|
||||||
"initialize" => {
|
"initialize" => {
|
||||||
@@ -661,7 +689,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"name": "nvim_execute_lua",
|
"name": "nvim_execute_lua",
|
||||||
"description": "Execute arbitrary Lua code in Neovim and return the result (JSON serialized).",
|
"description": "Execute arbitrary Lua code in Neovim and return the result (JSON serialized). CRITICAL: This is the primary tool for editing files. Use `vim.api.nvim_buf_set_lines` to modify buffers. ALWAYS trigger a notification (`require('notify')(...)`). ALWAYS auto-save (`vim.cmd('write')`) after edits. NEVER use interactive commands or 'c' flags in substitutions (e.g. use `vim.cmd('%s/old/new/ge')`).",
|
||||||
"inputSchema": {
|
"inputSchema": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
@@ -673,7 +701,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
]
|
]
|
||||||
})),
|
})),
|
||||||
error: None,
|
error: None,
|
||||||
}).await
|
}).await;
|
||||||
}
|
}
|
||||||
"tools/call" => {
|
"tools/call" => {
|
||||||
let params = msg.params.unwrap_or(json!({}));
|
let params = msg.params.unwrap_or(json!({}));
|
||||||
@@ -685,12 +713,12 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
"nvim_goto_line" => {
|
"nvim_goto_line" => {
|
||||||
if let (Some(file), Some(line)) = (
|
if let (Some(file), Some(line)) = (
|
||||||
args.get("file").and_then(|v| v.as_str()),
|
args.get("file").and_then(|v| v.as_str()),
|
||||||
args.get("line").and_then(|v| v.as_i64()),
|
args.get("line").and_then(serde_json::Value::as_i64),
|
||||||
) {
|
) {
|
||||||
let escaped_file = file.replace("\\", "\\\\");
|
let escaped_file = file.replace('\\', "\\\\");
|
||||||
let cmd = format!("e {} | {} | normal! zz", escaped_file, line);
|
let cmd = format!("e {escaped_file} | {line} | normal! zz");
|
||||||
match send_nvim_command(&cmd).await {
|
match send_nvim_command(&cmd).await {
|
||||||
Ok(_) => {
|
Ok(()) => {
|
||||||
send_text_result!(id.clone(), format!("Successfully jumped to {} line {}", file, line));
|
send_text_result!(id.clone(), format!("Successfully jumped to {} line {}", file, line));
|
||||||
}
|
}
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
@@ -719,11 +747,11 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
},
|
},
|
||||||
"nvim_set_diagnostics" => {
|
"nvim_set_diagnostics" => {
|
||||||
if let (Some(line), Some(message)) = (
|
if let (Some(line), Some(message)) = (
|
||||||
args.get("line").and_then(|v| v.as_i64()),
|
args.get("line").and_then(serde_json::Value::as_i64),
|
||||||
args.get("message").and_then(|v| v.as_str()),
|
args.get("message").and_then(|v| v.as_str()),
|
||||||
) {
|
) {
|
||||||
match set_nvim_diagnostics(line, message).await {
|
match set_nvim_diagnostics(line, message).await {
|
||||||
Ok(_) => {
|
Ok(()) => {
|
||||||
send_text_result!(id.clone(), format!("Successfully set diagnostic on line {}", line));
|
send_text_result!(id.clone(), format!("Successfully set diagnostic on line {}", line));
|
||||||
}
|
}
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
@@ -733,7 +761,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_list_buffers" => {
|
"nvim_list_buffers" => {
|
||||||
let lua_code = r#"
|
let lua_code = r"
|
||||||
local bufs = vim.api.nvim_list_bufs()
|
local bufs = vim.api.nvim_list_bufs()
|
||||||
local result = {}
|
local result = {}
|
||||||
for _, buf in ipairs(bufs) do
|
for _, buf in ipairs(bufs) do
|
||||||
@@ -743,7 +771,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
end
|
end
|
||||||
end
|
end
|
||||||
return vim.fn.json_encode(result)
|
return vim.fn.json_encode(result)
|
||||||
"#;
|
";
|
||||||
match execute_nvim_lua(lua_code).await {
|
match execute_nvim_lua(lua_code).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
send_text_result!(id.clone(), result);
|
send_text_result!(id.clone(), result);
|
||||||
@@ -752,7 +780,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_get_diagnostics" => {
|
"nvim_get_diagnostics" => {
|
||||||
let lua_code = r#"
|
let lua_code = r"
|
||||||
local bufnr = vim.api.nvim_get_current_buf()
|
local bufnr = vim.api.nvim_get_current_buf()
|
||||||
local diagnostics = vim.diagnostic.get(bufnr)
|
local diagnostics = vim.diagnostic.get(bufnr)
|
||||||
local result = {}
|
local result = {}
|
||||||
@@ -766,7 +794,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
})
|
})
|
||||||
end
|
end
|
||||||
return vim.fn.json_encode(result)
|
return vim.fn.json_encode(result)
|
||||||
"#;
|
";
|
||||||
match execute_nvim_lua(lua_code).await {
|
match execute_nvim_lua(lua_code).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
send_text_result!(id.clone(), result);
|
send_text_result!(id.clone(), result);
|
||||||
@@ -776,24 +804,24 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
"nvim_open_file" => {
|
"nvim_open_file" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
vim.cmd('edit ' .. vim.fn.fnameescape(args.file))
|
vim.cmd('edit ' .. vim.fn.fnameescape(args.file))
|
||||||
if args.filetype and args.filetype ~= '' then
|
if args.filetype and args.filetype ~= '' then
|
||||||
vim.bo.filetype = args.filetype
|
vim.bo.filetype = args.filetype
|
||||||
end
|
end
|
||||||
return 'Opened file ' .. args.file
|
return 'Opened file ' .. args.file
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_open_buffer" => {
|
"nvim_open_buffer" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local buf = vim.api.nvim_create_buf(true, true)
|
local buf = vim.api.nvim_create_buf(true, true)
|
||||||
if args.name and args.name ~= '' then
|
if args.name and args.name ~= '' then
|
||||||
pcall(vim.api.nvim_buf_set_name, buf, args.name)
|
pcall(vim.api.nvim_buf_set_name, buf, args.name)
|
||||||
@@ -807,30 +835,30 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
end
|
end
|
||||||
vim.api.nvim_win_set_buf(0, buf)
|
vim.api.nvim_win_set_buf(0, buf)
|
||||||
return 'Opened buffer ' .. tostring(buf)
|
return 'Opened buffer ' .. tostring(buf)
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_close_buffer" => {
|
"nvim_close_buffer" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||||
local force = args.force or false
|
local force = args.force or false
|
||||||
vim.api.nvim_buf_delete(buf, {{ force = force }})
|
vim.api.nvim_buf_delete(buf, {{ force = force }})
|
||||||
return 'Closed buffer ' .. tostring(buf)
|
return 'Closed buffer ' .. tostring(buf)
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_split_window" => {
|
"nvim_split_window" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local cmd = args.direction == 'horizontal' and 'split' or 'vsplit'
|
local cmd = args.direction == 'horizontal' and 'split' or 'vsplit'
|
||||||
vim.cmd(cmd)
|
vim.cmd(cmd)
|
||||||
if args.file and args.file ~= '' then
|
if args.file and args.file ~= '' then
|
||||||
@@ -839,16 +867,16 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
vim.api.nvim_win_set_buf(0, args.buf_id)
|
vim.api.nvim_win_set_buf(0, args.buf_id)
|
||||||
end
|
end
|
||||||
return 'Split window created'
|
return 'Split window created'
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_reload_buffer" => {
|
"nvim_reload_buffer" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||||
vim.api.nvim_buf_call(buf, function()
|
vim.api.nvim_buf_call(buf, function()
|
||||||
if args.force then
|
if args.force then
|
||||||
@@ -858,7 +886,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
end
|
end
|
||||||
end)
|
end)
|
||||||
return 'Reloaded buffer ' .. tostring(buf)
|
return 'Reloaded buffer ' .. tostring(buf)
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
@@ -875,9 +903,9 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_set_quickfix" => {
|
"nvim_set_quickfix" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local items = args.items or {{}}
|
local items = args.items or {{}}
|
||||||
local action = ' '
|
local action = ' '
|
||||||
if args.action == 'append' then action = 'a' end
|
if args.action == 'append' then action = 'a' end
|
||||||
@@ -886,16 +914,16 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
vim.fn.setqflist(items, action)
|
vim.fn.setqflist(items, action)
|
||||||
vim.cmd('copen')
|
vim.cmd('copen')
|
||||||
return 'Populated quickfix with ' .. tostring(#items) .. ' items'
|
return 'Populated quickfix with ' .. tostring(#items) .. ' items'
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_highlight_lines" => {
|
"nvim_highlight_lines" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||||
local group = args.group or 'IncSearch'
|
local group = args.group or 'IncSearch'
|
||||||
local ns = vim.api.nvim_create_namespace('antigravity_highlight')
|
local ns = vim.api.nvim_create_namespace('antigravity_highlight')
|
||||||
@@ -917,16 +945,16 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
end, duration)
|
end, duration)
|
||||||
end
|
end
|
||||||
return 'Highlighted lines ' .. tostring(args.start_line) .. ' to ' .. tostring(args.end_line)
|
return 'Highlighted lines ' .. tostring(args.start_line) .. ' to ' .. tostring(args.end_line)
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"nvim_get_messages" => {
|
"nvim_get_messages" => {
|
||||||
let json_str = serde_json::to_string(args).unwrap().replace("\\", "\\\\").replace("'", "\\'");
|
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||||
let code = format!("
|
let code = format!("
|
||||||
local args = vim.json.decode('{}')
|
local args = vim.json.decode('{json_str}')
|
||||||
local msg = vim.fn.execute('messages')
|
local msg = vim.fn.execute('messages')
|
||||||
local lines = vim.split(msg, '\\n')
|
local lines = vim.split(msg, '\\n')
|
||||||
if args.tail and args.tail > 0 and #lines > args.tail then
|
if args.tail and args.tail > 0 and #lines > args.tail then
|
||||||
@@ -937,7 +965,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
return table.concat(tail_lines, '\\n')
|
return table.concat(tail_lines, '\\n')
|
||||||
end
|
end
|
||||||
return msg
|
return msg
|
||||||
", json_str);
|
");
|
||||||
match execute_nvim_lua(&code).await {
|
match execute_nvim_lua(&code).await {
|
||||||
Ok(res) => send_text_result!(id.clone(), res),
|
Ok(res) => send_text_result!(id.clone(), res),
|
||||||
Err(e) => send_error(id, -32603, &e).await,
|
Err(e) => send_error(id, -32603, &e).await,
|
||||||
@@ -962,6 +990,18 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
|
|
||||||
"nvim_execute_lua" => {
|
"nvim_execute_lua" => {
|
||||||
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
||||||
|
// BAKE IN: Block interactive prompts that cause server deadlocks
|
||||||
|
let lower_code = code.to_lowercase();
|
||||||
|
if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Block the 'c' confirm flag in vim.cmd substitutions
|
||||||
|
if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && (code.contains("gc'") || code.contains("gc\"") || code.contains("gc\n") || code.contains("c'") || code.contains("c\"")) {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/g' or '/ge' instead.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
match execute_nvim_lua(code).await {
|
match execute_nvim_lua(code).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
send_text_result!(id.clone(), result);
|
send_text_result!(id.clone(), result);
|
||||||
@@ -976,28 +1016,31 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
if !id.is_null() {
|
if id.is_null() {
|
||||||
send_error(id, -32601, "Method not found").await;
|
|
||||||
} else {
|
|
||||||
// Ignore notifications silently
|
// Ignore notifications silently
|
||||||
|
} else {
|
||||||
|
send_error(id, -32601, "Method not found").await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let elapsed = start_time.elapsed();
|
let elapsed = start_time.elapsed();
|
||||||
tracing::info!("<<< [Nvim] {} (id: {}) completed in {:?}", method_clone, id_clone, elapsed);
|
tracing::info!(
|
||||||
|
"<<< [Nvim] {} (id: {}) completed in {:?}",
|
||||||
|
method_clone,
|
||||||
|
id_clone,
|
||||||
|
elapsed
|
||||||
|
);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
fn init_logging(app_name: &str) -> tracing_appender::non_blocking::WorkerGuard {
|
fn init_logging(app_name: &str) -> tracing_appender::non_blocking::WorkerGuard {
|
||||||
let log_dir = dirs::home_dir()
|
let log_dir = dirs::home_dir()
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.join(".gemini/mcp_memory/logs");
|
.join(".gemini/mcp_memory/logs");
|
||||||
std::fs::create_dir_all(&log_dir).unwrap_or_default();
|
std::fs::create_dir_all(&log_dir).unwrap_or_default();
|
||||||
|
|
||||||
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
|
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{app_name}.log"));
|
||||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||||
|
|
||||||
let _ = tracing_subscriber::fmt()
|
let _ = tracing_subscriber::fmt()
|
||||||
@@ -1036,11 +1079,10 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_rmpv_to_json_map() {
|
fn test_rmpv_to_json_map() {
|
||||||
let mut map = vec![];
|
let map = vec![(
|
||||||
map.push((
|
|
||||||
rmpv::Value::String("key1".into()),
|
rmpv::Value::String("key1".into()),
|
||||||
rmpv::Value::Integer(100.into()),
|
rmpv::Value::Integer(100.into()),
|
||||||
));
|
)];
|
||||||
let rmp_map = rmpv::Value::Map(map);
|
let rmp_map = rmpv::Value::Map(map);
|
||||||
|
|
||||||
let json_map = rmpv_to_json(&rmp_map);
|
let json_map = rmpv_to_json(&rmp_map);
|
||||||
@@ -1075,3 +1117,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
{"code": "local buf = vim.fn.bufnr('server/src/search.rs')\nif buf == -1 then\n vim.cmd('e server/src/search.rs')\n buf = vim.api.nvim_get_current_buf()\nend\n\nvim.api.nvim_buf_set_lines(buf, 216, 222, false, {\n ' &self,',\n ' entities: &[Entity],',\n ' tasks: &[Task],',\n ' snippets: &[Snippet],',\n ' adrs: &[Adr],'\n})\n\n-- We need to change the loop variables inside the task from values to clones if they are passed as slices\n-- Actually we can just clone the slice data before moving it into spawn_blocking\nvim.api.nvim_buf_set_lines(buf, 222, 223, false, {\n ' ) -> tokio::task::JoinHandle<tantivy::Result<()>> {',\n ' let entities = entities.to_vec();',\n ' let tasks = tasks.to_vec();',\n ' let snippets = snippets.to_vec();',\n ' let adrs = adrs.to_vec();'\n})\n\nvim.cmd('write')\nrequire('notify')('Updated search index_batch signature', 'info', { title = 'Antigravity' })\nreturn 'ok'\n"}
|
||||||
+3
-6
@@ -12,10 +12,7 @@ fn main() {
|
|||||||
|
|
||||||
let is_dirty = Command::new("git")
|
let is_dirty = Command::new("git")
|
||||||
.args(["status", "--porcelain"])
|
.args(["status", "--porcelain"])
|
||||||
.output()
|
.output().is_ok_and(|out| !out.stdout.is_empty());
|
||||||
.ok()
|
|
||||||
.map(|out| !out.stdout.is_empty())
|
|
||||||
.unwrap_or(false);
|
|
||||||
|
|
||||||
if is_dirty {
|
if is_dirty {
|
||||||
git_hash.push_str("-dirty");
|
git_hash.push_str("-dirty");
|
||||||
@@ -23,8 +20,8 @@ fn main() {
|
|||||||
|
|
||||||
let build_date = chrono::Local::now().format("%y.%m.%d.%H%M%S").to_string();
|
let build_date = chrono::Local::now().format("%y.%m.%d.%H%M%S").to_string();
|
||||||
|
|
||||||
let version = format!("{} ({})", build_date, git_hash);
|
let version = format!("{build_date} ({git_hash})");
|
||||||
println!("cargo:rustc-env=APP_VERSION={}", version);
|
println!("cargo:rustc-env=APP_VERSION={version}");
|
||||||
|
|
||||||
if std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default() == "windows" {
|
if std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default() == "windows" {
|
||||||
let mut res = winres::WindowsResource::new();
|
let mut res = winres::WindowsResource::new();
|
||||||
|
|||||||
@@ -0,0 +1,62 @@
|
|||||||
|
import os
|
||||||
|
import re
|
||||||
|
|
||||||
|
def bake_nvim_safeguards():
|
||||||
|
filepath = 'nvim-core/src/lib.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
# 1. We're going to inject an AST-level safeguard inside "nvim_execute_lua"
|
||||||
|
# Find the execute_lua block
|
||||||
|
target = """ "nvim_execute_lua" => {
|
||||||
|
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
||||||
|
match execute_nvim_lua(code).await {"""
|
||||||
|
|
||||||
|
safeguarded = """ "nvim_execute_lua" => {
|
||||||
|
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
||||||
|
// BAKE IN: Block interactive prompts that cause server deadlocks
|
||||||
|
let lower_code = code.to_lowercase();
|
||||||
|
if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Block the 'c' confirm flag in vim.cmd substitutions
|
||||||
|
if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && code.contains("c") {
|
||||||
|
// A rough heuristic to block `%s/foo/bar/gc`
|
||||||
|
if regex::Regex::new(r"%s.*/.*c").unwrap().is_match(code) {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/ge' instead.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match execute_nvim_lua(code).await {"""
|
||||||
|
|
||||||
|
if target in content:
|
||||||
|
# Also need to add regex as a dependency or just use basic string matching.
|
||||||
|
# Let's just use basic string matching for the %s block to avoid adding the regex crate to nvim-core if not present
|
||||||
|
|
||||||
|
safeguarded_no_regex = """ "nvim_execute_lua" => {
|
||||||
|
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
||||||
|
// BAKE IN: Block interactive prompts that cause server deadlocks
|
||||||
|
let lower_code = code.to_lowercase();
|
||||||
|
if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Block the 'c' confirm flag in vim.cmd substitutions
|
||||||
|
if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && (code.contains("gc'") || code.contains('gc"') || code.contains("gc\\n") || code.contains("c'") || code.contains('c"')) {
|
||||||
|
send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/g' or '/ge' instead.").await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
match execute_nvim_lua(code).await {"""
|
||||||
|
|
||||||
|
content = content.replace(target, safeguarded_no_regex)
|
||||||
|
print("Baked safeguards into nvim_execute_lua")
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
else:
|
||||||
|
print("Could not find nvim_execute_lua block")
|
||||||
|
|
||||||
|
bake_nvim_safeguards()
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def rewrite_graph():
|
||||||
|
with open('server/src/handlers_v2/graph.rs', 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_block = """ let mut included: HashSet<&str> = HashSet::new();
|
||||||
|
let mut to_draw = Vec::new();
|
||||||
|
|
||||||
|
state.read_graph(|full| {
|
||||||
|
for (name, e) in &full.entities {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& e.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty()
|
||||||
|
|| contains_ignore_ascii_case(name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&e.entity_type, &query)
|
||||||
|
{
|
||||||
|
included.insert(name.as_str());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for r in &full.relations {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& r.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty() || included.contains(r.from.as_str()) || included.contains(r.to.as_str()) {
|
||||||
|
included.insert(r.from.as_str());
|
||||||
|
included.insert(r.to.as_str());
|
||||||
|
to_draw.push(r.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
use std::fmt::Write;
|
||||||
|
let mut output = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
|
||||||
|
output.push_str("graph TD;\\n");
|
||||||
|
|
||||||
|
let sanitize = |s: &str, id_mode: bool| -> String {
|
||||||
|
let mut out = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
if c != '"' && c != '(' && c != ')' {
|
||||||
|
if id_mode && (c == ' ' || c == '-' || c == '.') {
|
||||||
|
out.push('_');
|
||||||
|
} else {
|
||||||
|
out.push(c);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out
|
||||||
|
};
|
||||||
|
|
||||||
|
for name in &included {
|
||||||
|
let _ = writeln!(
|
||||||
|
output,
|
||||||
|
" id_{}[\\"{}\\"];",
|
||||||
|
sanitize(name, true),
|
||||||
|
sanitize(name, false)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for r in to_draw {
|
||||||
|
let _ = writeln!(
|
||||||
|
output,
|
||||||
|
" id_{}-->|\\"{}\\"|id_{};",
|
||||||
|
sanitize(&r.from, true),
|
||||||
|
r.relation_type.replace("\\"", ""),
|
||||||
|
sanitize(&r.to, true)
|
||||||
|
);
|
||||||
|
}"""
|
||||||
|
|
||||||
|
new_block = """ use std::fmt::Write;
|
||||||
|
let mut output = state.read_graph(|full| {
|
||||||
|
let mut included: HashSet<&str> = HashSet::new();
|
||||||
|
let mut to_draw = Vec::new();
|
||||||
|
|
||||||
|
for (name, e) in &full.entities {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& e.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty()
|
||||||
|
|| contains_ignore_ascii_case(name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&e.entity_type, &query)
|
||||||
|
{
|
||||||
|
included.insert(name.as_str());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for r in &full.relations {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& r.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty() || included.contains(r.from.as_str()) || included.contains(r.to.as_str()) {
|
||||||
|
included.insert(r.from.as_str());
|
||||||
|
included.insert(r.to.as_str());
|
||||||
|
to_draw.push(r.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
|
||||||
|
out.push_str("graph TD;\\n");
|
||||||
|
|
||||||
|
let sanitize = |s: &str, id_mode: bool| -> String {
|
||||||
|
let mut o = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
if c != '"' && c != '(' && c != ')' {
|
||||||
|
if id_mode && (c == ' ' || c == '-' || c == '.') {
|
||||||
|
o.push('_');
|
||||||
|
} else {
|
||||||
|
o.push(c);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
o
|
||||||
|
};
|
||||||
|
|
||||||
|
for name in &included {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
" id_{}[\\"{}\\"];",
|
||||||
|
sanitize(name, true),
|
||||||
|
sanitize(name, false)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for r in to_draw {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
" id_{}-->|\\"{}\\"|id_{};",
|
||||||
|
sanitize(&r.from, true),
|
||||||
|
r.relation_type.replace("\\"", ""),
|
||||||
|
sanitize(&r.to, true)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
out
|
||||||
|
});"""
|
||||||
|
|
||||||
|
if old_block in content:
|
||||||
|
with open('server/src/handlers_v2/graph.rs', 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content.replace(old_block, new_block))
|
||||||
|
print("Replaced visualize_graph")
|
||||||
|
else:
|
||||||
|
print("Could not find old block")
|
||||||
|
|
||||||
|
rewrite_graph()
|
||||||
@@ -895,23 +895,17 @@
|
|||||||
// --- WebSocket Activity Feed ---
|
// --- WebSocket Activity Feed ---
|
||||||
const MAX_ACTIVITY_HISTORY = 100;
|
const MAX_ACTIVITY_HISTORY = 100;
|
||||||
|
|
||||||
function saveActivityHistory(time, htmlData) {
|
async function loadActivityHistory() {
|
||||||
try {
|
|
||||||
let history = JSON.parse(localStorage.getItem('activityHistory') || '[]');
|
|
||||||
history.push({ time, data: htmlData });
|
|
||||||
if (history.length > MAX_ACTIVITY_HISTORY) history = history.slice(-MAX_ACTIVITY_HISTORY);
|
|
||||||
localStorage.setItem('activityHistory', JSON.stringify(history));
|
|
||||||
} catch(e) {}
|
|
||||||
}
|
|
||||||
|
|
||||||
function loadActivityHistory() {
|
|
||||||
try {
|
try {
|
||||||
|
const response = await fetch('/api/activity');
|
||||||
|
const history = await response.json();
|
||||||
const feed = document.getElementById('activity-feed');
|
const feed = document.getElementById('activity-feed');
|
||||||
let history = JSON.parse(localStorage.getItem('activityHistory') || '[]');
|
feed.innerHTML = '';
|
||||||
history.forEach(item => {
|
history.forEach(item => {
|
||||||
const div = document.createElement('div');
|
const div = document.createElement('div');
|
||||||
div.className = 'feed-entry';
|
div.className = 'feed-entry';
|
||||||
div.innerHTML = `<span class="time">[${item.time}]</span> ${item.data}`;
|
const timeStr = new Date(item.time).toLocaleTimeString([], {hour: '2-digit', minute:'2-digit', second:'2-digit'});
|
||||||
|
div.innerHTML = `<span class="time">[${timeStr}]</span> ${item.message || item.data}`;
|
||||||
feed.appendChild(div);
|
feed.appendChild(div);
|
||||||
});
|
});
|
||||||
if (history.length > 0) {
|
if (history.length > 0) {
|
||||||
@@ -931,12 +925,10 @@
|
|||||||
if (data.type === 'activity') {
|
if (data.type === 'activity') {
|
||||||
const div = document.createElement('div');
|
const div = document.createElement('div');
|
||||||
div.className = 'feed-entry';
|
div.className = 'feed-entry';
|
||||||
const time = new Date().toLocaleTimeString();
|
const timeStr = new Date(data.data.time).toLocaleTimeString([], {hour: '2-digit', minute:'2-digit', second:'2-digit'});
|
||||||
div.innerHTML = `<span class="time">[${time}]</span> ${data.data}`;
|
div.innerHTML = `<span class="time">[${timeStr}]</span> ${data.data.message || data.data.data || data.data}`;
|
||||||
feed.appendChild(div);
|
feed.appendChild(div);
|
||||||
|
|
||||||
saveActivityHistory(time, data.data);
|
|
||||||
|
|
||||||
// Auto-scroll logic
|
// Auto-scroll logic
|
||||||
const isScrolledToBottom = feed.scrollHeight - feed.clientHeight <= feed.scrollTop + 20;
|
const isScrolledToBottom = feed.scrollHeight - feed.clientHeight <= feed.scrollTop + 20;
|
||||||
if (isScrolledToBottom) {
|
if (isScrolledToBottom) {
|
||||||
|
|||||||
@@ -0,0 +1,61 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def fix_env_and_workspaces():
|
||||||
|
filepath = 'server/src/handlers_v2/env.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_env = """ let fps = state.env_fingerprints.read();
|
||||||
|
if let Some(fp) = fps.get(&req.namespace) {"""
|
||||||
|
|
||||||
|
new_env = """ let data = state.env_fingerprints.read_with(|fps| {
|
||||||
|
fps.get(&req.namespace).cloned()
|
||||||
|
});
|
||||||
|
if let Some(fp) = data {"""
|
||||||
|
|
||||||
|
if old_env in content:
|
||||||
|
content = content.replace(old_env, new_env)
|
||||||
|
print('Replaced env read')
|
||||||
|
else:
|
||||||
|
print('Could not find old_env')
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
filepath = 'server/src/handlers_v2/workspaces.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_ws = """ let query = req.query.to_lowercase();
|
||||||
|
let snippets = state.snippets.read();
|
||||||
|
let mut results = Vec::new();
|
||||||
|
for s in snippets {
|
||||||
|
if contains_ignore_ascii_case(&s.name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.description, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.language, &query)
|
||||||
|
{
|
||||||
|
results.push(s);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let data = serde_json::to_string(&results).unwrap_or_default();"""
|
||||||
|
|
||||||
|
new_ws = """ let query = req.query.to_lowercase();
|
||||||
|
let data = state.snippets.read_with(|snippets| {
|
||||||
|
let results: Vec<_> = snippets.iter().filter(|s| {
|
||||||
|
contains_ignore_ascii_case(&s.name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.description, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.language, &query)
|
||||||
|
}).cloned().collect();
|
||||||
|
serde_json::to_string(&results).unwrap_or_default()
|
||||||
|
});"""
|
||||||
|
|
||||||
|
if old_ws in content:
|
||||||
|
content = content.replace(old_ws, new_ws)
|
||||||
|
print('Replaced workspaces read')
|
||||||
|
else:
|
||||||
|
print('Could not find old_ws')
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
fix_env_and_workspaces()
|
||||||
@@ -0,0 +1,60 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def fix_main_stats():
|
||||||
|
filepath = 'server/src/main.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_stats = """ let (entities, relations) = {
|
||||||
|
let graph = state_clone.get_full_graph();
|
||||||
|
(graph.entities.len(), graph.relations.len())
|
||||||
|
};
|
||||||
|
let tasks = state_clone.tasks.read().len();
|
||||||
|
let snippets = state_clone.snippets.read().len();
|
||||||
|
let tech_debts = state_clone.tech_debts.read().len();
|
||||||
|
let adrs = state_clone.adrs.read().len();
|
||||||
|
|
||||||
|
let ledger = state_clone.ledger.read().len();
|
||||||
|
let sticky = state_clone.sticky.read().len();
|
||||||
|
let error_fixes = state_clone.error_fixes.read().len();
|
||||||
|
let pinned_files = state_clone.pinned_files.read().len();
|
||||||
|
let session_summaries = state_clone.session_summaries.read().len();
|
||||||
|
let handoff_memos = state_clone.handoff_memos.read().len();
|
||||||
|
let env_fingerprints = state_clone.env_fingerprints.read().len();
|
||||||
|
let env_requirements = state_clone.env_requirements.read().len();
|
||||||
|
let milestones = state_clone.milestones.read().len();
|
||||||
|
let environments = state_clone.environments.read().len();
|
||||||
|
let pr_checklists = state_clone.pr_checklists.read().len();
|
||||||
|
let gates = state_clone.gates.read().len();
|
||||||
|
let context_workspaces = state_clone.context_workspaces.read().len();"""
|
||||||
|
|
||||||
|
new_stats = """ let (entities, relations) = state_clone.read_graph(|g| (g.entities.len(), g.relations.len()));
|
||||||
|
let tasks = state_clone.tasks.read_with(|items| items.len());
|
||||||
|
let snippets = state_clone.snippets.read_with(|items| items.len());
|
||||||
|
let tech_debts = state_clone.tech_debts.read_with(|items| items.len());
|
||||||
|
let adrs = state_clone.adrs.read_with(|items| items.len());
|
||||||
|
|
||||||
|
let ledger = state_clone.ledger.read_with(|items| items.len());
|
||||||
|
let sticky = state_clone.sticky.read_with(|items| items.len());
|
||||||
|
let error_fixes = state_clone.error_fixes.read_with(|items| items.len());
|
||||||
|
let pinned_files = state_clone.pinned_files.read_with(|items| items.len());
|
||||||
|
let session_summaries = state_clone.session_summaries.read_with(|items| items.len());
|
||||||
|
let handoff_memos = state_clone.handoff_memos.read_with(|items| items.len());
|
||||||
|
let env_fingerprints = state_clone.env_fingerprints.read_with(|items| items.len());
|
||||||
|
let env_requirements = state_clone.env_requirements.read_with(|items| items.len());
|
||||||
|
let milestones = state_clone.milestones.read_with(|items| items.len());
|
||||||
|
let environments = state_clone.environments.read_with(|items| items.len());
|
||||||
|
let pr_checklists = state_clone.pr_checklists.read_with(|items| items.len());
|
||||||
|
let gates = state_clone.gates.read_with(|items| items.len());
|
||||||
|
let context_workspaces = state_clone.context_workspaces.read_with(|items| items.len());"""
|
||||||
|
|
||||||
|
if old_stats in content:
|
||||||
|
content = content.replace(old_stats, new_stats)
|
||||||
|
print('Replaced stats reads')
|
||||||
|
else:
|
||||||
|
print('Could not find old_stats')
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
fix_main_stats()
|
||||||
@@ -0,0 +1,111 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def fix_meta():
|
||||||
|
filepath = 'server/src/handlers_v2/meta.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_omni = """ let q = req.query.to_lowercase();
|
||||||
|
let tech_debts: Vec<_> = state
|
||||||
|
.tech_debts
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|d| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| d.namespace == *ns)
|
||||||
|
&& (contains_ignore_ascii_case(&d.description, &q)
|
||||||
|
|| contains_ignore_ascii_case(&d.ideal_solution, &q))
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let memos: Vec<_> = state
|
||||||
|
.handoff_memos
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|m| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
|
||||||
|
&& contains_ignore_ascii_case(&m.content, &q)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let error_fixes: Vec<_> = state
|
||||||
|
.error_fixes
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|f| {
|
||||||
|
contains_ignore_ascii_case(&f.signature, &q)
|
||||||
|
|| contains_ignore_ascii_case(&f.solution, &q)
|
||||||
|
})
|
||||||
|
.collect();"""
|
||||||
|
|
||||||
|
new_omni = """ let q = req.query.to_lowercase();
|
||||||
|
let tech_debts = state.tech_debts.read_with(|debts| {
|
||||||
|
debts.iter().filter(|d| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| d.namespace == *ns)
|
||||||
|
&& (contains_ignore_ascii_case(&d.description, &q)
|
||||||
|
|| contains_ignore_ascii_case(&d.ideal_solution, &q))
|
||||||
|
}).cloned().collect::<Vec<_>>()
|
||||||
|
});
|
||||||
|
let memos = state.handoff_memos.read_with(|memos| {
|
||||||
|
memos.iter().filter(|m| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
|
||||||
|
&& contains_ignore_ascii_case(&m.content, &q)
|
||||||
|
}).cloned().collect::<Vec<_>>()
|
||||||
|
});
|
||||||
|
let error_fixes = state.error_fixes.read_with(|fixes| {
|
||||||
|
fixes.iter().filter(|f| {
|
||||||
|
contains_ignore_ascii_case(&f.signature, &q)
|
||||||
|
|| contains_ignore_ascii_case(&f.solution, &q)
|
||||||
|
}).cloned().collect::<Vec<_>>()
|
||||||
|
});"""
|
||||||
|
|
||||||
|
if old_omni in content:
|
||||||
|
content = content.replace(old_omni, new_omni)
|
||||||
|
print('Replaced omni_search vecs')
|
||||||
|
else:
|
||||||
|
print('Could not find old_omni')
|
||||||
|
|
||||||
|
old_health = """ let active_tasks = state
|
||||||
|
.tasks
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|t| t.status != "done")
|
||||||
|
.count();
|
||||||
|
let unresolved_debt = state
|
||||||
|
.tech_debts
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|d| d.namespace == req.namespace && !d.is_resolved)
|
||||||
|
.count();
|
||||||
|
let unread_memos = state
|
||||||
|
.handoff_memos
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|m| m.namespace == req.namespace)
|
||||||
|
.count();
|
||||||
|
let active_milestones = state
|
||||||
|
.milestones
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|m| m.namespace == req.namespace && m.status != "done")
|
||||||
|
.count();
|
||||||
|
let remaining_checklists = state
|
||||||
|
.pr_checklists
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|c| c.namespace == req.namespace)
|
||||||
|
.count();"""
|
||||||
|
|
||||||
|
new_health = """ let active_tasks = state.tasks.read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count());
|
||||||
|
let unresolved_debt = state.tech_debts.read_with(|debts| debts.iter().filter(|d| d.namespace == req.namespace && !d.is_resolved).count());
|
||||||
|
let unread_memos = state.handoff_memos.read_with(|memos| memos.iter().filter(|m| m.namespace == req.namespace).count());
|
||||||
|
let active_milestones = state.milestones.read_with(|milestones| milestones.iter().filter(|m| m.namespace == req.namespace && m.status != "done").count());
|
||||||
|
let remaining_checklists = state.pr_checklists.read_with(|checklists| checklists.iter().filter(|c| c.namespace == req.namespace).count());"""
|
||||||
|
|
||||||
|
if old_health in content:
|
||||||
|
content = content.replace(old_health, new_health)
|
||||||
|
print('Replaced health counters')
|
||||||
|
else:
|
||||||
|
print('Could not find old_health')
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
fix_meta()
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
import os
|
||||||
|
import re
|
||||||
|
|
||||||
|
def fix_nvim_msgid():
|
||||||
|
filepath = 'nvim-core/src/lib.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
# Add atomic import and static var if not exists
|
||||||
|
if 'static NEXT_MSGID' not in content:
|
||||||
|
atomic_def = "use std::sync::atomic::{AtomicU64, Ordering};\nstatic NEXT_MSGID: AtomicU64 = AtomicU64::new(1);\n"
|
||||||
|
# Find NVIM_CONN
|
||||||
|
conn_idx = content.find('static NVIM_CONN')
|
||||||
|
if conn_idx != -1:
|
||||||
|
content = content[:conn_idx] + atomic_def + content[conn_idx:]
|
||||||
|
|
||||||
|
# Replace all hardcoded msgid
|
||||||
|
# e.g., RmpValue::Integer(1.into()), // msgid
|
||||||
|
# with: let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst); ... RmpValue::Integer(msgid.into()),
|
||||||
|
|
||||||
|
# We need to insert `let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);` before `let req = RmpValue::Array(vec![`
|
||||||
|
|
||||||
|
# We will use regex to find the blocks
|
||||||
|
funcs = [
|
||||||
|
('send_nvim_command', '1'),
|
||||||
|
('get_nvim_active_buffer', '2'),
|
||||||
|
('get_nvim_cursor', '3'),
|
||||||
|
('get_nvim_visual_selection', '4'),
|
||||||
|
('set_nvim_diagnostics', '5'),
|
||||||
|
('execute_nvim_lua', '6'),
|
||||||
|
]
|
||||||
|
|
||||||
|
for func, old_id in funcs:
|
||||||
|
pattern = rf"let req = RmpValue::Array\(vec!\[\s*RmpValue::Integer\(0\.into\(\)\),\s*RmpValue::Integer\({old_id}\.into\(\)\), // msgid"
|
||||||
|
replacement = f"let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);\n let req = RmpValue::Array(vec![\n RmpValue::Integer(0.into()),\n RmpValue::Integer(msgid.into()), // msgid"
|
||||||
|
|
||||||
|
content = re.sub(pattern, replacement, content)
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
print('Fixed msgid allocations')
|
||||||
|
|
||||||
|
fix_nvim_msgid()
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def fix_notes():
|
||||||
|
filepath = 'server/src/handlers_v2/notes.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
old_standup = """ let tasks = state
|
||||||
|
.tasks
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|t| t.updated_at >= cutoff)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let changes = state
|
||||||
|
.ledger
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|c| c.timestamp >= cutoff)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let summaries = state
|
||||||
|
.session_summaries
|
||||||
|
.read()
|
||||||
|
.into_iter()
|
||||||
|
.filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff)
|
||||||
|
.collect::<Vec<_>>();"""
|
||||||
|
|
||||||
|
new_standup = """ let tasks = state.tasks.read_with(|items| items.iter().filter(|t| t.updated_at >= cutoff).cloned().collect::<Vec<_>>());
|
||||||
|
let changes = state.ledger.read_with(|items| items.iter().filter(|c| c.timestamp >= cutoff).cloned().collect::<Vec<_>>());
|
||||||
|
let summaries = state.session_summaries.read_with(|items| items.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).cloned().collect::<Vec<_>>());"""
|
||||||
|
|
||||||
|
if old_standup in content:
|
||||||
|
content = content.replace(old_standup, new_standup)
|
||||||
|
print('Replaced standup in notes.rs')
|
||||||
|
else:
|
||||||
|
print('Could not find old_standup in notes.rs')
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
fix_notes()
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
def fix_stub_leaks():
|
||||||
|
filepath = 'stub/src/main.rs'
|
||||||
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
# 1. Fix connect_async to handle shutdown and timeout
|
||||||
|
old_connect = """ match tokio_tungstenite::connect_async(request).await {"""
|
||||||
|
new_connect = """ let connect_result = tokio::select! {
|
||||||
|
_ = shutdown_rx.recv() => {
|
||||||
|
tracing::info!("Shutdown received during connect");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
res = tokio::time::timeout(
|
||||||
|
tokio::time::Duration::from_secs(5),
|
||||||
|
tokio_tungstenite::connect_async(request)
|
||||||
|
) => res,
|
||||||
|
};
|
||||||
|
|
||||||
|
match connect_result {
|
||||||
|
Ok(Ok((ws_stream, _))) => {"""
|
||||||
|
|
||||||
|
if old_connect in content:
|
||||||
|
content = content.replace(old_connect, new_connect)
|
||||||
|
print("Replaced connect_async")
|
||||||
|
|
||||||
|
# Fix Err block to match the new match structure
|
||||||
|
old_err = """ Err(e) => {
|
||||||
|
tracing::error!("Failed to connect via WSS: {}", e);
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
|
}"""
|
||||||
|
new_err = """ Ok(Err(e)) => {
|
||||||
|
tracing::error!("Failed to connect via WSS: {}", e);
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
tracing::error!("Connection attempt timed out");
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
|
}"""
|
||||||
|
|
||||||
|
if old_err in content:
|
||||||
|
content = content.replace(old_err, new_err)
|
||||||
|
print("Replaced Err branch")
|
||||||
|
|
||||||
|
# 2. Fix the break in send_task that exits the stub instead of reconnecting
|
||||||
|
old_select_send = """ _ = &mut send_task => {
|
||||||
|
tracing::error!("Send task exited");
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||||
|
recv_task.abort();
|
||||||
|
break;
|
||||||
|
}"""
|
||||||
|
new_select_send = """ _ = &mut send_task => {
|
||||||
|
tracing::error!("Send task exited");
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||||
|
recv_task.abort();
|
||||||
|
}"""
|
||||||
|
|
||||||
|
if old_select_send in content:
|
||||||
|
content = content.replace(old_select_send, new_select_send)
|
||||||
|
print("Replaced select send_task")
|
||||||
|
|
||||||
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
fix_stub_leaks()
|
||||||
+120
-2985
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,169 @@
|
|||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
pub struct UpdateEnvFingerprintHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for UpdateEnvFingerprintHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"update_env_fingerprint"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
|
||||||
|
"update_env_fingerprint",
|
||||||
|
"Execute update_env_fingerprint",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: UpdateEnvFingerprintTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.env_fingerprints.modify(|fps| {
|
||||||
|
fps.insert(
|
||||||
|
req.namespace.clone(),
|
||||||
|
crate::models::EnvFingerprint {
|
||||||
|
namespace: req.namespace.clone(),
|
||||||
|
os: std::env::consts::OS.to_string(),
|
||||||
|
shell: std::env::var("SHELL").unwrap_or_else(|_| "unknown".to_string()),
|
||||||
|
tool_versions: req.tool_versions,
|
||||||
|
updated_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
Ok("Env fingerprint updated".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ReadEnvFingerprintHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ReadEnvFingerprintHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"read_env_fingerprint"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ReadEnvFingerprintTool>(
|
||||||
|
"read_env_fingerprint",
|
||||||
|
"Execute read_env_fingerprint",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ReadEnvFingerprintTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.env_fingerprints.read_with(|fps| {
|
||||||
|
fps.get(&req.namespace).cloned()
|
||||||
|
});
|
||||||
|
if let Some(fp) = data {
|
||||||
|
let data = serde_json::to_string(&fp).unwrap_or_default();
|
||||||
|
Ok(data.to_string())
|
||||||
|
} else {
|
||||||
|
Ok("{}".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LogEnvRequirementHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LogEnvRequirementHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"log_env_requirement"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LogEnvRequirementTool>(
|
||||||
|
"log_env_requirement",
|
||||||
|
"Execute log_env_requirement",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LogEnvRequirementTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.env_requirements.modify(|reqs| {
|
||||||
|
reqs.retain(|r| !(r.namespace == req.namespace && r.key == req.key));
|
||||||
|
reqs.push(crate::models::EnvRequirement {
|
||||||
|
namespace: req.namespace,
|
||||||
|
key: req.key,
|
||||||
|
description: req.description,
|
||||||
|
is_secret: req.is_secret,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("Env requirement logged".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct RegisterEnvironmentHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for RegisterEnvironmentHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"register_environment"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<RegisterEnvironmentTool>(
|
||||||
|
"register_environment",
|
||||||
|
"Execute register_environment",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: RegisterEnvironmentTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.environments.modify(|envs| {
|
||||||
|
envs.retain(|e| !(e.namespace == req.namespace && e.name == req.name));
|
||||||
|
envs.push(crate::models::EnvironmentDetail {
|
||||||
|
namespace: req.namespace,
|
||||||
|
name: req.name,
|
||||||
|
url: req.url,
|
||||||
|
description: req.description,
|
||||||
|
requires_vpn: req.requires_vpn,
|
||||||
|
updated_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("Environment registered".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct GetEnvironmentDetailsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for GetEnvironmentDetailsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"get_environment_details"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
|
||||||
|
"get_environment_details",
|
||||||
|
"Execute get_environment_details",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: GetEnvironmentDetailsTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.environments.read_with(|envs| {
|
||||||
|
let filtered: Vec<_> = envs.iter().filter(|e| e.namespace == req.namespace).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1,556 @@
|
|||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(serde::Serialize)]
|
||||||
|
#[derive(Default)]
|
||||||
|
struct BorrowedGraph<'a> {
|
||||||
|
entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
|
||||||
|
relations: Vec<&'a crate::models::Relation>,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
pub struct QueryGraphPathHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for QueryGraphPathHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"query_graph_path"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: crate::tools::QueryGraphPathTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.read_graph(|graph| {
|
||||||
|
let max_depth = req.max_depth.unwrap_or(5);
|
||||||
|
let mut queue = std::collections::VecDeque::new();
|
||||||
|
let mut visited = std::collections::HashSet::new();
|
||||||
|
let mut parents: std::collections::HashMap<String, (String, String)> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
|
|
||||||
|
queue.push_back(req.start_node.clone());
|
||||||
|
visited.insert(req.start_node.clone());
|
||||||
|
|
||||||
|
let mut found = false;
|
||||||
|
let mut current_depth = 0;
|
||||||
|
let mut nodes_at_current_depth = 1;
|
||||||
|
let mut nodes_at_next_depth = 0;
|
||||||
|
|
||||||
|
while let Some(current) = queue.pop_front() {
|
||||||
|
if current == req.end_node {
|
||||||
|
found = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
nodes_at_current_depth -= 1;
|
||||||
|
if current_depth < max_depth {
|
||||||
|
for rel in &graph.relations {
|
||||||
|
if rel.from == current && !visited.contains(&rel.to) {
|
||||||
|
visited.insert(rel.to.clone());
|
||||||
|
parents.insert(
|
||||||
|
rel.to.clone(),
|
||||||
|
(current.clone(), rel.relation_type.clone()),
|
||||||
|
);
|
||||||
|
queue.push_back(rel.to.clone());
|
||||||
|
nodes_at_next_depth += 1;
|
||||||
|
} else if rel.to == current && !visited.contains(&rel.from) {
|
||||||
|
visited.insert(rel.from.clone());
|
||||||
|
parents.insert(
|
||||||
|
rel.from.clone(),
|
||||||
|
(current.clone(), format!("inverse({})", rel.relation_type)),
|
||||||
|
);
|
||||||
|
queue.push_back(rel.from.clone());
|
||||||
|
nodes_at_next_depth += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if nodes_at_current_depth == 0 {
|
||||||
|
current_depth += 1;
|
||||||
|
nodes_at_current_depth = nodes_at_next_depth;
|
||||||
|
nodes_at_next_depth = 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if found {
|
||||||
|
let mut path = Vec::new();
|
||||||
|
let mut curr = req.end_node.clone();
|
||||||
|
while curr != req.start_node {
|
||||||
|
if let Some((parent, rel_type)) = parents.get(&curr) {
|
||||||
|
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
|
||||||
|
curr = parent.clone();
|
||||||
|
} else {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
path.reverse();
|
||||||
|
Ok(format!("Path found:\n{}", path.join("\n")))
|
||||||
|
} else {
|
||||||
|
Ok(format!(
|
||||||
|
"No path found between {} and {} within depth {}",
|
||||||
|
req.start_node, req.end_node, max_depth
|
||||||
|
))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct CreateEntitiesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for CreateEntitiesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"create_entities"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut inserted = Vec::new();
|
||||||
|
state.modify_graph(|g| {
|
||||||
|
for entity in req.entities {
|
||||||
|
if !entity.name.is_empty() {
|
||||||
|
inserted.push(entity.clone());
|
||||||
|
g.entities.insert(entity.name.clone(), entity);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||||
|
for entity in inserted {
|
||||||
|
let _ = idx.index_entity(&entity).await;
|
||||||
|
}
|
||||||
|
Ok("Entities created".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct CreateRelationsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for CreateRelationsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"create_relations"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: CreateRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|g| {
|
||||||
|
for relation in req.relations {
|
||||||
|
if !relation.from.is_empty() && !relation.to.is_empty() {
|
||||||
|
g.relations.push(relation);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Ok("Relations created".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct AddObservationsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddObservationsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_observations"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|g| {
|
||||||
|
for o in req.observations {
|
||||||
|
if let Some(e) = g.entities.get_mut(&o.entity_name) {
|
||||||
|
e.observations.extend(o.contents);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Ok("Observations added".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteEntitiesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteEntitiesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_entities"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteEntitiesTool>("delete_entities", "Execute delete_entities")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let to_delete: HashSet<_> = req.entity_names.into_iter().collect();
|
||||||
|
state.modify_graph(|master| {
|
||||||
|
for name in &to_delete {
|
||||||
|
master.entities.remove(name);
|
||||||
|
}
|
||||||
|
master
|
||||||
|
.relations
|
||||||
|
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
|
||||||
|
});
|
||||||
|
Ok("Entities deleted".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteObservationsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteObservationsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_observations"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteObservationsTool>(
|
||||||
|
"delete_observations",
|
||||||
|
"Execute delete_observations",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteObservationsTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|master| {
|
||||||
|
for d in req.deletions {
|
||||||
|
if let Some(e) = master.entities.get_mut(&d.entity_name) {
|
||||||
|
let to_rem: HashSet<_> = d.observations.into_iter().collect();
|
||||||
|
e.observations.retain(|o| !to_rem.contains(o));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Ok("Observations deleted".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteRelationsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteRelationsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_relations"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteRelationsTool>("delete_relations", "Execute delete_relations")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|master| {
|
||||||
|
let to_rem: HashSet<_> = req.relations.into_iter().collect();
|
||||||
|
master.relations.retain(|r| !to_rem.contains(r));
|
||||||
|
});
|
||||||
|
Ok("Relations deleted".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ReadGraphHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ReadGraphHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"read_graph"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.read_graph(|full| {
|
||||||
|
if let Some(ns) = req.namespace {
|
||||||
|
let mut filtered = BorrowedGraph::default();
|
||||||
|
for (k, v) in &full.entities {
|
||||||
|
if v.namespace == ns {
|
||||||
|
filtered.entities.insert(k, v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for r in &full.relations {
|
||||||
|
if r.namespace == ns {
|
||||||
|
filtered.relations.push(r);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
serde_json::to_string(full).unwrap_or_default()
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SearchNodesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for SearchNodesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"search_nodes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<SearchNodesTool>("search_nodes", "Execute search_nodes")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let matches = if let Ok(idx) = state.search_index.read() {
|
||||||
|
idx.search(&req.query, req.namespace.as_deref())
|
||||||
|
.unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
vec![]
|
||||||
|
};
|
||||||
|
|
||||||
|
let data = state.read_graph(|full| {
|
||||||
|
let mut result = BorrowedGraph::default();
|
||||||
|
for (id, doc_type, _, _, _) in &matches {
|
||||||
|
if doc_type == "entity"
|
||||||
|
&& let Some(e) = full.entities.get(id)
|
||||||
|
{
|
||||||
|
result.entities.insert(id, e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::to_string(&result).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct OpenNodesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for OpenNodesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"open_nodes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<OpenNodesTool>("open_nodes", "Execute open_nodes")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.read_graph(|full| {
|
||||||
|
let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
|
||||||
|
let mut result = BorrowedGraph::default();
|
||||||
|
let mut connected: HashSet<&str> = HashSet::new();
|
||||||
|
for r in &full.relations {
|
||||||
|
if targets.contains(r.from.as_str()) {
|
||||||
|
connected.insert(r.to.as_str());
|
||||||
|
result.relations.push(r);
|
||||||
|
} else if targets.contains(r.to.as_str()) {
|
||||||
|
connected.insert(r.from.as_str());
|
||||||
|
result.relations.push(r);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (name, e) in &full.entities {
|
||||||
|
if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
|
||||||
|
result.entities.insert(name, e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::to_string(&result).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct VisualizeGraphHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for VisualizeGraphHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"visualize_graph"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<VisualizeGraphTool>("visualize_graph", "Execute visualize_graph")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let query = req.query.unwrap_or_default();
|
||||||
|
use std::fmt::Write;
|
||||||
|
let mut output = state.read_graph(|full| {
|
||||||
|
let mut included: HashSet<&str> = HashSet::new();
|
||||||
|
let mut to_draw = Vec::new();
|
||||||
|
|
||||||
|
for (name, e) in &full.entities {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& e.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty()
|
||||||
|
|| contains_ignore_ascii_case(name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&e.entity_type, &query)
|
||||||
|
{
|
||||||
|
included.insert(name.as_str());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for r in &full.relations {
|
||||||
|
if let Some(ns) = &req.namespace
|
||||||
|
&& r.namespace != *ns
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if query.is_empty() || included.contains(r.from.as_str()) || included.contains(r.to.as_str()) {
|
||||||
|
included.insert(r.from.as_str());
|
||||||
|
included.insert(r.to.as_str());
|
||||||
|
to_draw.push(r.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
|
||||||
|
out.push_str("graph TD;\n");
|
||||||
|
|
||||||
|
let sanitize = |s: &str, id_mode: bool| -> String {
|
||||||
|
let mut o = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
if c != '"' && c != '(' && c != ')' {
|
||||||
|
if id_mode && (c == ' ' || c == '-' || c == '.') {
|
||||||
|
o.push('_');
|
||||||
|
} else {
|
||||||
|
o.push(c);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
o
|
||||||
|
};
|
||||||
|
|
||||||
|
for name in &included {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
" id_{}[\"{}\"];",
|
||||||
|
sanitize(name, true),
|
||||||
|
sanitize(name, false)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for r in to_draw {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
" id_{}-->|\"{}\"|id_{};",
|
||||||
|
sanitize(&r.from, true),
|
||||||
|
r.relation_type.replace("\"", ""),
|
||||||
|
sanitize(&r.to, true)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
out
|
||||||
|
});
|
||||||
|
if output == "graph TD;\n" {
|
||||||
|
output = "No nodes found to visualize.".to_string();
|
||||||
|
}
|
||||||
|
Ok(output.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct CondenseEntityHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for CondenseEntityHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"condense_entity"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<CondenseEntityTool>("condense_entity", "Execute condense_entity")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|master| {
|
||||||
|
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
||||||
|
e.observations = req.summarized_observations;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Ok("Entity condensed".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct MergeEntitiesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for MergeEntitiesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"merge_entities"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<MergeEntitiesTool>("merge_entities", "Execute merge_entities")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.modify_graph(|master| {
|
||||||
|
if let Some(src) = master.entities.remove(&req.source_entity) {
|
||||||
|
if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
|
||||||
|
tgt.observations.extend(src.observations);
|
||||||
|
MemoryState::deduplicate(&mut tgt.observations);
|
||||||
|
} else {
|
||||||
|
let mut new_tgt = src.clone();
|
||||||
|
new_tgt.name = req.target_entity.clone();
|
||||||
|
master.entities.insert(req.target_entity.clone(), new_tgt);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for r in &mut master.relations {
|
||||||
|
if r.from == req.source_entity {
|
||||||
|
r.from = req.target_entity.clone();
|
||||||
|
}
|
||||||
|
if r.to == req.source_entity {
|
||||||
|
r.to = req.target_entity.clone();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
MemoryState::deduplicate(&mut master.relations);
|
||||||
|
});
|
||||||
|
Ok("Entities merged".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct FindOrphansHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for FindOrphansHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"find_orphans"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<FindOrphansTool>("find_orphans", "Execute find_orphans")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let orphans = state.read_graph(|full| {
|
||||||
|
let mut connected = std::collections::HashSet::new();
|
||||||
|
for r in &full.relations {
|
||||||
|
connected.insert(r.from.as_str());
|
||||||
|
connected.insert(r.to.as_str());
|
||||||
|
}
|
||||||
|
full.entities
|
||||||
|
.keys()
|
||||||
|
.filter(|k| !connected.contains(k.as_str()))
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<String>>()
|
||||||
|
});
|
||||||
|
let data = serde_json::to_string(&orphans).unwrap_or_default();
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
use crate::handlers_v2::utils::*;
|
||||||
@@ -0,0 +1,463 @@
|
|||||||
|
use crate::models::*;
|
||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
pub struct LogDecisionHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LogDecisionHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"log_decision"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Execute log_decision")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut adr_id = String::new();
|
||||||
|
let mut new_adr = None;
|
||||||
|
|
||||||
|
state.adrs.modify(|adrs| {
|
||||||
|
adr_id = format!("ADR-{:04}", adrs.len() + 1);
|
||||||
|
let a = Adr {
|
||||||
|
id: adr_id.clone(),
|
||||||
|
title: req.title,
|
||||||
|
context: req.context,
|
||||||
|
decision: req.decision,
|
||||||
|
consequence: req.consequence,
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
};
|
||||||
|
new_adr = Some(a.clone());
|
||||||
|
adrs.push(a);
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(adr) = new_adr {
|
||||||
|
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||||
|
let _ = idx.index_adr(&adr).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(format!("Decision logged as {}", adr_id).to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct QueryDecisionsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for QueryDecisionsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"query_decisions"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Execute query_decisions")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.adrs.read_with(|adrs| {
|
||||||
|
let filtered: Vec<_> = adrs.iter().filter(|a| {
|
||||||
|
if let Some(q) = &req.query {
|
||||||
|
contains_ignore_ascii_case(&a.title, q)
|
||||||
|
|| contains_ignore_ascii_case(&a.context, q)
|
||||||
|
|| contains_ignore_ascii_case(&a.decision, q)
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LogErrorFixHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LogErrorFixHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"log_error_fix"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Execute log_error_fix")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LogErrorFixTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.error_fixes.modify(|fixes| {
|
||||||
|
fixes.push(crate::models::ErrorFix {
|
||||||
|
signature: req.signature,
|
||||||
|
solution: req.solution,
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
git_commit: req.git_commit,
|
||||||
|
git_branch: req.git_branch,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("Error fix logged".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SearchErrorFixesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for SearchErrorFixesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"search_error_fixes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<SearchErrorFixesTool>(
|
||||||
|
"search_error_fixes",
|
||||||
|
"Execute search_error_fixes",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: SearchErrorFixesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let q = req.query;
|
||||||
|
let data = state.error_fixes.read_with(|fixes| {
|
||||||
|
let filtered: Vec<_> = fixes.iter().filter(|f| {
|
||||||
|
contains_ignore_ascii_case(&f.signature, &q)
|
||||||
|
|| contains_ignore_ascii_case(&f.solution, &q)
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LogCodeChangeHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LogCodeChangeHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"log_code_change"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Execute log_code_change")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LogCodeChangeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.ledger.modify(|ledger| {
|
||||||
|
ledger.push(CodeChange {
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
file_path: req.file_path,
|
||||||
|
description: req.description,
|
||||||
|
git_commit: req.git_commit,
|
||||||
|
git_branch: req.git_branch,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("Code change logged".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct QueryRecentChangesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for QueryRecentChangesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"query_recent_changes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<QueryRecentChangesTool>(
|
||||||
|
"query_recent_changes",
|
||||||
|
"Execute query_recent_changes",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let data = state.ledger.read_with(|l| serde_json::to_string(l).unwrap_or_else(|_| "[]".to_string()));
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LearnPreferenceHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LearnPreferenceHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"learn_preference"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Execute learn_preference")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LearnPreferenceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.prefs.modify(|prefs| {
|
||||||
|
prefs.insert(
|
||||||
|
req.key.clone(),
|
||||||
|
crate::models::Preference {
|
||||||
|
key: req.key.clone(),
|
||||||
|
value: req.value,
|
||||||
|
updated_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
Ok("Preference learned".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ReadPreferencesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ReadPreferencesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"read_preferences"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Execute read_preferences")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let data = state.prefs.read_with(|prefs| serde_json::to_string(prefs).unwrap_or_default());
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LogTechDebtHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LogTechDebtHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"log_tech_debt"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Execute log_tech_debt")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LogTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.tech_debts.modify(|debts| {
|
||||||
|
debts.push(crate::models::TechDebt {
|
||||||
|
id: uuid::Uuid::new_v4().to_string(),
|
||||||
|
namespace: req.namespace,
|
||||||
|
description: req.description,
|
||||||
|
ideal_solution: req.ideal_solution,
|
||||||
|
is_resolved: false,
|
||||||
|
created_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
git_commit: req.git_commit,
|
||||||
|
git_branch: req.git_branch,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("Tech debt logged".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ResolveTechDebtHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ResolveTechDebtHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"resolve_tech_debt"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ResolveTechDebtTool>(
|
||||||
|
"resolve_tech_debt",
|
||||||
|
"Execute resolve_tech_debt",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ResolveTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut found = false;
|
||||||
|
state.tech_debts.modify(|debts| {
|
||||||
|
for d in debts.iter_mut() {
|
||||||
|
if d.id == req.id {
|
||||||
|
d.is_resolved = true;
|
||||||
|
found = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if found {
|
||||||
|
Ok("Tech debt resolved".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Tech debt not found".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ListTechDebtHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ListTechDebtHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"list_tech_debt"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "Execute list_tech_debt")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ListTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.tech_debts.read_with(|debts| {
|
||||||
|
let filtered: Vec<_> = debts.iter().filter(|d| {
|
||||||
|
d.namespace == req.namespace && (req.include_resolved || !d.is_resolved)
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct OmniSearchHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for OmniSearchHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"omni_search"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Execute omni_search")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let matches = {
|
||||||
|
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner());
|
||||||
|
idx.search(&req.query, req.namespace.as_deref())
|
||||||
|
.unwrap_or_default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let kg_json = state.read_graph(|full| {
|
||||||
|
let mut kg_entities = std::collections::HashMap::new();
|
||||||
|
for (id, doc_type, _, _, _) in &matches {
|
||||||
|
if doc_type == "entity"
|
||||||
|
&& let Some(e) = full.entities.get(id)
|
||||||
|
{
|
||||||
|
kg_entities.insert(id, e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::to_value(&kg_entities).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let tasks_json = state.tasks.read_with(|all_tasks| {
|
||||||
|
let filtered: Vec<_> = all_tasks.iter().filter(|t| {
|
||||||
|
matches.iter().any(|(id, typ, _, _, _)| id == &t.id && typ == "task")
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let snippets_json = state.snippets.read_with(|all_snippets| {
|
||||||
|
let filtered: Vec<_> = all_snippets.iter().filter(|s| {
|
||||||
|
matches.iter().any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet")
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let adrs_json = state.adrs.read_with(|all_adrs| {
|
||||||
|
let filtered: Vec<_> = all_adrs.iter().filter(|a| {
|
||||||
|
matches.iter().any(|(id, typ, _, _, _)| id == &a.id && typ == "adr")
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let q = req.query;
|
||||||
|
let tech_debts_json = state.tech_debts.read_with(|debts| {
|
||||||
|
let filtered: Vec<_> = debts.iter().filter(|d| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| d.namespace == *ns)
|
||||||
|
&& (contains_ignore_ascii_case(&d.description, &q)
|
||||||
|
|| contains_ignore_ascii_case(&d.ideal_solution, &q))
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let memos_json = state.handoff_memos.read_with(|memos| {
|
||||||
|
let filtered: Vec<_> = memos.iter().filter(|m| {
|
||||||
|
req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
|
||||||
|
&& contains_ignore_ascii_case(&m.content, &q)
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let error_fixes_json = state.error_fixes.read_with(|fixes| {
|
||||||
|
let filtered: Vec<_> = fixes.iter().filter(|f| {
|
||||||
|
contains_ignore_ascii_case(&f.signature, &q)
|
||||||
|
|| contains_ignore_ascii_case(&f.solution, &q)
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_value(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
|
||||||
|
let report = serde_json::json!({
|
||||||
|
"knowledge_graph": kg_json,
|
||||||
|
"tasks": tasks_json,
|
||||||
|
"snippets": snippets_json,
|
||||||
|
"adrs": adrs_json,
|
||||||
|
"tech_debts": tech_debts_json,
|
||||||
|
"handoff_memos": memos_json,
|
||||||
|
"error_fixes": error_fixes_json
|
||||||
|
});
|
||||||
|
Ok(report.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct GetProjectHealthHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for GetProjectHealthHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"get_project_health"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<GetProjectHealthTool>(
|
||||||
|
"get_project_health",
|
||||||
|
"Execute get_project_health",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: GetProjectHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let active_tasks = state.tasks.read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count());
|
||||||
|
let unresolved_debt = state.tech_debts.read_with(|debts| debts.iter().filter(|d| d.namespace == req.namespace && !d.is_resolved).count());
|
||||||
|
let unread_memos = state.handoff_memos.read_with(|memos| memos.iter().filter(|m| m.namespace == req.namespace).count());
|
||||||
|
let active_milestones = state.milestones.read_with(|milestones| milestones.iter().filter(|m| m.namespace == req.namespace && m.status != "done").count());
|
||||||
|
let remaining_checklists = state.pr_checklists.read_with(|checklists| checklists.iter().filter(|c| c.namespace == req.namespace).count());
|
||||||
|
|
||||||
|
let report = serde_json::json!({
|
||||||
|
"active_tasks": active_tasks,
|
||||||
|
"unresolved_tech_debt": unresolved_debt,
|
||||||
|
"unread_handoff_memos": unread_memos,
|
||||||
|
"active_milestones": active_milestones,
|
||||||
|
"remaining_pr_checklist_items": remaining_checklists
|
||||||
|
});
|
||||||
|
Ok(report.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
use crate::handlers_v2::utils::*;
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
pub mod env;
|
||||||
|
pub mod graph;
|
||||||
|
pub mod meta;
|
||||||
|
pub mod notes;
|
||||||
|
pub mod tasks;
|
||||||
|
pub mod utils;
|
||||||
|
pub mod workspaces;
|
||||||
@@ -0,0 +1,264 @@
|
|||||||
|
use crate::models::*;
|
||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
pub struct AddStickyNoteHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddStickyNoteHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_sticky_note"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddStickyNoteTool>("add_sticky_note", "Execute add_sticky_note")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddStickyNoteTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.sticky.modify(|notes| {
|
||||||
|
notes.push(StickyNote {
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
content: req.content,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("Sticky note added.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ReadStickyNotesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ReadStickyNotesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"read_sticky_notes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ReadStickyNotesTool>(
|
||||||
|
"read_sticky_notes",
|
||||||
|
"Execute read_sticky_notes",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let data = state.sticky.read_with(|s| serde_json::to_string(s).unwrap_or_else(|_| "[]".to_string()));
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteStickyNoteHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteStickyNoteHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_sticky_note"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteStickyNoteTool>(
|
||||||
|
"delete_sticky_note",
|
||||||
|
"Execute delete_sticky_note",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteStickyNoteTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut success = false;
|
||||||
|
state.sticky.modify(|notes| {
|
||||||
|
if req.index > 0 && req.index <= notes.len() {
|
||||||
|
notes.remove(req.index - 1);
|
||||||
|
success = true;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if success {
|
||||||
|
Ok("Sticky note deleted.".to_string())
|
||||||
|
} else {
|
||||||
|
Err("Invalid sticky note index.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ClearStickyNotesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ClearStickyNotesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"clear_sticky_notes"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ClearStickyNotesTool>(
|
||||||
|
"clear_sticky_notes",
|
||||||
|
"Execute clear_sticky_notes",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
state.sticky.modify(|notes| {
|
||||||
|
notes.clear();
|
||||||
|
});
|
||||||
|
Ok("All sticky notes cleared.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LeaveHandoffMemoHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LeaveHandoffMemoHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"leave_handoff_memo"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LeaveHandoffMemoTool>(
|
||||||
|
"leave_handoff_memo",
|
||||||
|
"Execute leave_handoff_memo",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LeaveHandoffMemoTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.handoff_memos.modify(|memos| {
|
||||||
|
memos.push(crate::models::HandoffMemo {
|
||||||
|
id: uuid::Uuid::new_v4().to_string(),
|
||||||
|
author: "agy".to_string(),
|
||||||
|
content: req.content,
|
||||||
|
namespace: req.namespace,
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("Handoff memo left".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ReadHandoffMemosHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ReadHandoffMemosHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"read_handoff_memos"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ReadHandoffMemosTool>(
|
||||||
|
"read_handoff_memos",
|
||||||
|
"Execute read_handoff_memos",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ReadHandoffMemosTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.handoff_memos.read_with(|items| {
|
||||||
|
let filtered: Vec<_> = items.iter().filter(|i| {
|
||||||
|
if let Some(ns) = &req.namespace {
|
||||||
|
&i.namespace == ns
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ClearHandoffMemosHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ClearHandoffMemosHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"clear_handoff_memos"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ClearHandoffMemosTool>(
|
||||||
|
"clear_handoff_memos",
|
||||||
|
"Execute clear_handoff_memos",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ClearHandoffMemosTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let ids: HashSet<_> = req.ids.into_iter().collect();
|
||||||
|
state
|
||||||
|
.handoff_memos
|
||||||
|
.modify(|memos| memos.retain(|m| !ids.contains(&m.id)));
|
||||||
|
Ok("Handoff memos cleared".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct AddSessionSummaryHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddSessionSummaryHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_session_summary"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddSessionSummaryTool>(
|
||||||
|
"add_session_summary",
|
||||||
|
"Execute add_session_summary",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddSessionSummaryTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.session_summaries.modify(|summaries| {
|
||||||
|
summaries.push(crate::models::SessionSummary {
|
||||||
|
summary: req.summary,
|
||||||
|
namespace: req.namespace,
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("Session summary added".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct GenerateStandupReportHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for GenerateStandupReportHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"generate_standup_report"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<GenerateStandupReportTool>(
|
||||||
|
"generate_standup_report",
|
||||||
|
"Execute generate_standup_report",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: GenerateStandupReportTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let cutoff = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs()
|
||||||
|
.saturating_sub(req.hours_lookback * 3600);
|
||||||
|
let tasks = state.tasks.read_with(|items| items.iter().filter(|t| t.updated_at >= cutoff).cloned().collect::<Vec<_>>());
|
||||||
|
let changes = state.ledger.read_with(|items| items.iter().filter(|c| c.timestamp >= cutoff).cloned().collect::<Vec<_>>());
|
||||||
|
let summaries = state.session_summaries.read_with(|items| items.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).cloned().collect::<Vec<_>>());
|
||||||
|
let report = serde_json::json!({ "tasks_updated": tasks, "code_changes": changes, "session_summaries": summaries });
|
||||||
|
Ok(report.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
@@ -0,0 +1,472 @@
|
|||||||
|
use crate::models::*;
|
||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
pub struct AddTaskHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddTaskHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_task"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddTaskTool>("add_task", "Execute add_task")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let now = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
let task_id = uuid::Uuid::new_v4().to_string();
|
||||||
|
|
||||||
|
let parent_id = req.parent_id.clone();
|
||||||
|
let deps = req.dependencies.clone().unwrap_or_default();
|
||||||
|
|
||||||
|
let task = Task {
|
||||||
|
id: task_id.clone(),
|
||||||
|
title: req.title,
|
||||||
|
status: "pending".to_string(),
|
||||||
|
description: req.description,
|
||||||
|
created_at: now,
|
||||||
|
updated_at: now,
|
||||||
|
git_branch: req.git_branch,
|
||||||
|
parent_id,
|
||||||
|
dependencies: deps,
|
||||||
|
acceptance_criteria: vec![],
|
||||||
|
};
|
||||||
|
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||||
|
let _ = idx.index_task(&task).await;
|
||||||
|
state.tasks.modify(|tasks| {
|
||||||
|
tasks.push(task);
|
||||||
|
});
|
||||||
|
Ok(format!("Task added with ID: {}", task_id).to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteTaskHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteTaskHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_task"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteTaskTool>("delete_task", "Execute delete_task")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut deleted_count = 0;
|
||||||
|
state.tasks.modify(|tasks| {
|
||||||
|
let initial_len = tasks.len();
|
||||||
|
// Collect IDs of tasks to delete (this task + all its recursive children)
|
||||||
|
let mut to_delete = std::collections::HashSet::new();
|
||||||
|
to_delete.insert(req.id.clone());
|
||||||
|
|
||||||
|
let mut children_map: std::collections::HashMap<String, Vec<String>> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
|
for t in tasks.iter() {
|
||||||
|
if let Some(pid) = &t.parent_id {
|
||||||
|
children_map
|
||||||
|
.entry(pid.clone())
|
||||||
|
.or_default()
|
||||||
|
.push(t.id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut queue = std::collections::VecDeque::new();
|
||||||
|
queue.push_back(req.id.clone());
|
||||||
|
|
||||||
|
while let Some(curr) = queue.pop_front() {
|
||||||
|
if to_delete.insert(curr.clone())
|
||||||
|
&& let Some(children) = children_map.get(&curr)
|
||||||
|
{
|
||||||
|
queue.extend(children.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tasks.retain(|t| !to_delete.contains(&t.id));
|
||||||
|
deleted_count = initial_len - tasks.len();
|
||||||
|
});
|
||||||
|
|
||||||
|
if deleted_count > 0 {
|
||||||
|
Ok(vec![
|
||||||
|
format!("Deleted task and its children ({} total).", deleted_count).to_string(),
|
||||||
|
][0]
|
||||||
|
.clone())
|
||||||
|
} else {
|
||||||
|
Ok("Task not found.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct UpdateTaskStatusHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for UpdateTaskStatusHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"update_task_status"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<UpdateTaskStatusTool>(
|
||||||
|
"update_task_status",
|
||||||
|
"Execute update_task_status",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: UpdateTaskStatusTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut found = false;
|
||||||
|
let mut blocked = false;
|
||||||
|
let mut blocker_details = String::new();
|
||||||
|
let target_status = req.status.to_lowercase();
|
||||||
|
|
||||||
|
state.tasks.modify(|tasks| {
|
||||||
|
// Find target task
|
||||||
|
let mut target_id = String::new();
|
||||||
|
if let Some(t) = tasks.iter().find(|t| t.id == req.id || t.title == req.id) {
|
||||||
|
target_id = t.id.clone();
|
||||||
|
}
|
||||||
|
|
||||||
|
if target_id.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
found = true;
|
||||||
|
|
||||||
|
if target_status == "done" || target_status == "completed" {
|
||||||
|
// 1. Check Acceptance Criteria
|
||||||
|
if let Some(t) = tasks.iter().find(|t| t.id == target_id)
|
||||||
|
&& t.acceptance_criteria.iter().any(|c| !c.is_met)
|
||||||
|
{
|
||||||
|
blocked = true;
|
||||||
|
blocker_details = "Unmet acceptance criteria exist.".to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Check dependencies
|
||||||
|
if !blocked {
|
||||||
|
let mut uncompleted_deps = Vec::new();
|
||||||
|
if let Some(t) = tasks.iter().find(|t| t.id == target_id) {
|
||||||
|
for dep_id in &t.dependencies {
|
||||||
|
if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id)
|
||||||
|
&& dep_task.status != "completed"
|
||||||
|
&& dep_task.status != "done"
|
||||||
|
{
|
||||||
|
uncompleted_deps.push(dep_task.title.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !uncompleted_deps.is_empty() {
|
||||||
|
blocked = true;
|
||||||
|
blocker_details =
|
||||||
|
format!("Blocked by dependencies: {}", uncompleted_deps.join(", "));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Check child tasks
|
||||||
|
if !blocked {
|
||||||
|
let mut uncompleted_children = Vec::new();
|
||||||
|
for child in tasks
|
||||||
|
.iter()
|
||||||
|
.filter(|t| t.parent_id.as_ref() == Some(&target_id))
|
||||||
|
{
|
||||||
|
if child.status != "completed" && child.status != "done" {
|
||||||
|
uncompleted_children.push(child.title.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !uncompleted_children.is_empty() {
|
||||||
|
blocked = true;
|
||||||
|
blocker_details = format!(
|
||||||
|
"Blocked by child tasks: {}",
|
||||||
|
uncompleted_children.join(", ")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !blocked {
|
||||||
|
// Apply update
|
||||||
|
if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) {
|
||||||
|
t.status = target_status.clone();
|
||||||
|
t.updated_at = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cascade cancellation to children
|
||||||
|
if target_status == "cancelled" || target_status == "abandoned" {
|
||||||
|
let mut children_map: std::collections::HashMap<usize, Vec<usize>> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
|
|
||||||
|
// First pass: map string ID to index to build the adjacency list by index
|
||||||
|
let mut id_to_idx = std::collections::HashMap::new();
|
||||||
|
for (idx, t) in tasks.iter().enumerate() {
|
||||||
|
id_to_idx.insert(t.id.as_str(), idx);
|
||||||
|
}
|
||||||
|
|
||||||
|
for (idx, t) in tasks.iter().enumerate() {
|
||||||
|
if let Some(pid) = &t.parent_id
|
||||||
|
&& let Some(&p_idx) = id_to_idx.get(pid.as_str()) {
|
||||||
|
children_map.entry(p_idx).or_default().push(idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(&start_idx) = id_to_idx.get(target_id.as_str()) {
|
||||||
|
let mut queue = std::collections::VecDeque::new();
|
||||||
|
queue.push_back(start_idx);
|
||||||
|
|
||||||
|
while let Some(curr) = queue.pop_front() {
|
||||||
|
if let Some(child_indices) = children_map.get(&curr) {
|
||||||
|
for &idx in child_indices {
|
||||||
|
if tasks[idx].status != "completed"
|
||||||
|
&& tasks[idx].status != target_status
|
||||||
|
{
|
||||||
|
tasks[idx].status = target_status.clone();
|
||||||
|
queue.push_back(idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
if blocked {
|
||||||
|
Ok(format!(
|
||||||
|
"Error: Cannot transition task. {}",
|
||||||
|
blocker_details
|
||||||
|
))
|
||||||
|
} else if found {
|
||||||
|
Ok("Task status updated.".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Task not found.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ListActiveTasksHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ListActiveTasksHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"list_active_tasks"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ListActiveTasksTool>(
|
||||||
|
"list_active_tasks",
|
||||||
|
"Execute list_active_tasks",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ListActiveTasksTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.tasks.read_with(|tasks| {
|
||||||
|
let filtered: Vec<_> = tasks.iter().filter(|t| {
|
||||||
|
let status_match = t.status != "done";
|
||||||
|
let branch_match = match &req.git_branch {
|
||||||
|
Some(branch) => t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str()),
|
||||||
|
None => true,
|
||||||
|
};
|
||||||
|
status_match && branch_match
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SetAcceptanceCriteriaHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for SetAcceptanceCriteriaHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"set_acceptance_criteria"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<SetAcceptanceCriteriaTool>(
|
||||||
|
"set_acceptance_criteria",
|
||||||
|
"Execute set_acceptance_criteria",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: SetAcceptanceCriteriaTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut success = false;
|
||||||
|
state.tasks.modify(|tasks| {
|
||||||
|
if let Some(task) = tasks.iter_mut().rev().find(|t| t.title == req.task_title) {
|
||||||
|
task.acceptance_criteria = req
|
||||||
|
.criteria
|
||||||
|
.into_iter()
|
||||||
|
.map(|desc| crate::models::AcceptanceCriteria {
|
||||||
|
id: uuid::Uuid::new_v4().to_string(),
|
||||||
|
description: desc,
|
||||||
|
is_met: false,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
task.updated_at = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
success = true;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if success {
|
||||||
|
Ok("Acceptance criteria set successfully.".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Task not found.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct VerifyAcceptanceCriteriaHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for VerifyAcceptanceCriteriaHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"verify_acceptance_criteria"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<VerifyAcceptanceCriteriaTool>(
|
||||||
|
"verify_acceptance_criteria",
|
||||||
|
"Execute verify_acceptance_criteria",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: VerifyAcceptanceCriteriaTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut success = false;
|
||||||
|
let mut already_met = false;
|
||||||
|
state.tasks.modify(|tasks| {
|
||||||
|
if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id)
|
||||||
|
&& let Some(ac) = task
|
||||||
|
.acceptance_criteria
|
||||||
|
.iter_mut()
|
||||||
|
.find(|c| c.id == req.criteria || c.description == req.criteria)
|
||||||
|
{
|
||||||
|
if ac.is_met {
|
||||||
|
already_met = true;
|
||||||
|
} else {
|
||||||
|
ac.is_met = true;
|
||||||
|
success = true;
|
||||||
|
task.updated_at = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if success {
|
||||||
|
Ok(format!(
|
||||||
|
"Acceptance criteria verified with proof: {}",
|
||||||
|
req.proof
|
||||||
|
))
|
||||||
|
} else if already_met {
|
||||||
|
Ok("Acceptance criteria was already met.".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Acceptance criteria or task not found.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct AddMilestoneHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddMilestoneHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_milestone"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddMilestoneTool>("add_milestone", "Execute add_milestone")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddMilestoneTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.milestones.modify(|ms| {
|
||||||
|
ms.push(crate::models::Milestone {
|
||||||
|
id: uuid::Uuid::new_v4().to_string(),
|
||||||
|
title: req.title,
|
||||||
|
status: "pending".to_string(),
|
||||||
|
namespace: req.namespace,
|
||||||
|
target_date: None,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("Milestone added".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct UpdateMilestoneHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for UpdateMilestoneHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"update_milestone"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<UpdateMilestoneTool>("update_milestone", "Execute update_milestone")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: UpdateMilestoneTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut found = false;
|
||||||
|
state.milestones.modify(|ms| {
|
||||||
|
for m in ms.iter_mut() {
|
||||||
|
if m.id == req.id {
|
||||||
|
m.status = req.status.clone();
|
||||||
|
found = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if found {
|
||||||
|
Ok("Milestone updated".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Milestone not found".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ListMilestonesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ListMilestonesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"list_milestones"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ListMilestonesTool>("list_milestones", "Execute list_milestones")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ListMilestonesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.milestones.read_with(|items| {
|
||||||
|
let filtered: Vec<_> = items.iter().filter(|i| {
|
||||||
|
if let Some(ns) = &req.namespace {
|
||||||
|
&i.namespace == ns
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
pub fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool {
|
||||||
|
if needle.is_empty() {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
haystack
|
||||||
|
.as_bytes()
|
||||||
|
.windows(needle.len())
|
||||||
|
.any(|w| w.eq_ignore_ascii_case(needle.as_bytes()))
|
||||||
|
}
|
||||||
@@ -0,0 +1,359 @@
|
|||||||
|
use crate::models::*;
|
||||||
|
use crate::router::McpTool;
|
||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::tools::*;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
pub struct PinFileHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for PinFileHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"pin_file"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: PinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.pinned_files.modify(|pinned| {
|
||||||
|
pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path);
|
||||||
|
pinned.push(crate::models::PinnedFile {
|
||||||
|
namespace: req.namespace,
|
||||||
|
file_path: req.file_path,
|
||||||
|
timestamp: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
git_branch: req.git_branch,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("File pinned".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct UnpinFileHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for UnpinFileHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"unpin_file"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: UnpinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.pinned_files.modify(|pinned| {
|
||||||
|
pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path)
|
||||||
|
});
|
||||||
|
Ok("File unpinned".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ListPinnedFilesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ListPinnedFilesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"list_pinned_files"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ListPinnedFilesTool>(
|
||||||
|
"list_pinned_files",
|
||||||
|
"Execute list_pinned_files",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ListPinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.pinned_files.read_with(|pinned| {
|
||||||
|
let filtered: Vec<_> = pinned.iter().filter(|p| {
|
||||||
|
let ns_match = match &req.namespace {
|
||||||
|
Some(ns) => &p.namespace == ns,
|
||||||
|
std::option::Option::None => true,
|
||||||
|
};
|
||||||
|
let branch_match = match &req.git_branch {
|
||||||
|
Some(branch) => p.git_branch.is_none() || p.git_branch.as_deref() == Some(branch.as_str()),
|
||||||
|
std::option::Option::None => true,
|
||||||
|
};
|
||||||
|
ns_match && branch_match
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct StoreSnippetHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for StoreSnippetHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"store_snippet"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let snippet = Snippet {
|
||||||
|
name: req.name.clone(),
|
||||||
|
language: req.language,
|
||||||
|
code: req.code,
|
||||||
|
description: req.description,
|
||||||
|
updated_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let s_clone = snippet.clone();
|
||||||
|
state.snippets.modify(|snippets| {
|
||||||
|
snippets.retain(|s| s.name != req.name);
|
||||||
|
snippets.push(s_clone);
|
||||||
|
});
|
||||||
|
|
||||||
|
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||||
|
let _ = idx.index_snippet(&snippet).await;
|
||||||
|
|
||||||
|
Ok(format!("Snippet '{}' stored.", req.name).to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SearchSnippetsHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for SearchSnippetsHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"search_snippets"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<SearchSnippetsTool>("search_snippets", "Execute search_snippets")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: SearchSnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let query = req.query;
|
||||||
|
let data = state.snippets.read_with(|snippets| {
|
||||||
|
let results: Vec<_> = snippets.iter().filter(|s| {
|
||||||
|
contains_ignore_ascii_case(&s.name, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.description, &query)
|
||||||
|
|| contains_ignore_ascii_case(&s.language, &query)
|
||||||
|
}).cloned().collect();
|
||||||
|
serde_json::to_string(&results).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DeleteSnippetHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for DeleteSnippetHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"delete_snippet"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<DeleteSnippetTool>("delete_snippet", "Execute delete_snippet")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: DeleteSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut deleted = false;
|
||||||
|
state.snippets.modify(|snippets| {
|
||||||
|
let orig = snippets.len();
|
||||||
|
snippets.retain(|s| s.name != req.name);
|
||||||
|
deleted = snippets.len() < orig;
|
||||||
|
});
|
||||||
|
if deleted {
|
||||||
|
Ok("Snippet deleted.".to_string())
|
||||||
|
} else {
|
||||||
|
Ok("Snippet not found.".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SaveContextWorkspaceHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for SaveContextWorkspaceHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"save_context_workspace"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<SaveContextWorkspaceTool>(
|
||||||
|
"save_context_workspace",
|
||||||
|
"Execute save_context_workspace",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: SaveContextWorkspaceTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.context_workspaces.modify(|ws| {
|
||||||
|
ws.retain(|w| w.namespace != req.namespace || w.name != req.name);
|
||||||
|
ws.push(crate::models::ContextWorkspace {
|
||||||
|
namespace: req.namespace,
|
||||||
|
name: req.name,
|
||||||
|
pinned_files: req.pinned_files,
|
||||||
|
active_task_ids: req.active_task_ids,
|
||||||
|
saved_at: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
});
|
||||||
|
});
|
||||||
|
Ok("Context workspace saved".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LoadContextWorkspaceHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for LoadContextWorkspaceHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"load_context_workspace"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<LoadContextWorkspaceTool>(
|
||||||
|
"load_context_workspace",
|
||||||
|
"Execute load_context_workspace",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: LoadContextWorkspaceTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.context_workspaces.read_with(|ws| {
|
||||||
|
let filtered: Vec<_> = ws.iter().filter(|w| {
|
||||||
|
w.namespace == req.namespace && w.name == req.name
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered.first()).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ListContextWorkspacesHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ListContextWorkspacesHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"list_context_workspaces"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ListContextWorkspacesTool>(
|
||||||
|
"list_context_workspaces",
|
||||||
|
"Execute list_context_workspaces",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ListContextWorkspacesTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.context_workspaces.read_with(|ws| {
|
||||||
|
let filtered: Vec<_> = ws.iter().filter(|w| {
|
||||||
|
w.namespace == req.namespace
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct AddPrChecklistItemHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for AddPrChecklistItemHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"add_pr_checklist_item"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<AddPrChecklistItemTool>(
|
||||||
|
"add_pr_checklist_item",
|
||||||
|
"Execute add_pr_checklist_item",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: AddPrChecklistItemTool =
|
||||||
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state.pr_checklists.modify(|items| {
|
||||||
|
items.push(crate::models::PrChecklistItem {
|
||||||
|
namespace: req.namespace,
|
||||||
|
id: uuid::Uuid::new_v4().to_string(),
|
||||||
|
description: req.description,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
Ok("PR checklist item added".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct GetPrChecklistHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for GetPrChecklistHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"get_pr_checklist"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<GetPrChecklistTool>("get_pr_checklist", "Execute get_pr_checklist")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: GetPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let data = state.pr_checklists.read_with(|items| {
|
||||||
|
let filtered: Vec<_> = items.iter().filter(|i| {
|
||||||
|
i.namespace == req.namespace
|
||||||
|
}).collect();
|
||||||
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
|
});
|
||||||
|
Ok(data.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ClearPrChecklistHandler;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpTool for ClearPrChecklistHandler {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"clear_pr_checklist"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
crate::mcp::tool_def::<ClearPrChecklistTool>(
|
||||||
|
"clear_pr_checklist",
|
||||||
|
"Execute clear_pr_checklist",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
||||||
|
let req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
state
|
||||||
|
.pr_checklists
|
||||||
|
.modify(|items| items.retain(|i| i.namespace != req.namespace));
|
||||||
|
Ok("PR checklist cleared".to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
use crate::handlers_v2::utils::*;
|
||||||
|
|
||||||
|
|
||||||
+104
-102
@@ -4,8 +4,10 @@
|
|||||||
)]
|
)]
|
||||||
|
|
||||||
mod handlers;
|
mod handlers;
|
||||||
|
mod handlers_v2;
|
||||||
mod mcp;
|
mod mcp;
|
||||||
mod models;
|
mod models;
|
||||||
|
mod router;
|
||||||
mod search;
|
mod search;
|
||||||
mod state;
|
mod state;
|
||||||
mod store;
|
mod store;
|
||||||
@@ -204,7 +206,7 @@ async fn gate_set_handler(
|
|||||||
reason: body.reason.clone(),
|
reason: body.reason.clone(),
|
||||||
timestamp: SystemTime::now()
|
timestamp: SystemTime::now()
|
||||||
.duration_since(UNIX_EPOCH)
|
.duration_since(UNIX_EPOCH)
|
||||||
.unwrap()
|
.unwrap_or_default()
|
||||||
.as_secs(),
|
.as_secs(),
|
||||||
};
|
};
|
||||||
app_state.handler.state.gates.modify(|gates| {
|
app_state.handler.state.gates.modify(|gates| {
|
||||||
@@ -218,13 +220,28 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
state.rebuild_index().await;
|
state.rebuild_index().await;
|
||||||
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
||||||
let app_state = Arc::new(AppState {
|
let app_state = Arc::new(AppState {
|
||||||
handler: Arc::new(MemoryHandler {
|
handler: Arc::new(MemoryHandler::new(Arc::clone(&state))),
|
||||||
state: Arc::clone(&state),
|
|
||||||
}),
|
|
||||||
clients: RwLock::new(HashMap::new()),
|
clients: RwLock::new(HashMap::new()),
|
||||||
next_id: AtomicUsize::new(1),
|
next_id: AtomicUsize::new(1),
|
||||||
});
|
});
|
||||||
|
|
||||||
|
let app_state_clone = Arc::clone(&app_state);
|
||||||
|
let mut rx = state.activity_tx.subscribe();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
while let Ok(msg) = rx.recv().await {
|
||||||
|
let senders: Vec<_> = app_state_clone
|
||||||
|
.clients
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
|
.values()
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
for client_tx in senders {
|
||||||
|
let _ = client_tx.try_send(msg.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
let app = Router::new()
|
let app = Router::new()
|
||||||
.route(
|
.route(
|
||||||
"/api/version",
|
"/api/version",
|
||||||
@@ -307,8 +324,8 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
get({
|
get({
|
||||||
let state_clone = app_state.handler.state.clone();
|
let state_clone = app_state.handler.state.clone();
|
||||||
move || async move {
|
move || async move {
|
||||||
let tasks = state_clone.tasks.read();
|
let tasks = state_clone.tasks.read_with(|t| t.clone());
|
||||||
axum::Json(tasks.clone())
|
axum::Json(tasks)
|
||||||
}
|
}
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
@@ -317,8 +334,8 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
get({
|
get({
|
||||||
let state_clone = app_state.handler.state.clone();
|
let state_clone = app_state.handler.state.clone();
|
||||||
move || async move {
|
move || async move {
|
||||||
let sticky = state_clone.sticky.read();
|
let sticky = state_clone.sticky.read_with(|s| s.clone());
|
||||||
axum::Json(sticky.clone())
|
axum::Json(sticky)
|
||||||
}
|
}
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
@@ -350,33 +367,42 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
}
|
}
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
|
.route(
|
||||||
|
"/api/activity",
|
||||||
|
get({
|
||||||
|
let state_clone = app_state.handler.state.clone();
|
||||||
|
move || async move {
|
||||||
|
let activities = state_clone.recent_activities.read_with(|a| {
|
||||||
|
a.iter().cloned().collect::<Vec<_>>()
|
||||||
|
});
|
||||||
|
axum::Json(serde_json::json!(activities))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
.route(
|
.route(
|
||||||
"/api/stats",
|
"/api/stats",
|
||||||
get({
|
get({
|
||||||
let state_clone = app_state.handler.state.clone();
|
let state_clone = app_state.handler.state.clone();
|
||||||
move || async move {
|
move || async move {
|
||||||
let (entities, relations) = {
|
let (entities, relations) = state_clone.read_graph(|g| (g.entities.len(), g.relations.len()));
|
||||||
let graph = state_clone.get_full_graph();
|
let tasks = state_clone.tasks.read_with(|items| items.len());
|
||||||
(graph.entities.len(), graph.relations.len())
|
let snippets = state_clone.snippets.read_with(|items| items.len());
|
||||||
};
|
let tech_debts = state_clone.tech_debts.read_with(|items| items.len());
|
||||||
let tasks = state_clone.tasks.read().len();
|
let adrs = state_clone.adrs.read_with(|items| items.len());
|
||||||
let snippets = state_clone.snippets.read().len();
|
|
||||||
let tech_debts = state_clone.tech_debts.read().len();
|
|
||||||
let adrs = state_clone.adrs.read().len();
|
|
||||||
|
|
||||||
let ledger = state_clone.ledger.read().len();
|
let ledger = state_clone.ledger.read_with(|items| items.len());
|
||||||
let sticky = state_clone.sticky.read().len();
|
let sticky = state_clone.sticky.read_with(|items| items.len());
|
||||||
let error_fixes = state_clone.error_fixes.read().len();
|
let error_fixes = state_clone.error_fixes.read_with(|items| items.len());
|
||||||
let pinned_files = state_clone.pinned_files.read().len();
|
let pinned_files = state_clone.pinned_files.read_with(|items| items.len());
|
||||||
let session_summaries = state_clone.session_summaries.read().len();
|
let session_summaries = state_clone.session_summaries.read_with(|items| items.len());
|
||||||
let handoff_memos = state_clone.handoff_memos.read().len();
|
let handoff_memos = state_clone.handoff_memos.read_with(|items| items.len());
|
||||||
let env_fingerprints = state_clone.env_fingerprints.read().len();
|
let env_fingerprints = state_clone.env_fingerprints.read_with(|items| items.len());
|
||||||
let env_requirements = state_clone.env_requirements.read().len();
|
let env_requirements = state_clone.env_requirements.read_with(|items| items.len());
|
||||||
let milestones = state_clone.milestones.read().len();
|
let milestones = state_clone.milestones.read_with(|items| items.len());
|
||||||
let environments = state_clone.environments.read().len();
|
let environments = state_clone.environments.read_with(|items| items.len());
|
||||||
let pr_checklists = state_clone.pr_checklists.read().len();
|
let pr_checklists = state_clone.pr_checklists.read_with(|items| items.len());
|
||||||
let gates = state_clone.gates.read().len();
|
let gates = state_clone.gates.read_with(|items| items.len());
|
||||||
let context_workspaces = state_clone.context_workspaces.read().len();
|
let context_workspaces = state_clone.context_workspaces.read_with(|items| items.len());
|
||||||
|
|
||||||
axum::Json(serde_json::json!({
|
axum::Json(serde_json::json!({
|
||||||
"entities": entities,
|
"entities": entities,
|
||||||
@@ -406,7 +432,7 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
|
|
||||||
tracing::info!("MCP Memory Server running on http://127.0.0.1:3000/sse");
|
tracing::info!("MCP Memory Server running on http://127.0.0.1:3000/sse");
|
||||||
let addr = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
let addr = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||||
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr).parse().unwrap();
|
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr).parse().expect("Invalid bind address");
|
||||||
|
|
||||||
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
||||||
Ok(l) => l,
|
Ok(l) => l,
|
||||||
@@ -414,7 +440,8 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
let log_path = dirs::home_dir()
|
let log_path = dirs::home_dir()
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.join(".gemini/mcp_memory/daemon_error.log");
|
.join(".gemini/mcp_memory/daemon_error.log");
|
||||||
let _ = tokio::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e)).await;
|
let _ =
|
||||||
|
tokio::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e)).await;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -441,14 +468,14 @@ async fn ws_handler(
|
|||||||
.into_response()
|
.into_response()
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: String) {
|
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, _client_type: String) {
|
||||||
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
|
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
|
||||||
let (tx, mut rx) = mpsc::channel::<String>(100);
|
let (tx, mut rx) = mpsc::channel::<String>(100);
|
||||||
|
|
||||||
state
|
state
|
||||||
.clients
|
.clients
|
||||||
.write()
|
.write()
|
||||||
.unwrap()
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
.insert(session_id.clone(), tx.clone());
|
.insert(session_id.clone(), tx.clone());
|
||||||
|
|
||||||
let (mut sender, mut receiver) = socket.split();
|
let (mut sender, mut receiver) = socket.split();
|
||||||
@@ -483,49 +510,13 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
|
|||||||
);
|
);
|
||||||
tracing::trace!("Message content: {}", text);
|
tracing::trace!("Message content: {}", text);
|
||||||
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
|
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||||
if client_type == "proxy" {
|
|
||||||
// Send activity broadcast to UI clients
|
|
||||||
if let Some(method) = payload.get("method").and_then(|m| m.as_str())
|
|
||||||
&& method == "tools/call" {
|
|
||||||
let name = payload
|
|
||||||
.get("params")
|
|
||||||
.and_then(|p| p.get("name"))
|
|
||||||
.and_then(|n| n.as_str())
|
|
||||||
.unwrap_or("unknown_tool");
|
|
||||||
let activity_msg = format!("Agent executed tool: {}", name);
|
|
||||||
|
|
||||||
let event = serde_json::json!({
|
|
||||||
"type": "activity",
|
|
||||||
"data": activity_msg
|
|
||||||
});
|
|
||||||
|
|
||||||
let senders: Vec<_> = state_clone
|
|
||||||
.clients
|
|
||||||
.read()
|
|
||||||
.unwrap()
|
|
||||||
.iter()
|
|
||||||
.filter_map(|(id, tx)| {
|
|
||||||
if id != &session_id_clone {
|
|
||||||
Some(tx.clone())
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
for client_tx in senders {
|
|
||||||
let _ = client_tx.try_send(event.to_string());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} // End if proxy
|
|
||||||
|
|
||||||
// Process MCP request
|
// Process MCP request
|
||||||
if let Some(response) = handler.handle_request(payload).await {
|
if let Some(response) = handler.handle_request(payload).await {
|
||||||
let res_str = serde_json::to_string(&response).unwrap();
|
let res_str = serde_json::to_string(&response).unwrap_or_default();
|
||||||
let tx_opt = state_clone
|
let tx_opt = state_clone
|
||||||
.clients
|
.clients
|
||||||
.read()
|
.read()
|
||||||
.unwrap()
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
.get(&session_id_clone)
|
.get(&session_id_clone)
|
||||||
.cloned();
|
.cloned();
|
||||||
if let Some(client_tx) = tx_opt {
|
if let Some(client_tx) = tx_opt {
|
||||||
@@ -570,19 +561,19 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
|
|||||||
struct SessionCleanup {
|
struct SessionCleanup {
|
||||||
session_id: String,
|
session_id: String,
|
||||||
state: Arc<AppState>,
|
state: Arc<AppState>,
|
||||||
send_task: Option<tokio::task::JoinHandle<()>>,
|
send_task: tokio::task::JoinHandle<()>,
|
||||||
recv_task: Option<tokio::task::JoinHandle<()>>,
|
recv_task: tokio::task::JoinHandle<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for SessionCleanup {
|
impl Drop for SessionCleanup {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
self.state.clients.write().unwrap().remove(&self.session_id);
|
self.state
|
||||||
if let Some(task) = self.send_task.take() {
|
.clients
|
||||||
task.abort();
|
.write()
|
||||||
}
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
if let Some(task) = self.recv_task.take() {
|
.remove(&self.session_id);
|
||||||
task.abort();
|
self.send_task.abort();
|
||||||
}
|
self.recv_task.abort();
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Websocket session {} closed and cleaned up",
|
"Websocket session {} closed and cleaned up",
|
||||||
self.session_id
|
self.session_id
|
||||||
@@ -593,15 +584,15 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
|
|||||||
let mut cleanup = SessionCleanup {
|
let mut cleanup = SessionCleanup {
|
||||||
session_id: session_id.clone(),
|
session_id: session_id.clone(),
|
||||||
state: Arc::clone(&state),
|
state: Arc::clone(&state),
|
||||||
send_task: Some(send_task),
|
send_task,
|
||||||
recv_task: Some(recv_task),
|
recv_task,
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = cleanup.send_task.as_mut().unwrap() => {
|
_ = &mut cleanup.send_task => {
|
||||||
tracing::info!("Websocket send task finished for session {}", session_id);
|
tracing::info!("Websocket send task finished for session {}", session_id);
|
||||||
},
|
},
|
||||||
_ = cleanup.recv_task.as_mut().unwrap() => {
|
_ = &mut cleanup.recv_task => {
|
||||||
tracing::info!("Websocket recv task finished for session {}", session_id);
|
tracing::info!("Websocket recv task finished for session {}", session_id);
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
@@ -640,7 +631,13 @@ async fn nvim_telemetry_handler(
|
|||||||
});
|
});
|
||||||
|
|
||||||
let msg_str = ws_msg.to_string();
|
let msg_str = ws_msg.to_string();
|
||||||
let senders: Vec<_> = state.clients.read().unwrap().values().cloned().collect();
|
let senders: Vec<_> = state
|
||||||
|
.clients
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
|
.values()
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
for tx in senders {
|
for tx in senders {
|
||||||
let _ = tx.try_send(msg_str.clone());
|
let _ = tx.try_send(msg_str.clone());
|
||||||
}
|
}
|
||||||
@@ -698,9 +695,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
let mut cmd = std::process::Command::new("curl");
|
let mut cmd = std::process::Command::new("curl");
|
||||||
cmd.arg("-k").arg("-X").arg("POST");
|
cmd.arg("-k").arg("-X").arg("POST");
|
||||||
if !token.is_empty() {
|
if !token.is_empty() {
|
||||||
cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
|
cmd.arg("-H")
|
||||||
|
.arg(format!("Authorization: Bearer {}", token.trim()));
|
||||||
}
|
}
|
||||||
let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
|
let _ = cmd
|
||||||
|
.arg(format!("https://127.0.0.1:{}/shutdown", port))
|
||||||
|
.output();
|
||||||
println!("Sent shutdown request to server.");
|
println!("Sent shutdown request to server.");
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
@@ -711,9 +711,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
let mut cmd = std::process::Command::new("curl");
|
let mut cmd = std::process::Command::new("curl");
|
||||||
cmd.arg("-k").arg("-X").arg("POST");
|
cmd.arg("-k").arg("-X").arg("POST");
|
||||||
if !token.is_empty() {
|
if !token.is_empty() {
|
||||||
cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
|
cmd.arg("-H")
|
||||||
|
.arg(format!("Authorization: Bearer {}", token.trim()));
|
||||||
}
|
}
|
||||||
let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
|
let _ = cmd
|
||||||
|
.arg(format!("https://127.0.0.1:{}/shutdown", port))
|
||||||
|
.output();
|
||||||
println!("Sent shutdown request to existing server. Waiting for it to exit...");
|
println!("Sent shutdown request to existing server. Waiting for it to exit...");
|
||||||
std::thread::sleep(std::time::Duration::from_millis(1500));
|
std::thread::sleep(std::time::Duration::from_millis(1500));
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -725,7 +728,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
if !cli.daemon {
|
if !cli.daemon {
|
||||||
// Just spawn the daemon and exit. We no longer act as a proxy.
|
// Just spawn the daemon and exit. We no longer act as a proxy.
|
||||||
#[allow(clippy::zombie_processes)]
|
#[allow(clippy::zombie_processes)]
|
||||||
let _ = std::process::Command::new(std::env::current_exe().unwrap())
|
let _ = std::process::Command::new(std::env::current_exe().expect("Failed to get current executable path"))
|
||||||
.arg("--daemon")
|
.arg("--daemon")
|
||||||
.stdin(std::process::Stdio::null())
|
.stdin(std::process::Stdio::null())
|
||||||
.stdout(std::process::Stdio::null())
|
.stdout(std::process::Stdio::null())
|
||||||
@@ -745,13 +748,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
|
|
||||||
let redb_path = base.join("mcp_store.redb");
|
let redb_path = base.join("mcp_store.redb");
|
||||||
|
|
||||||
let db = Arc::new(redb::Database::create(&redb_path).unwrap());
|
let db = Arc::new(redb::Database::create(&redb_path).expect("Failed to create redb database"));
|
||||||
|
|
||||||
// Ensure table exists and migrate old JSON files
|
// Ensure table exists and migrate old JSON files
|
||||||
{
|
{
|
||||||
let write_txn = db.begin_write().unwrap();
|
let write_txn = db.begin_write().expect("Failed to begin write txn on redb");
|
||||||
{
|
{
|
||||||
let mut table = write_txn.open_table(crate::store::STORE_TABLE).unwrap();
|
let mut table = write_txn.open_table(crate::store::STORE_TABLE).expect("Failed to open STORE_TABLE");
|
||||||
|
|
||||||
let stores = vec![
|
let stores = vec![
|
||||||
("knowledge_graph_master", "knowledge_graph_master.json"),
|
("knowledge_graph_master", "knowledge_graph_master.json"),
|
||||||
@@ -776,24 +779,22 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
];
|
];
|
||||||
|
|
||||||
for (key, file_name) in stores.iter() {
|
for (key, file_name) in stores.iter() {
|
||||||
if table.get(*key).unwrap().is_none() {
|
if table.get(*key).expect("Failed to read from table").is_none() {
|
||||||
let json_path = base.join(file_name);
|
let json_path = base.join(file_name);
|
||||||
if json_path.exists()
|
if json_path.exists()
|
||||||
&& let Ok(data) = fs::read(&json_path)
|
&& let Ok(data) = fs::read(&json_path)
|
||||||
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok() {
|
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
|
||||||
table.insert(*key, data.as_slice()).unwrap();
|
{
|
||||||
let _ = fs::rename(
|
table.insert(*key, data.as_slice()).expect("Failed to insert migrated data");
|
||||||
&json_path,
|
let _ = fs::rename(&json_path, json_path.with_extension("json.migrated"));
|
||||||
json_path.with_extension("json.migrated"),
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
write_txn.commit().unwrap();
|
write_txn.commit().expect("Failed to commit db migration");
|
||||||
}
|
}
|
||||||
|
|
||||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
let rt = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
|
||||||
let _guard = rt.enter();
|
let _guard = rt.enter();
|
||||||
|
|
||||||
let state = Arc::new(MemoryState {
|
let state = Arc::new(MemoryState {
|
||||||
@@ -827,6 +828,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
tech_debts: Store::new("tech_debts", db.clone()),
|
tech_debts: Store::new("tech_debts", db.clone()),
|
||||||
gates: Store::new("gates", db.clone()),
|
gates: Store::new("gates", db.clone()),
|
||||||
context_workspaces: Store::new("context_workspaces", db.clone()),
|
context_workspaces: Store::new("context_workspaces", db.clone()),
|
||||||
|
recent_activities: Store::new("recent_activities", db.clone()),
|
||||||
activity_tx: tokio::sync::broadcast::channel(100).0,
|
activity_tx: tokio::sync::broadcast::channel(100).0,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
+2
-1
@@ -22,7 +22,7 @@ pub fn error(id: serde_json::Value, code: i32, message: &str) -> serde_json::Val
|
|||||||
|
|
||||||
pub fn tool_def<T: JsonSchema>(name: &str, description: &str) -> serde_json::Value {
|
pub fn tool_def<T: JsonSchema>(name: &str, description: &str) -> serde_json::Value {
|
||||||
let schema = schemars::schema_for!(T);
|
let schema = schemars::schema_for!(T);
|
||||||
let schema_val = serde_json::to_value(&schema).unwrap();
|
let schema_val = serde_json::to_value(&schema).unwrap_or_default();
|
||||||
// MCP expects standard JSON schema. Schemars returns draft-07.
|
// MCP expects standard JSON schema. Schemars returns draft-07.
|
||||||
json!({
|
json!({
|
||||||
"name": name,
|
"name": name,
|
||||||
@@ -30,3 +30,4 @@ pub fn tool_def<T: JsonSchema>(name: &str, description: &str) -> serde_json::Val
|
|||||||
"inputSchema": schema_val
|
"inputSchema": schema_val
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -0,0 +1,173 @@
|
|||||||
|
import os
|
||||||
|
import re
|
||||||
|
|
||||||
|
GROUPS = {
|
||||||
|
"graph": [
|
||||||
|
"query_graph_path", "create_entities", "create_relations", "add_observations",
|
||||||
|
"delete_entities", "delete_observations", "delete_relations", "read_graph",
|
||||||
|
"search_nodes", "open_nodes", "visualize_graph", "condense_entity",
|
||||||
|
"merge_entities", "find_orphans"
|
||||||
|
],
|
||||||
|
"tasks": [
|
||||||
|
"add_task", "delete_task", "update_task_status", "list_active_tasks",
|
||||||
|
"set_acceptance_criteria", "verify_acceptance_criteria",
|
||||||
|
"add_milestone", "update_milestone", "list_milestones"
|
||||||
|
],
|
||||||
|
"notes": [
|
||||||
|
"add_sticky_note", "read_sticky_notes", "delete_sticky_note", "clear_sticky_notes",
|
||||||
|
"leave_handoff_memo", "read_handoff_memos", "clear_handoff_memos",
|
||||||
|
"add_session_summary", "generate_standup_report"
|
||||||
|
],
|
||||||
|
"meta": [
|
||||||
|
"log_decision", "query_decisions", "log_error_fix", "search_error_fixes",
|
||||||
|
"log_code_change", "query_recent_changes", "learn_preference", "read_preferences",
|
||||||
|
"log_tech_debt", "resolve_tech_debt", "list_tech_debt", "omni_search", "get_project_health"
|
||||||
|
],
|
||||||
|
"env": [
|
||||||
|
"update_env_fingerprint", "read_env_fingerprint", "log_env_requirement",
|
||||||
|
"register_environment", "get_environment_details"
|
||||||
|
],
|
||||||
|
"workspaces": [
|
||||||
|
"pin_file", "unpin_file", "list_pinned_files", "store_snippet", "search_snippets",
|
||||||
|
"delete_snippet", "save_context_workspace", "load_context_workspace",
|
||||||
|
"list_context_workspaces", "add_pr_checklist_item", "get_pr_checklist", "clear_pr_checklist"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
def to_camel_case(snake_str):
|
||||||
|
components = snake_str.split('_')
|
||||||
|
return "".join(x.title() for x in components)
|
||||||
|
|
||||||
|
def parse_rust_match(file_path):
|
||||||
|
with open(file_path, "r", encoding="utf-8") as f:
|
||||||
|
lines = f.readlines()
|
||||||
|
|
||||||
|
start_idx = -1
|
||||||
|
for i, line in enumerate(lines):
|
||||||
|
if "let result: Result<String, String> = match name {" in line:
|
||||||
|
start_idx = i
|
||||||
|
break
|
||||||
|
|
||||||
|
if start_idx == -1:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
brace_depth = 1
|
||||||
|
i = start_idx + 1
|
||||||
|
|
||||||
|
tools = {}
|
||||||
|
current_tool = None
|
||||||
|
current_body = []
|
||||||
|
|
||||||
|
while i < len(lines):
|
||||||
|
line = lines[i]
|
||||||
|
|
||||||
|
if brace_depth == 1 and "=>" in line and '"' in line:
|
||||||
|
parts = line.strip().split('"')
|
||||||
|
if len(parts) >= 3:
|
||||||
|
tool_name = parts[1]
|
||||||
|
current_tool = tool_name
|
||||||
|
current_body = []
|
||||||
|
# Don't add the "name" => { line
|
||||||
|
|
||||||
|
if current_tool is not None and not (brace_depth == 1 and "=>" in line and '"' in line):
|
||||||
|
# check if this line closes the block
|
||||||
|
next_depth = brace_depth + line.count('{') - line.count('}')
|
||||||
|
if next_depth == 1 and current_tool is not None:
|
||||||
|
# This is the closing brace
|
||||||
|
tools[current_tool] = "".join(current_body)
|
||||||
|
current_tool = None
|
||||||
|
else:
|
||||||
|
current_body.append(line)
|
||||||
|
|
||||||
|
brace_depth += line.count('{')
|
||||||
|
brace_depth -= line.count('}')
|
||||||
|
|
||||||
|
if brace_depth == 0:
|
||||||
|
break
|
||||||
|
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
return tools
|
||||||
|
|
||||||
|
def transform_body(body):
|
||||||
|
# Transform parse_tool!
|
||||||
|
body = re.sub(
|
||||||
|
r'let req = parse_tool!\(args, id, ([^)]+)\);',
|
||||||
|
r'let req: \1 = serde_json::from_value(args).map_err(|e| e.to_string())?;',
|
||||||
|
body
|
||||||
|
)
|
||||||
|
# Transform handle_list_with_namespace!
|
||||||
|
def repl_handle_list(m):
|
||||||
|
store = m.group(1)
|
||||||
|
tool_type = m.group(2)
|
||||||
|
return f"""
|
||||||
|
let req: {tool_type} = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
|
let mut items = state.{store}.read();
|
||||||
|
if let Some(ns) = req.namespace {{
|
||||||
|
items.retain(|i| i.namespace == ns);
|
||||||
|
}}
|
||||||
|
let data = serde_json::to_string(&items).unwrap_or_default();
|
||||||
|
return Ok(data.to_string());
|
||||||
|
"""
|
||||||
|
body = re.sub(
|
||||||
|
r'return handle_list_with_namespace!\(self, ([^,]+), ([^,]+), args, id\);',
|
||||||
|
repl_handle_list,
|
||||||
|
body
|
||||||
|
)
|
||||||
|
|
||||||
|
# Replace self.state with state
|
||||||
|
body = body.replace("self.state.", "state.")
|
||||||
|
|
||||||
|
return body
|
||||||
|
|
||||||
|
|
||||||
|
tools = parse_rust_match("server/src/handlers.rs")
|
||||||
|
|
||||||
|
for group, tool_names in GROUPS.items():
|
||||||
|
file_path = f"server/src/handlers_v2/{group}.rs"
|
||||||
|
with open(file_path, "w", encoding="utf-8") as f:
|
||||||
|
f.write("use crate::router::McpTool;\n")
|
||||||
|
f.write("use crate::state::MemoryState;\n")
|
||||||
|
f.write("use crate::tools::*;\n")
|
||||||
|
f.write("use async_trait::async_trait;\n")
|
||||||
|
f.write("use serde_json::Value;\n")
|
||||||
|
f.write("use std::sync::Arc;\n")
|
||||||
|
f.write("use std::time::{SystemTime, UNIX_EPOCH};\n\n")
|
||||||
|
|
||||||
|
for name in tool_names:
|
||||||
|
if name not in tools:
|
||||||
|
continue
|
||||||
|
|
||||||
|
body = tools[name]
|
||||||
|
# special case for query_graph_path which we already wrote properly?
|
||||||
|
# actually we will just overwrite it with the transformed body
|
||||||
|
body = transform_body(body)
|
||||||
|
|
||||||
|
struct_name = to_camel_case(name) + "Handler"
|
||||||
|
tool_type = to_camel_case(name) + "Tool"
|
||||||
|
|
||||||
|
f.write(f"pub struct {struct_name};\n\n")
|
||||||
|
f.write(f"#[async_trait]\n")
|
||||||
|
f.write(f"impl McpTool for {struct_name} {{\n")
|
||||||
|
f.write(f" fn name(&self) -> &'static str {{\n")
|
||||||
|
f.write(f' "{name}"\n')
|
||||||
|
f.write(f" }}\n\n")
|
||||||
|
f.write(f" fn schema(&self) -> Value {{\n")
|
||||||
|
# For schema description we can just put a generic one or extract it.
|
||||||
|
# I will use a generic one for now, or you can extract it from tools/list.
|
||||||
|
f.write(f' crate::mcp::tool_def::<{tool_type}>(\n')
|
||||||
|
f.write(f' "{name}",\n')
|
||||||
|
f.write(f' "Execute {name}",\n')
|
||||||
|
f.write(f' )\n')
|
||||||
|
f.write(f" }}\n\n")
|
||||||
|
f.write(f" async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {{\n")
|
||||||
|
f.write(body)
|
||||||
|
f.write(f" }}\n")
|
||||||
|
f.write(f"}}\n\n")
|
||||||
|
|
||||||
|
print("Generated handlers_v2 modules")
|
||||||
|
|
||||||
|
# generate mod.rs
|
||||||
|
with open("server/src/handlers_v2/mod.rs", "w", encoding="utf-8") as f:
|
||||||
|
for group in GROUPS.keys():
|
||||||
|
f.write(f"pub mod {group};\n")
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
use std::fs;
|
||||||
|
use std::io::Write;
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
let content = fs::read_to_string("server/src/handlers.rs").unwrap();
|
||||||
|
println!("Read {} bytes", content.len());
|
||||||
|
// Find the match name { block
|
||||||
|
let match_start = content.find("match name {").unwrap();
|
||||||
|
// naive extraction
|
||||||
|
println!("Found match block at {}", match_start);
|
||||||
|
}
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
import re
|
||||||
|
import os
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "r", encoding="utf-8") as f:
|
||||||
|
lines = f.readlines()
|
||||||
|
|
||||||
|
list_start = -1
|
||||||
|
for i, line in enumerate(lines):
|
||||||
|
if '"tools/list" => {' in line:
|
||||||
|
list_start = i
|
||||||
|
break
|
||||||
|
|
||||||
|
# Find end of tools/call
|
||||||
|
call_start = -1
|
||||||
|
for i in range(list_start, len(lines)):
|
||||||
|
if '"tools/call" => {' in line:
|
||||||
|
call_start = i
|
||||||
|
break
|
||||||
|
|
||||||
|
# Find end of tools/call
|
||||||
|
# Match brace depth from call_start
|
||||||
|
brace_depth = 1
|
||||||
|
call_end = -1
|
||||||
|
for i in range(call_start + 1, len(lines)):
|
||||||
|
brace_depth += lines[i].count('{')
|
||||||
|
brace_depth -= lines[i].count('}')
|
||||||
|
if brace_depth == 0:
|
||||||
|
call_end = i
|
||||||
|
break
|
||||||
|
|
||||||
|
# replacement block
|
||||||
|
replacement = """ "tools/list" => {
|
||||||
|
let mut tools: Vec<serde_json::Value> = self.tools.values().map(|t| t.schema()).collect();
|
||||||
|
tools.sort_by_key(|t| t.get("name").and_then(|n| n.as_str()).unwrap_or("").to_string());
|
||||||
|
Some(crate::mcp::success(
|
||||||
|
id,
|
||||||
|
serde_json::json!({ "tools": tools }),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
"tools/call" => {
|
||||||
|
let params = req.get("params").unwrap_or(&serde_json::Value::Null);
|
||||||
|
let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
|
||||||
|
let args = params
|
||||||
|
.get("arguments")
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||||
|
|
||||||
|
self.state
|
||||||
|
.broadcast_activity(&format!("Agent executed tool: {}", name));
|
||||||
|
|
||||||
|
let result: Result<String, String> = if let Some(tool) = self.tools.get(name) {
|
||||||
|
tool.execute(args, self.state.clone()).await
|
||||||
|
} else {
|
||||||
|
Err(format!("Unknown tool: {}", name))
|
||||||
|
};
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(text) => {
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"content": [{"type": "text", "text": text}],
|
||||||
|
"isError": false
|
||||||
|
});
|
||||||
|
Some(crate::mcp::success(id_clone, payload))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Tool {} failed: {}", name, e);
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"content": [{"type": "text", "text": e}],
|
||||||
|
"isError": true
|
||||||
|
});
|
||||||
|
Some(crate::mcp::success(id_clone, payload))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
new_lines = lines[:list_start] + [replacement] + lines[call_end+1:]
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "w", encoding="utf-8") as f:
|
||||||
|
f.writelines(new_lines)
|
||||||
|
|
||||||
|
print("tools/list and tools/call replaced.")
|
||||||
@@ -0,0 +1,106 @@
|
|||||||
|
import re
|
||||||
|
import os
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "r", encoding="utf-8") as f:
|
||||||
|
content = f.read()
|
||||||
|
|
||||||
|
# Replace MemoryHandler struct
|
||||||
|
struct_pattern = r'pub struct MemoryHandler \{\s*pub state: Arc<MemoryState>,\s*\}'
|
||||||
|
|
||||||
|
new_struct = """use crate::router::McpTool;
|
||||||
|
|
||||||
|
pub struct MemoryHandler {
|
||||||
|
pub state: Arc<MemoryState>,
|
||||||
|
pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MemoryHandler {
|
||||||
|
pub fn new(state: Arc<MemoryState>) -> Self {
|
||||||
|
let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> = std::collections::HashMap::new();
|
||||||
|
|
||||||
|
macro_rules! register {
|
||||||
|
($module:ident::$handler:ident) => {
|
||||||
|
let h = crate::handlers_v2::$module::$handler;
|
||||||
|
tools.insert(h.name().to_string(), Box::new(h));
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
register!(graph::QueryGraphPathHandler);
|
||||||
|
register!(graph::CreateEntitiesHandler);
|
||||||
|
register!(graph::CreateRelationsHandler);
|
||||||
|
register!(graph::AddObservationsHandler);
|
||||||
|
register!(graph::DeleteEntitiesHandler);
|
||||||
|
register!(graph::DeleteObservationsHandler);
|
||||||
|
register!(graph::DeleteRelationsHandler);
|
||||||
|
register!(graph::ReadGraphHandler);
|
||||||
|
register!(graph::SearchNodesHandler);
|
||||||
|
register!(graph::OpenNodesHandler);
|
||||||
|
register!(graph::VisualizeGraphHandler);
|
||||||
|
register!(graph::CondenseEntityHandler);
|
||||||
|
register!(graph::MergeEntitiesHandler);
|
||||||
|
register!(graph::FindOrphansHandler);
|
||||||
|
|
||||||
|
register!(tasks::AddTaskHandler);
|
||||||
|
register!(tasks::DeleteTaskHandler);
|
||||||
|
register!(tasks::UpdateTaskStatusHandler);
|
||||||
|
register!(tasks::ListActiveTasksHandler);
|
||||||
|
register!(tasks::SetAcceptanceCriteriaHandler);
|
||||||
|
register!(tasks::VerifyAcceptanceCriteriaHandler);
|
||||||
|
register!(tasks::AddMilestoneHandler);
|
||||||
|
register!(tasks::UpdateMilestoneHandler);
|
||||||
|
register!(tasks::ListMilestonesHandler);
|
||||||
|
|
||||||
|
register!(notes::AddStickyNoteHandler);
|
||||||
|
register!(notes::ReadStickyNotesHandler);
|
||||||
|
register!(notes::DeleteStickyNoteHandler);
|
||||||
|
register!(notes::ClearStickyNotesHandler);
|
||||||
|
register!(notes::LeaveHandoffMemoHandler);
|
||||||
|
register!(notes::ReadHandoffMemosHandler);
|
||||||
|
register!(notes::ClearHandoffMemosHandler);
|
||||||
|
register!(notes::AddSessionSummaryHandler);
|
||||||
|
register!(notes::GenerateStandupReportHandler);
|
||||||
|
|
||||||
|
register!(meta::LogDecisionHandler);
|
||||||
|
register!(meta::QueryDecisionsHandler);
|
||||||
|
register!(meta::LogErrorFixHandler);
|
||||||
|
register!(meta::SearchErrorFixesHandler);
|
||||||
|
register!(meta::LogCodeChangeHandler);
|
||||||
|
register!(meta::QueryRecentChangesHandler);
|
||||||
|
register!(meta::LearnPreferenceHandler);
|
||||||
|
register!(meta::ReadPreferencesHandler);
|
||||||
|
register!(meta::LogTechDebtHandler);
|
||||||
|
register!(meta::ResolveTechDebtHandler);
|
||||||
|
register!(meta::ListTechDebtHandler);
|
||||||
|
register!(meta::OmniSearchHandler);
|
||||||
|
register!(meta::GetProjectHealthHandler);
|
||||||
|
|
||||||
|
register!(env::UpdateEnvFingerprintHandler);
|
||||||
|
register!(env::ReadEnvFingerprintHandler);
|
||||||
|
register!(env::LogEnvRequirementHandler);
|
||||||
|
register!(env::RegisterEnvironmentHandler);
|
||||||
|
register!(env::GetEnvironmentDetailsHandler);
|
||||||
|
|
||||||
|
register!(workspaces::PinFileHandler);
|
||||||
|
register!(workspaces::UnpinFileHandler);
|
||||||
|
register!(workspaces::ListPinnedFilesHandler);
|
||||||
|
register!(workspaces::StoreSnippetHandler);
|
||||||
|
register!(workspaces::SearchSnippetsHandler);
|
||||||
|
register!(workspaces::DeleteSnippetHandler);
|
||||||
|
register!(workspaces::SaveContextWorkspaceHandler);
|
||||||
|
register!(workspaces::LoadContextWorkspaceHandler);
|
||||||
|
register!(workspaces::ListContextWorkspacesHandler);
|
||||||
|
register!(workspaces::AddPrChecklistItemHandler);
|
||||||
|
register!(workspaces::GetPrChecklistHandler);
|
||||||
|
register!(workspaces::ClearPrChecklistHandler);
|
||||||
|
|
||||||
|
Self { state, tools }
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
content = re.sub(struct_pattern, new_struct, content)
|
||||||
|
content = content.replace("impl MemoryHandler {\n pub async fn handle_request", " pub async fn handle_request")
|
||||||
|
|
||||||
|
with open("server/src/handlers.rs", "w", encoding="utf-8") as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
print("MemoryHandler struct updated.")
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
use crate::state::MemoryState;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
pub trait McpTool: Send + Sync {
|
||||||
|
/// The unique name of the tool
|
||||||
|
fn name(&self) -> &'static str;
|
||||||
|
|
||||||
|
/// The JSON schema for the tool
|
||||||
|
fn schema(&self) -> Value;
|
||||||
|
|
||||||
|
/// Execute the tool with the given arguments
|
||||||
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String>;
|
||||||
|
}
|
||||||
+87
-113
@@ -30,11 +30,14 @@ impl MemoryIndex {
|
|||||||
let schema = schema_builder.build();
|
let schema = schema_builder.build();
|
||||||
|
|
||||||
let index_dir = store_dir.join("tantivy_index");
|
let index_dir = store_dir.join("tantivy_index");
|
||||||
std::fs::create_dir_all(&index_dir).unwrap();
|
std::fs::create_dir_all(&index_dir)
|
||||||
|
.map_err(|e| tantivy::TantivyError::SystemError(e.to_string()))?;
|
||||||
let index = Index::open_in_dir(&index_dir)
|
let index = Index::open_in_dir(&index_dir)
|
||||||
.unwrap_or_else(|_| Index::create_in_dir(&index_dir, schema.clone()).unwrap());
|
.or_else(|_| Index::create_in_dir(&index_dir, schema.clone()))?;
|
||||||
|
|
||||||
let writer = index.writer(50_000_000)?;
|
let mut writer = index.writer(50_000_000)?;
|
||||||
|
writer.delete_all_documents()?;
|
||||||
|
writer.commit()?;
|
||||||
let reader = index
|
let reader = index
|
||||||
.reader_builder()
|
.reader_builder()
|
||||||
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
||||||
@@ -54,22 +57,17 @@ impl MemoryIndex {
|
|||||||
|
|
||||||
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||||
let writer = Arc::clone(&self.writer);
|
let writer = Arc::clone(&self.writer);
|
||||||
let e = e.clone();
|
|
||||||
let id_field = self.id_field;
|
let doc = doc!(
|
||||||
let title_field = self.title_field;
|
self.id_field => e.name.as_str(),
|
||||||
let body_field = self.body_field;
|
self.title_field => e.name.as_str(),
|
||||||
let type_field = self.type_field;
|
self.body_field => e.observations.join(" "),
|
||||||
let namespace_field = self.namespace_field;
|
self.type_field => "entity",
|
||||||
|
self.namespace_field => e.namespace.as_str()
|
||||||
|
);
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let doc = doc!(
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
id_field => e.name.clone(),
|
|
||||||
title_field => e.name.clone(),
|
|
||||||
body_field => e.observations.join(" "),
|
|
||||||
type_field => "entity",
|
|
||||||
namespace_field => e.namespace.clone()
|
|
||||||
);
|
|
||||||
let writer = writer.lock().unwrap();
|
|
||||||
writer.add_document(doc)?;
|
writer.add_document(doc)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
@@ -77,22 +75,17 @@ impl MemoryIndex {
|
|||||||
|
|
||||||
pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||||
let writer = Arc::clone(&self.writer);
|
let writer = Arc::clone(&self.writer);
|
||||||
let t = t.clone();
|
|
||||||
let id_field = self.id_field;
|
let doc = doc!(
|
||||||
let title_field = self.title_field;
|
self.id_field => t.id.as_str(),
|
||||||
let body_field = self.body_field;
|
self.title_field => t.title.as_str(),
|
||||||
let type_field = self.type_field;
|
self.body_field => t.description.as_str(),
|
||||||
let namespace_field = self.namespace_field;
|
self.type_field => "task",
|
||||||
|
self.namespace_field => "global"
|
||||||
|
);
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let doc = doc!(
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
id_field => t.id.clone(),
|
|
||||||
title_field => t.title.clone(),
|
|
||||||
body_field => t.description.clone(),
|
|
||||||
type_field => "task",
|
|
||||||
namespace_field => "global"
|
|
||||||
);
|
|
||||||
let writer = writer.lock().unwrap();
|
|
||||||
writer.add_document(doc)?;
|
writer.add_document(doc)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
@@ -101,12 +94,16 @@ impl MemoryIndex {
|
|||||||
pub async fn commit(&self) -> tantivy::Result<()> {
|
pub async fn commit(&self) -> tantivy::Result<()> {
|
||||||
let writer = Arc::clone(&self.writer);
|
let writer = Arc::clone(&self.writer);
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let mut writer = writer.lock().unwrap();
|
let mut writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
writer.commit()?;
|
writer.commit()?;
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
.unwrap_or_else(|_| Err(tantivy::TantivyError::SystemError("Commit task panicked".to_string())))
|
.unwrap_or_else(|_| {
|
||||||
|
Err(tantivy::TantivyError::SystemError(
|
||||||
|
"Commit task panicked".to_string(),
|
||||||
|
))
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn search(
|
pub fn search(
|
||||||
@@ -165,22 +162,17 @@ impl MemoryIndex {
|
|||||||
|
|
||||||
pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||||
let writer = Arc::clone(&self.writer);
|
let writer = Arc::clone(&self.writer);
|
||||||
let s = s.clone();
|
|
||||||
let id_field = self.id_field;
|
let doc = doc!(
|
||||||
let title_field = self.title_field;
|
self.id_field => s.name.as_str(),
|
||||||
let body_field = self.body_field;
|
self.title_field => s.name.as_str(),
|
||||||
let type_field = self.type_field;
|
self.body_field => format!("{} {}", s.language, s.description),
|
||||||
let namespace_field = self.namespace_field;
|
self.type_field => "snippet",
|
||||||
|
self.namespace_field => "global"
|
||||||
|
);
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let doc = doc!(
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
id_field => s.name.clone(),
|
|
||||||
title_field => s.name.clone(),
|
|
||||||
body_field => format!("{} {}", s.language, s.description),
|
|
||||||
type_field => "snippet",
|
|
||||||
namespace_field => "global"
|
|
||||||
);
|
|
||||||
let writer = writer.lock().unwrap();
|
|
||||||
writer.add_document(doc)?;
|
writer.add_document(doc)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
@@ -188,86 +180,68 @@ impl MemoryIndex {
|
|||||||
|
|
||||||
pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||||
let writer = Arc::clone(&self.writer);
|
let writer = Arc::clone(&self.writer);
|
||||||
let a = a.clone();
|
|
||||||
let id_field = self.id_field;
|
let doc = doc!(
|
||||||
let title_field = self.title_field;
|
self.id_field => a.id.as_str(),
|
||||||
let body_field = self.body_field;
|
self.title_field => a.title.as_str(),
|
||||||
let type_field = self.type_field;
|
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
||||||
let namespace_field = self.namespace_field;
|
self.type_field => "adr",
|
||||||
|
self.namespace_field => "global"
|
||||||
|
);
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let doc = doc!(
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
id_field => a.id.clone(),
|
|
||||||
title_field => a.title.clone(),
|
|
||||||
body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
|
||||||
type_field => "adr",
|
|
||||||
namespace_field => "global"
|
|
||||||
);
|
|
||||||
let writer = writer.lock().unwrap();
|
|
||||||
writer.add_document(doc)?;
|
writer.add_document(doc)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn index_batch(
|
pub fn add_entity_sync(&self, e: &Entity) {
|
||||||
&self,
|
if let Ok(writer) = self.writer.lock() {
|
||||||
entities: Vec<Entity>,
|
let _ = writer.add_document(doc!(
|
||||||
tasks: Vec<Task>,
|
self.id_field => e.name.as_str(),
|
||||||
snippets: Vec<Snippet>,
|
self.title_field => e.name.as_str(),
|
||||||
adrs: Vec<Adr>,
|
self.body_field => e.observations.join(" "),
|
||||||
) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
self.type_field => "entity",
|
||||||
let writer = Arc::clone(&self.writer);
|
self.namespace_field => e.namespace.as_str()
|
||||||
let id_field = self.id_field;
|
));
|
||||||
let title_field = self.title_field;
|
}
|
||||||
let body_field = self.body_field;
|
|
||||||
let type_field = self.type_field;
|
|
||||||
let namespace_field = self.namespace_field;
|
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
|
||||||
let writer = writer.lock().unwrap();
|
|
||||||
|
|
||||||
for e in entities {
|
|
||||||
writer.add_document(doc!(
|
|
||||||
id_field => e.name.clone(),
|
|
||||||
title_field => e.name.clone(),
|
|
||||||
body_field => e.observations.join(" "),
|
|
||||||
type_field => "entity",
|
|
||||||
namespace_field => e.namespace.clone()
|
|
||||||
))?;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
for t in tasks {
|
pub fn add_task_sync(&self, t: &Task) {
|
||||||
writer.add_document(doc!(
|
if let Ok(writer) = self.writer.lock() {
|
||||||
id_field => t.id.clone(),
|
let _ = writer.add_document(doc!(
|
||||||
title_field => t.title.clone(),
|
self.id_field => t.id.as_str(),
|
||||||
body_field => t.description.clone(),
|
self.title_field => t.title.as_str(),
|
||||||
type_field => "task",
|
self.body_field => t.description.as_str(),
|
||||||
namespace_field => "global"
|
self.type_field => "task",
|
||||||
))?;
|
self.namespace_field => "global"
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for s in snippets {
|
pub fn add_snippet_sync(&self, s: &Snippet) {
|
||||||
writer.add_document(doc!(
|
if let Ok(writer) = self.writer.lock() {
|
||||||
id_field => s.name.clone(),
|
let _ = writer.add_document(doc!(
|
||||||
title_field => s.name.clone(),
|
self.id_field => s.name.as_str(),
|
||||||
body_field => format!("{} {}", s.language, s.description),
|
self.title_field => s.name.as_str(),
|
||||||
type_field => "snippet",
|
self.body_field => format!("{} {}", s.language, s.description),
|
||||||
namespace_field => "global"
|
self.type_field => "snippet",
|
||||||
))?;
|
self.namespace_field => "global"
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for a in adrs {
|
pub fn add_adr_sync(&self, a: &Adr) {
|
||||||
writer.add_document(doc!(
|
if let Ok(writer) = self.writer.lock() {
|
||||||
id_field => a.id.clone(),
|
let _ = writer.add_document(doc!(
|
||||||
title_field => a.title.clone(),
|
self.id_field => a.id.as_str(),
|
||||||
body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
self.title_field => a.title.as_str(),
|
||||||
type_field => "adr",
|
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
||||||
namespace_field => "global"
|
self.type_field => "adr",
|
||||||
))?;
|
self.namespace_field => "global"
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+48
-11
@@ -3,7 +3,7 @@ use crate::search::MemoryIndex;
|
|||||||
use crate::store::Store;
|
use crate::store::Store;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use std::sync::RwLock;
|
use std::sync::{Arc, RwLock};
|
||||||
|
|
||||||
pub struct MemoryState {
|
pub struct MemoryState {
|
||||||
pub base_dir: PathBuf,
|
pub base_dir: PathBuf,
|
||||||
@@ -27,6 +27,7 @@ pub struct MemoryState {
|
|||||||
pub tech_debts: Store<Vec<TechDebt>>,
|
pub tech_debts: Store<Vec<TechDebt>>,
|
||||||
pub gates: Store<Vec<GateRecord>>,
|
pub gates: Store<Vec<GateRecord>>,
|
||||||
pub context_workspaces: Store<Vec<ContextWorkspace>>,
|
pub context_workspaces: Store<Vec<ContextWorkspace>>,
|
||||||
|
pub recent_activities: Store<std::collections::VecDeque<serde_json::Value>>,
|
||||||
pub activity_tx: tokio::sync::broadcast::Sender<String>,
|
pub activity_tx: tokio::sync::broadcast::Sender<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -37,16 +38,33 @@ impl MemoryState {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn broadcast_activity(&self, message: &str) {
|
pub fn broadcast_activity(&self, message: &str) {
|
||||||
|
let time = std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_millis() as u64;
|
||||||
|
|
||||||
|
let item = serde_json::json!({
|
||||||
|
"time": time,
|
||||||
|
"message": message
|
||||||
|
});
|
||||||
|
|
||||||
|
self.recent_activities.modify(|activities| {
|
||||||
|
activities.push_back(item.clone());
|
||||||
|
if activities.len() > 100 {
|
||||||
|
activities.pop_front();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
let payload = serde_json::json!({
|
let payload = serde_json::json!({
|
||||||
"type": "activity",
|
"type": "activity",
|
||||||
"data": message
|
"data": item
|
||||||
})
|
})
|
||||||
.to_string();
|
.to_string();
|
||||||
let _ = self.activity_tx.send(payload);
|
let _ = self.activity_tx.send(payload);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_full_graph(&self) -> KnowledgeGraph {
|
pub fn get_full_graph(&self) -> KnowledgeGraph {
|
||||||
self.graph.read()
|
self.graph.read_with(|g| g.clone())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn read_graph<F, R>(&self, f: F) -> R
|
pub fn read_graph<F, R>(&self, f: F) -> R
|
||||||
@@ -60,16 +78,35 @@ impl MemoryState {
|
|||||||
self.graph.modify(update_fn);
|
self.graph.modify(update_fn);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn rebuild_index(&self) {
|
pub async fn rebuild_index(self: &Arc<Self>) {
|
||||||
if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
|
if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
|
||||||
|
let state = Arc::clone(self);
|
||||||
|
let idx = new_idx.clone();
|
||||||
|
|
||||||
let entities = self.graph.read().entities.into_values().collect();
|
tokio::task::spawn_blocking(move || {
|
||||||
let tasks = self.tasks.read();
|
state.graph.read_with(|g| {
|
||||||
let snippets = self.snippets.read();
|
for e in g.entities.values() {
|
||||||
let adrs = self.adrs.read();
|
idx.add_entity_sync(e);
|
||||||
|
}
|
||||||
let handle = new_idx.index_batch(entities, tasks, snippets, adrs);
|
});
|
||||||
let _ = handle.await;
|
state.tasks.read_with(|tasks| {
|
||||||
|
for t in tasks {
|
||||||
|
idx.add_task_sync(t);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
state.snippets.read_with(|snippets| {
|
||||||
|
for s in snippets {
|
||||||
|
idx.add_snippet_sync(s);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
state.adrs.read_with(|adrs| {
|
||||||
|
for a in adrs {
|
||||||
|
idx.add_adr_sync(a);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}).await.unwrap_or_else(|e| {
|
||||||
|
tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
|
||||||
|
});
|
||||||
|
|
||||||
let _ = new_idx.commit().await;
|
let _ = new_idx.commit().await;
|
||||||
if let Ok(mut w) = self.search_index.write() {
|
if let Ok(mut w) = self.search_index.write() {
|
||||||
|
|||||||
+22
-26
@@ -22,58 +22,54 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
|
|||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
while rx.recv().await.is_some() {
|
while rx.recv().await.is_some() {
|
||||||
// Drain any other pending notifications so we batch writes
|
// Drain any other pending notifications so we batch writes
|
||||||
while let Ok(_) = rx.try_recv() {}
|
while rx.try_recv().is_ok() {}
|
||||||
|
|
||||||
let db_inner = db_clone.clone();
|
let db_inner = db_clone.clone();
|
||||||
let key_inner = key_clone.clone();
|
let key_inner = key_clone.clone();
|
||||||
let json_data = {
|
let json_data = {
|
||||||
let lock = cache_clone.read().unwrap();
|
let lock = cache_clone.read().unwrap_or_else(|e| e.into_inner());
|
||||||
serde_json::to_vec(&*lock).unwrap()
|
serde_json::to_vec(&*lock).unwrap_or_default()
|
||||||
};
|
};
|
||||||
|
|
||||||
let _ = tokio::task::spawn_blocking(move || {
|
let _ = tokio::task::spawn_blocking(move || {
|
||||||
let write_txn = db_inner.begin_write().unwrap();
|
if let Ok(write_txn) = db_inner.begin_write() {
|
||||||
{
|
if let Ok(mut table) = write_txn.open_table(STORE_TABLE) {
|
||||||
let mut table = write_txn.open_table(STORE_TABLE).unwrap();
|
let _ = table.insert(key_inner.as_str(), json_data.as_slice());
|
||||||
table.insert(key_inner.as_str(), json_data.as_slice()).unwrap();
|
|
||||||
}
|
}
|
||||||
write_txn.commit().unwrap();
|
let _ = write_txn.commit();
|
||||||
}).await;
|
}
|
||||||
|
})
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
Self {
|
Self { cache, tx }
|
||||||
cache,
|
|
||||||
tx,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn load_from_db(key: &str, db: &Database) -> T {
|
fn load_from_db(key: &str, db: &Database) -> T {
|
||||||
let read_txn = db.begin_read().unwrap();
|
let Ok(read_txn) = db.begin_read() else {
|
||||||
|
return T::default();
|
||||||
|
};
|
||||||
if let Ok(table) = read_txn.open_table(STORE_TABLE)
|
if let Ok(table) = read_txn.open_table(STORE_TABLE)
|
||||||
&& let Ok(Some(value)) = table.get(key)
|
&& let Ok(Some(value)) = table.get(key)
|
||||||
&& let Ok(parsed) = serde_json::from_slice::<T>(value.value()) {
|
&& let Ok(parsed) = serde_json::from_slice::<T>(value.value())
|
||||||
|
{
|
||||||
return parsed;
|
return parsed;
|
||||||
}
|
}
|
||||||
T::default()
|
T::default()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn read(&self) -> T {
|
|
||||||
let lock = self.cache.read().unwrap();
|
|
||||||
lock.clone()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn read_with<F, R>(&self, f: F) -> R
|
pub fn read_with<F, R>(&self, f: F) -> R
|
||||||
where
|
where
|
||||||
F: FnOnce(&T) -> R,
|
F: FnOnce(&T) -> R,
|
||||||
{
|
{
|
||||||
let lock = self.cache.read().unwrap();
|
let lock = self.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||||
f(&lock)
|
f(&lock)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn modify<F: FnOnce(&mut T)>(&self, f: F) {
|
pub fn modify<F: FnOnce(&mut T)>(&self, f: F) {
|
||||||
{
|
{
|
||||||
let mut lock = self.cache.write().unwrap();
|
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
|
||||||
f(&mut lock);
|
f(&mut lock);
|
||||||
}
|
}
|
||||||
let _ = self.tx.try_send(());
|
let _ = self.tx.try_send(());
|
||||||
@@ -105,7 +101,7 @@ mod tests {
|
|||||||
let db = Arc::new(db);
|
let db = Arc::new(db);
|
||||||
let store = Store::<TestData>::new("test_key", db.clone());
|
let store = Store::<TestData>::new("test_key", db.clone());
|
||||||
|
|
||||||
assert_eq!(store.read(), TestData::default());
|
assert_eq!(store.read_with(|s| s.clone()), TestData::default());
|
||||||
|
|
||||||
store.modify(|data| {
|
store.modify(|data| {
|
||||||
data.name = "Hello".to_string();
|
data.name = "Hello".to_string();
|
||||||
@@ -116,7 +112,7 @@ mod tests {
|
|||||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
store.read(),
|
store.read_with(|s| s.clone()),
|
||||||
TestData {
|
TestData {
|
||||||
name: "Hello".to_string(),
|
name: "Hello".to_string(),
|
||||||
value: 42
|
value: 42
|
||||||
@@ -126,7 +122,7 @@ mod tests {
|
|||||||
// Load again to verify persistence
|
// Load again to verify persistence
|
||||||
let store2 = Store::<TestData>::new("test_key", db.clone());
|
let store2 = Store::<TestData>::new("test_key", db.clone());
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
store2.read(),
|
store2.read_with(|s| s.clone()),
|
||||||
TestData {
|
TestData {
|
||||||
name: "Hello".to_string(),
|
name: "Hello".to_string(),
|
||||||
value: 42
|
value: 42
|
||||||
@@ -165,6 +161,6 @@ mod tests {
|
|||||||
// Wait for all blocking writes to flush
|
// Wait for all blocking writes to flush
|
||||||
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
||||||
|
|
||||||
assert_eq!(store.read().value, 50);
|
assert_eq!(store.read_with(|s| s.value), 50);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -314,6 +314,7 @@ pub struct AddSessionSummaryTool {
|
|||||||
|
|
||||||
/// Get a timeline of major project events.
|
/// Get a timeline of major project events.
|
||||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||||
|
#[allow(dead_code)]
|
||||||
pub struct GetProjectTimelineTool {
|
pub struct GetProjectTimelineTool {
|
||||||
/// Optional namespace to restrict the timeline to.
|
/// Optional namespace to restrict the timeline to.
|
||||||
pub namespace: Option<String>,
|
pub namespace: Option<String>,
|
||||||
|
|||||||
+16
-12
@@ -2,22 +2,27 @@ use std::collections::HashSet;
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_eager_tools_parity() {
|
fn test_eager_tools_parity() {
|
||||||
// 1. Read handlers.rs to get memory tools
|
// 1. Read handlers_v2/*.rs to get memory tools
|
||||||
let memory_source =
|
|
||||||
std::fs::read_to_string("src/handlers.rs").expect("Failed to read handlers.rs");
|
|
||||||
let mut memory_tools = HashSet::new();
|
let mut memory_tools = HashSet::new();
|
||||||
|
let entries = std::fs::read_dir("src/handlers_v2").expect("Failed to read handlers_v2 dir");
|
||||||
|
for entry in entries {
|
||||||
|
let entry = entry.unwrap();
|
||||||
|
if entry.path().extension().unwrap_or_default() == "rs" {
|
||||||
|
let memory_source = std::fs::read_to_string(entry.path()).unwrap();
|
||||||
let parts: Vec<&str> = memory_source.split("crate::mcp::tool_def").collect();
|
let parts: Vec<&str> = memory_source.split("crate::mcp::tool_def").collect();
|
||||||
for part in parts.iter().skip(1) {
|
for part in parts.iter().skip(1) {
|
||||||
if let Some(start) = part.find("\"") {
|
if let Some(start) = part.find('"') {
|
||||||
let rest = &part[start + 1..];
|
let rest = &part[start + 1..];
|
||||||
if let Some(end) = rest.find("\"") {
|
if let Some(end) = rest.find('"') {
|
||||||
memory_tools.insert(rest[..end].to_string());
|
memory_tools.insert(rest[..end].to_string());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
assert!(
|
assert!(
|
||||||
!memory_tools.is_empty(),
|
!memory_tools.is_empty(),
|
||||||
"Could not find memory tools in handlers.rs"
|
"Could not find memory tools in handlers_v2 directory"
|
||||||
);
|
);
|
||||||
|
|
||||||
// 2. Read nvim-core/src/lib.rs to get nvim tools
|
// 2. Read nvim-core/src/lib.rs to get nvim tools
|
||||||
@@ -26,9 +31,10 @@ fn test_eager_tools_parity() {
|
|||||||
let mut nvim_tools = HashSet::new();
|
let mut nvim_tools = HashSet::new();
|
||||||
for line in nvim_source.lines() {
|
for line in nvim_source.lines() {
|
||||||
if line.contains("\"name\": \"nvim_")
|
if line.contains("\"name\": \"nvim_")
|
||||||
&& let Some(start) = line.find("\"name\": \"") {
|
&& let Some(start) = line.find("\"name\": \"")
|
||||||
|
{
|
||||||
let rest = &line[start + 9..];
|
let rest = &line[start + 9..];
|
||||||
if let Some(end) = rest.find("\"") {
|
if let Some(end) = rest.find('"') {
|
||||||
nvim_tools.insert(rest[..end].to_string());
|
nvim_tools.insert(rest[..end].to_string());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -51,8 +57,7 @@ fn test_eager_tools_parity() {
|
|||||||
let name = tool.as_str().unwrap();
|
let name = tool.as_str().unwrap();
|
||||||
assert!(
|
assert!(
|
||||||
memory_tools.contains(name),
|
memory_tools.contains(name),
|
||||||
"Windows config Memory tool '{}' not implemented in handlers.rs!",
|
"Windows config Memory tool '{name}' not implemented in handlers.rs!"
|
||||||
name
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -62,8 +67,7 @@ fn test_eager_tools_parity() {
|
|||||||
let name = tool.as_str().unwrap();
|
let name = tool.as_str().unwrap();
|
||||||
assert!(
|
assert!(
|
||||||
nvim_tools.contains(name),
|
nvim_tools.contains(name),
|
||||||
"Windows config Nvim tool '{}' not implemented in nvim-core!",
|
"Windows config Nvim tool '{name}' not implemented in nvim-core!"
|
||||||
name
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+2
-6
@@ -15,10 +15,6 @@ fn main() {
|
|||||||
.and_then(|out| String::from_utf8(out.stdout).ok())
|
.and_then(|out| String::from_utf8(out.stdout).ok())
|
||||||
.unwrap_or_else(|| "unknown".to_string());
|
.unwrap_or_else(|| "unknown".to_string());
|
||||||
|
|
||||||
let version = format!(
|
let version = format!("{} ({})", git_date.trim(), git_hash.trim());
|
||||||
"{} ({})",
|
println!("cargo:rustc-env=APP_VERSION={version}");
|
||||||
git_date.trim(),
|
|
||||||
git_hash.trim()
|
|
||||||
);
|
|
||||||
println!("cargo:rustc-env=APP_VERSION={}", version);
|
|
||||||
}
|
}
|
||||||
@@ -17,7 +17,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
.danger_accept_invalid_certs(true)
|
.danger_accept_invalid_certs(true)
|
||||||
.build()?;
|
.build()?;
|
||||||
|
|
||||||
let sse_url = format!("{}/sse", target);
|
let sse_url = format!("{target}/sse");
|
||||||
|
|
||||||
tracing::info!("Connecting to SSE: {}", sse_url);
|
tracing::info!("Connecting to SSE: {}", sse_url);
|
||||||
|
|
||||||
@@ -59,9 +59,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
tracing::info!("Received POST endpoint: {}", data_content);
|
tracing::info!("Received POST endpoint: {}", data_content);
|
||||||
post_endpoint = Some(data_content);
|
post_endpoint = Some(data_content);
|
||||||
break;
|
break;
|
||||||
} else {
|
|
||||||
tracing::info!("Received early SSE data: {}", text);
|
|
||||||
}
|
}
|
||||||
|
tracing::info!("Received early SSE data: {}", text);
|
||||||
}
|
}
|
||||||
if post_endpoint.is_some() {
|
if post_endpoint.is_some() {
|
||||||
break;
|
break;
|
||||||
@@ -69,7 +68,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let post_endpoint = post_endpoint.ok_or("Did not receive endpoint from SSE stream")?;
|
let post_endpoint = post_endpoint.ok_or("Did not receive endpoint from SSE stream")?;
|
||||||
let post_url = format!("{}{}", target, post_endpoint);
|
let post_url = format!("{target}{post_endpoint}");
|
||||||
|
|
||||||
let payload = r#"{"jsonrpc":"2.0","id":999,"method":"server/discover","params":{}}"#;
|
let payload = r#"{"jsonrpc":"2.0","id":999,"method":"server/discover","params":{}}"#;
|
||||||
|
|
||||||
@@ -98,7 +97,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
_ = timeout.tick() => {
|
_ = timeout.tick() => {
|
||||||
tracing::warn!("Timed out waiting for SSE response.");
|
tracing::warn!("Timed out waiting for SSE response.");
|
||||||
}
|
}
|
||||||
_ = async {
|
() = async {
|
||||||
while let Some(chunk) = stream.next().await {
|
while let Some(chunk) = stream.next().await {
|
||||||
if let Ok(bytes) = chunk {
|
if let Ok(bytes) = chunk {
|
||||||
tracing::info!("Received SSE Chunk: {}", String::from_utf8_lossy(&bytes));
|
tracing::info!("Received SSE Chunk: {}", String::from_utf8_lossy(&bytes));
|
||||||
|
|||||||
+23
-11
@@ -1,29 +1,41 @@
|
|||||||
use std::sync::LazyLock;
|
|
||||||
use regex::Regex;
|
use regex::Regex;
|
||||||
|
use std::sync::LazyLock;
|
||||||
|
|
||||||
static ID_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""id"\s*:\s*([^,}]+)"#).unwrap());
|
static ID_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""id"\s*:\s*([^,}]+)"#).unwrap());
|
||||||
static METHOD_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""method"\s*:\s*"([^"]+)""#).unwrap());
|
static METHOD_REGEX: LazyLock<Regex> =
|
||||||
static TOOL_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""name"\s*:\s*"([^"]+)""#).unwrap());
|
LazyLock::new(|| Regex::new(r#""method"\s*:\s*"([^"]+)""#).unwrap());
|
||||||
|
static TOOL_REGEX: LazyLock<Regex> =
|
||||||
|
LazyLock::new(|| Regex::new(r#""name"\s*:\s*"([^"]+)""#).unwrap());
|
||||||
static ERROR_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""error"\s*:\s*\{"#).unwrap());
|
static ERROR_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""error"\s*:\s*\{"#).unwrap());
|
||||||
static IS_ERROR_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""isError"\s*:\s*true"#).unwrap());
|
static IS_ERROR_REGEX: LazyLock<Regex> =
|
||||||
|
LazyLock::new(|| Regex::new(r#""isError"\s*:\s*true"#).unwrap());
|
||||||
|
|
||||||
pub fn extract_log_prefix(json_str: &str, is_response: bool) -> String {
|
pub fn extract_log_prefix(json_str: &str, is_response: bool) -> String {
|
||||||
let id = ID_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("null");
|
let id = ID_REGEX
|
||||||
|
.captures(json_str)
|
||||||
|
.and_then(|c| c.get(1))
|
||||||
|
.map_or("null", |m| m.as_str());
|
||||||
|
|
||||||
if is_response {
|
if is_response {
|
||||||
let is_error = ERROR_REGEX.is_match(json_str) || IS_ERROR_REGEX.is_match(json_str);
|
let is_error = ERROR_REGEX.is_match(json_str) || IS_ERROR_REGEX.is_match(json_str);
|
||||||
return format!("Response id={} [Error: {}]", id, is_error);
|
return format!("Response id={id} [Error: {is_error}]");
|
||||||
}
|
}
|
||||||
|
|
||||||
let method = METHOD_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("");
|
let method = METHOD_REGEX
|
||||||
|
.captures(json_str)
|
||||||
|
.and_then(|c| c.get(1))
|
||||||
|
.map_or("", |m| m.as_str());
|
||||||
|
|
||||||
if method == "tools/call" {
|
if method == "tools/call" {
|
||||||
let tool = TOOL_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("unknown");
|
let tool = TOOL_REGEX
|
||||||
format!("ToolCall[{}] id={}", tool, id)
|
.captures(json_str)
|
||||||
|
.and_then(|c| c.get(1))
|
||||||
|
.map_or("unknown", |m| m.as_str());
|
||||||
|
format!("ToolCall[{tool}] id={id}")
|
||||||
} else if !method.is_empty() {
|
} else if !method.is_empty() {
|
||||||
format!("Request[{}] id={}", method, id)
|
format!("Request[{method}] id={id}")
|
||||||
} else if id != "null" {
|
} else if id != "null" {
|
||||||
format!("Response id={}", id)
|
format!("Response id={id}")
|
||||||
} else {
|
} else {
|
||||||
"Message".to_string()
|
"Message".to_string()
|
||||||
}
|
}
|
||||||
|
|||||||
+53
-16
@@ -10,8 +10,6 @@ struct Cli {
|
|||||||
target: String,
|
target: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
mod logger;
|
mod logger;
|
||||||
|
|
||||||
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
||||||
@@ -19,7 +17,7 @@ fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::Worker
|
|||||||
base_dir.push(".gemini/mcp_memory/logs");
|
base_dir.push(".gemini/mcp_memory/logs");
|
||||||
std::fs::create_dir_all(&base_dir).unwrap_or_default();
|
std::fs::create_dir_all(&base_dir).unwrap_or_default();
|
||||||
|
|
||||||
let file_appender = tracing_appender::rolling::daily(base_dir, format!("{}.log", app_name));
|
let file_appender = tracing_appender::rolling::daily(base_dir, format!("{app_name}.log"));
|
||||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||||
|
|
||||||
let _ = tracing_subscriber::fmt()
|
let _ = tracing_subscriber::fmt()
|
||||||
@@ -47,14 +45,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
let _ = shutdown_tx.send(()).await;
|
let _ = shutdown_tx.send(()).await;
|
||||||
});
|
});
|
||||||
|
|
||||||
let target_url = if cli.target != "http://localhost:3000" {
|
let target_url = if cli.target == "http://localhost:3000" {
|
||||||
cli.target
|
|
||||||
} else {
|
|
||||||
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||||
format!("http://127.0.0.1:{}", port)
|
format!("http://127.0.0.1:{port}")
|
||||||
|
} else {
|
||||||
|
cli.target
|
||||||
};
|
};
|
||||||
let ws_url = target_url.replace("http://", "ws://").replace("https://", "wss://");
|
let ws_url = target_url
|
||||||
let ws_url = format!("{}/ws?client=proxy", ws_url);
|
.replace("http://", "ws://")
|
||||||
|
.replace("https://", "wss://");
|
||||||
|
let ws_url = format!("{ws_url}/ws?client=proxy");
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if shutdown_rx.try_recv().is_ok() {
|
if shutdown_rx.try_recv().is_ok() {
|
||||||
@@ -74,8 +74,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
match tokio_tungstenite::connect_async(request).await {
|
let connect_result = tokio::select! {
|
||||||
Ok((ws_stream, _)) => {
|
_ = shutdown_rx.recv() => {
|
||||||
|
tracing::info!("Shutdown received during connect");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
res = tokio::time::timeout(
|
||||||
|
tokio::time::Duration::from_secs(5),
|
||||||
|
tokio_tungstenite::connect_async(request)
|
||||||
|
) => res,
|
||||||
|
};
|
||||||
|
|
||||||
|
match connect_result {
|
||||||
|
Ok(Ok((ws_stream, _))) => {
|
||||||
tracing::info!("Successfully connected to target server");
|
tracing::info!("Successfully connected to target server");
|
||||||
let (mut write, mut read) = ws_stream.split();
|
let (mut write, mut read) = ws_stream.split();
|
||||||
|
|
||||||
@@ -83,8 +94,21 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
let mut send_task = tokio::spawn(async move {
|
let mut send_task = tokio::spawn(async move {
|
||||||
while let Ok(msg) = rx.recv().await {
|
while let Ok(msg) = rx.recv().await {
|
||||||
let log_prefix = logger::extract_log_prefix(&msg, false);
|
let log_prefix = logger::extract_log_prefix(&msg, false);
|
||||||
tracing::info!(">>> [Stub] Forwarding {} to server (length: {}): {}", log_prefix, msg.len(), if msg.len() > 1000 { format!("{}...", &msg[..1000]) } else { msg.clone() });
|
tracing::info!(
|
||||||
if write.send(tokio_tungstenite::tungstenite::Message::Text(msg)).await.is_err() {
|
">>> [Stub] Forwarding {} to server (length: {}): {}",
|
||||||
|
log_prefix,
|
||||||
|
msg.len(),
|
||||||
|
if msg.len() > 1000 {
|
||||||
|
format!("{}...", &msg[..1000])
|
||||||
|
} else {
|
||||||
|
msg.clone()
|
||||||
|
}
|
||||||
|
);
|
||||||
|
if write
|
||||||
|
.send(tokio_tungstenite::tungstenite::Message::Text(msg))
|
||||||
|
.await
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
tracing::error!("Failed to write to websocket");
|
tracing::error!("Failed to write to websocket");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
@@ -95,7 +119,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
while let Some(Ok(msg)) = read.next().await {
|
while let Some(Ok(msg)) = read.next().await {
|
||||||
if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
|
if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
|
||||||
let log_prefix = logger::extract_log_prefix(&text, true);
|
let log_prefix = logger::extract_log_prefix(&text, true);
|
||||||
tracing::info!("<<< [Stub] Received {} from server (length: {}): {}", log_prefix, text.len(), if text.len() > 1000 { format!("{}...", &text[..1000]) } else { text.clone() });
|
tracing::info!(
|
||||||
|
"<<< [Stub] Received {} from server (length: {}): {}",
|
||||||
|
log_prefix,
|
||||||
|
text.len(),
|
||||||
|
if text.len() > 1000 {
|
||||||
|
format!("{}...", &text[..1000])
|
||||||
|
} else {
|
||||||
|
text.clone()
|
||||||
|
}
|
||||||
|
);
|
||||||
use tokio::io::AsyncWriteExt;
|
use tokio::io::AsyncWriteExt;
|
||||||
let mut stdout = tokio::io::stdout();
|
let mut stdout = tokio::io::stdout();
|
||||||
let _ = stdout.write_all(text.as_bytes()).await;
|
let _ = stdout.write_all(text.as_bytes()).await;
|
||||||
@@ -116,7 +149,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
tracing::error!("Send task exited");
|
tracing::error!("Send task exited");
|
||||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||||
recv_task.abort();
|
recv_task.abort();
|
||||||
break;
|
|
||||||
}
|
}
|
||||||
_ = &mut recv_task => {
|
_ = &mut recv_task => {
|
||||||
tracing::error!("Recv task exited");
|
tracing::error!("Recv task exited");
|
||||||
@@ -125,13 +157,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Ok(Err(e)) => {
|
||||||
tracing::error!("Failed to connect via WSS: {}", e);
|
tracing::error!("Failed to connect via WSS: {}", e);
|
||||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
}
|
}
|
||||||
|
Err(_) => {
|
||||||
|
tracing::error!("Connection attempt timed out");
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
+23
-17
@@ -14,7 +14,7 @@ impl Drop for ChildGuard {
|
|||||||
|
|
||||||
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
|
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
|
||||||
let s = serde_json::to_string(&msg).unwrap();
|
let s = serde_json::to_string(&msg).unwrap();
|
||||||
let payload = format!("{}\n", s);
|
let payload = format!("{s}\n");
|
||||||
stdin.write_all(payload.as_bytes()).unwrap();
|
stdin.write_all(payload.as_bytes()).unwrap();
|
||||||
stdin.flush().unwrap();
|
stdin.flush().unwrap();
|
||||||
}
|
}
|
||||||
@@ -56,12 +56,13 @@ async fn test_full_system_e2e_performance() {
|
|||||||
let nvim_exe = debug_dir.join(format!("{}{}", nvim_name, std::env::consts::EXE_SUFFIX));
|
let nvim_exe = debug_dir.join(format!("{}{}", nvim_name, std::env::consts::EXE_SUFFIX));
|
||||||
let stub_exe = debug_dir.join(format!("mcp-memory-stub{}", std::env::consts::EXE_SUFFIX));
|
let stub_exe = debug_dir.join(format!("mcp-memory-stub{}", std::env::consts::EXE_SUFFIX));
|
||||||
|
|
||||||
assert!(server_exe.exists(), "Server not found at {:?}", server_exe);
|
assert!(server_exe.exists(), "Server not found at {server_exe:?}");
|
||||||
assert!(nvim_exe.exists(), "Nvim not found at {:?}", nvim_exe);
|
assert!(nvim_exe.exists(), "Nvim not found at {nvim_exe:?}");
|
||||||
assert!(stub_exe.exists(), "Stub not found at {:?}", stub_exe);
|
assert!(stub_exe.exists(), "Stub not found at {stub_exe:?}");
|
||||||
|
|
||||||
// 1. Start Server
|
// 1. Start Server
|
||||||
let _server = ChildGuard(Command::new(&server_exe)
|
let _server = ChildGuard(
|
||||||
|
Command::new(&server_exe)
|
||||||
.arg("--daemon")
|
.arg("--daemon")
|
||||||
.env("MCP_PORT", test_port)
|
.env("MCP_PORT", test_port)
|
||||||
.env("RUST_LOG", "debug")
|
.env("RUST_LOG", "debug")
|
||||||
@@ -71,7 +72,8 @@ async fn test_full_system_e2e_performance() {
|
|||||||
.stdout(Stdio::inherit())
|
.stdout(Stdio::inherit())
|
||||||
.stderr(Stdio::inherit())
|
.stderr(Stdio::inherit())
|
||||||
.spawn()
|
.spawn()
|
||||||
.expect("Failed to start server"));
|
.expect("Failed to start server"),
|
||||||
|
);
|
||||||
|
|
||||||
// Give server time to generate TLS cert and start
|
// Give server time to generate TLS cert and start
|
||||||
let client = reqwest::Client::builder()
|
let client = reqwest::Client::builder()
|
||||||
@@ -81,7 +83,7 @@ async fn test_full_system_e2e_performance() {
|
|||||||
let mut started = false;
|
let mut started = false;
|
||||||
for _ in 0..30 {
|
for _ in 0..30 {
|
||||||
if let Ok(resp) = client
|
if let Ok(resp) = client
|
||||||
.get(format!("http://127.0.0.1:{}/health", test_port))
|
.get(format!("http://127.0.0.1:{test_port}/health"))
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
&& resp.status().is_success()
|
&& resp.status().is_success()
|
||||||
@@ -94,9 +96,10 @@ async fn test_full_system_e2e_performance() {
|
|||||||
assert!(started, "Server failed to start in time");
|
assert!(started, "Server failed to start in time");
|
||||||
|
|
||||||
// 2. Start Stub
|
// 2. Start Stub
|
||||||
let mut stub = ChildGuard(Command::new(&stub_exe)
|
let mut stub = ChildGuard(
|
||||||
|
Command::new(&stub_exe)
|
||||||
.arg("--target")
|
.arg("--target")
|
||||||
.arg(format!("http://127.0.0.1:{}", test_port))
|
.arg(format!("http://127.0.0.1:{test_port}"))
|
||||||
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
|
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
|
||||||
.env("MCP_AUTH_TOKEN", test_auth_token)
|
.env("MCP_AUTH_TOKEN", test_auth_token)
|
||||||
.env("RUST_LOG", "debug")
|
.env("RUST_LOG", "debug")
|
||||||
@@ -104,18 +107,21 @@ async fn test_full_system_e2e_performance() {
|
|||||||
.stdout(Stdio::piped())
|
.stdout(Stdio::piped())
|
||||||
.stderr(Stdio::inherit())
|
.stderr(Stdio::inherit())
|
||||||
.spawn()
|
.spawn()
|
||||||
.expect("Failed to start stub"));
|
.expect("Failed to start stub"),
|
||||||
|
);
|
||||||
|
|
||||||
let mut stub_stdin = stub.0.stdin.take().unwrap();
|
let mut stub_stdin = stub.0.stdin.take().unwrap();
|
||||||
let mut stub_stdout = BufReader::new(stub.0.stdout.take().unwrap());
|
let mut stub_stdout = BufReader::new(stub.0.stdout.take().unwrap());
|
||||||
|
|
||||||
// 3. Start Nvim Bridge
|
// 3. Start Nvim Bridge
|
||||||
let mut nvim = ChildGuard(Command::new(&nvim_exe)
|
let mut nvim = ChildGuard(
|
||||||
|
Command::new(&nvim_exe)
|
||||||
.stdin(Stdio::piped())
|
.stdin(Stdio::piped())
|
||||||
.stdout(Stdio::piped())
|
.stdout(Stdio::piped())
|
||||||
.stderr(Stdio::inherit())
|
.stderr(Stdio::inherit())
|
||||||
.spawn()
|
.spawn()
|
||||||
.expect("Failed to start nvim bridge"));
|
.expect("Failed to start nvim bridge"),
|
||||||
|
);
|
||||||
|
|
||||||
let mut nvim_stdin = nvim.0.stdin.take().unwrap();
|
let mut nvim_stdin = nvim.0.stdin.take().unwrap();
|
||||||
let mut nvim_stdout = BufReader::new(nvim.0.stdout.take().unwrap());
|
let mut nvim_stdout = BufReader::new(nvim.0.stdout.take().unwrap());
|
||||||
@@ -138,7 +144,7 @@ async fn test_full_system_e2e_performance() {
|
|||||||
send_message(&mut stub_stdin, tools_req);
|
send_message(&mut stub_stdin, tools_req);
|
||||||
} else {
|
} else {
|
||||||
let s = serde_json::to_string(&tools_req).unwrap();
|
let s = serde_json::to_string(&tools_req).unwrap();
|
||||||
stub_stdin.write_all(format!("{}\n", s).as_bytes()).unwrap();
|
stub_stdin.write_all(format!("{s}\n").as_bytes()).unwrap();
|
||||||
stub_stdin.flush().unwrap();
|
stub_stdin.flush().unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -150,7 +156,7 @@ async fn test_full_system_e2e_performance() {
|
|||||||
assert_eq!(resp["id"], i);
|
assert_eq!(resp["id"], i);
|
||||||
}
|
}
|
||||||
let stub_duration = start_time.elapsed();
|
let stub_duration = start_time.elapsed();
|
||||||
println!("Stub 100 requests: {:?}", stub_duration);
|
println!("Stub 100 requests: {stub_duration:?}");
|
||||||
|
|
||||||
// Send 100 concurrent-like sequential rapid requests to Win-Nvim
|
// Send 100 concurrent-like sequential rapid requests to Win-Nvim
|
||||||
println!("Starting 100 requests to nvim...");
|
println!("Starting 100 requests to nvim...");
|
||||||
@@ -167,7 +173,7 @@ async fn test_full_system_e2e_performance() {
|
|||||||
send_message(&mut nvim_stdin, tools_req);
|
send_message(&mut nvim_stdin, tools_req);
|
||||||
} else {
|
} else {
|
||||||
let s = serde_json::to_string(&tools_req).unwrap();
|
let s = serde_json::to_string(&tools_req).unwrap();
|
||||||
nvim_stdin.write_all(format!("{}\n", s).as_bytes()).unwrap();
|
nvim_stdin.write_all(format!("{s}\n").as_bytes()).unwrap();
|
||||||
nvim_stdin.flush().unwrap();
|
nvim_stdin.flush().unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -181,8 +187,8 @@ async fn test_full_system_e2e_performance() {
|
|||||||
}
|
}
|
||||||
let nvim_duration = start_time_nvim.elapsed();
|
let nvim_duration = start_time_nvim.elapsed();
|
||||||
|
|
||||||
println!("Stub 100 requests: {:?}", stub_duration);
|
println!("Stub 100 requests: {stub_duration:?}");
|
||||||
println!("Win-Nvim 100 requests: {:?}", nvim_duration);
|
println!("Win-Nvim 100 requests: {nvim_duration:?}");
|
||||||
|
|
||||||
let _ = std::fs::remove_dir_all(temp_dir);
|
let _ = std::fs::remove_dir_all(temp_dir);
|
||||||
}
|
}
|
||||||
@@ -31,8 +31,7 @@ async fn test_stub_connection_refused() {
|
|||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
res.is_ok(),
|
res.is_ok(),
|
||||||
"Stub hung on connection refused! Took {:?}",
|
"Stub hung on connection refused! Took {elapsed:?}"
|
||||||
elapsed
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -65,7 +64,7 @@ async fn test_stub_handles_eof_cleanly() {
|
|||||||
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
|
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
|
||||||
let elapsed = start.elapsed();
|
let elapsed = start.elapsed();
|
||||||
|
|
||||||
assert!(res.is_ok(), "Stub hung after EOF! Took {:?}", elapsed);
|
assert!(res.is_ok(), "Stub hung after EOF! Took {elapsed:?}");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -102,7 +101,6 @@ async fn test_stub_sse_fallback_failure() {
|
|||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
res.is_ok(),
|
res.is_ok(),
|
||||||
"Stub hung on fallback failure! Took {:?}",
|
"Stub hung on fallback failure! Took {elapsed:?}"
|
||||||
elapsed
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
local buf = vim.fn.bufnr('server/src/search.rs')
|
||||||
|
if buf == -1 then
|
||||||
|
vim.cmd('e server/src/search.rs')
|
||||||
|
buf = vim.api.nvim_get_current_buf()
|
||||||
|
end
|
||||||
|
|
||||||
|
vim.api.nvim_buf_set_lines(buf, 216, 222, false, {
|
||||||
|
' &self,',
|
||||||
|
' entities: &[Entity],',
|
||||||
|
' tasks: &[Task],',
|
||||||
|
' snippets: &[Snippet],',
|
||||||
|
' adrs: &[Adr],'
|
||||||
|
})
|
||||||
|
|
||||||
|
-- We need to change the loop variables inside the task from values to clones if they are passed as slices
|
||||||
|
-- Actually we can just clone the slice data before moving it into spawn_blocking
|
||||||
|
vim.api.nvim_buf_set_lines(buf, 222, 223, false, {
|
||||||
|
' ) -> tokio::task::JoinHandle<tantivy::Result<()>> {',
|
||||||
|
' let entities = entities.to_vec();',
|
||||||
|
' let tasks = tasks.to_vec();',
|
||||||
|
' let snippets = snippets.to_vec();',
|
||||||
|
' let adrs = adrs.to_vec();'
|
||||||
|
})
|
||||||
|
|
||||||
|
vim.cmd('write')
|
||||||
|
require('notify')('Updated search index_batch signature', 'info', { title = 'Antigravity' })
|
||||||
|
return 'ok'
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
local buf = vim.fn.bufnr('server/src/search.rs')
|
||||||
|
if buf == -1 then
|
||||||
|
vim.cmd('e server/src/search.rs')
|
||||||
|
buf = vim.api.nvim_get_current_buf()
|
||||||
|
end
|
||||||
|
|
||||||
|
vim.api.nvim_buf_set_lines(buf, 216, 222, false, {
|
||||||
|
' &self,',
|
||||||
|
' entities: &[Entity],',
|
||||||
|
' tasks: &[Task],',
|
||||||
|
' snippets: &[Snippet],',
|
||||||
|
' adrs: &[Adr],'
|
||||||
|
})
|
||||||
|
|
||||||
|
-- We need to change the loop variables inside the task from values to clones if they are passed as slices
|
||||||
|
-- Actually we can just clone the slice data before moving it into spawn_blocking
|
||||||
|
vim.api.nvim_buf_set_lines(buf, 222, 223, false, {
|
||||||
|
' ) -> tokio::task::JoinHandle<tantivy::Result<()>> {',
|
||||||
|
' let entities = entities.to_vec();',
|
||||||
|
' let tasks = tasks.to_vec();',
|
||||||
|
' let snippets = snippets.to_vec();',
|
||||||
|
' let adrs = adrs.to_vec();'
|
||||||
|
})
|
||||||
|
|
||||||
|
vim.cmd('write')
|
||||||
|
require('notify')('Updated search index_batch signature', 'info', { title = 'Antigravity' })
|
||||||
|
return 'ok'
|
||||||
|
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
local code = vim.fn.readfile('temp.lua')
|
||||||
|
return table.concat(code, '\n')
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
local buf = vim.fn.bufnr('server/src/handlers_v2/graph.rs')
|
||||||
|
if buf == -1 then
|
||||||
|
vim.cmd('e server/src/handlers_v2/graph.rs')
|
||||||
|
buf = vim.api.nvim_get_current_buf()
|
||||||
|
end
|
||||||
|
return vim.api.nvim_buf_get_name(buf)
|
||||||
@@ -0,0 +1,25 @@
|
|||||||
|
--- server/src/main.rs
|
||||||
|
+++ server/src/main.rs
|
||||||
|
@@ -225,6 +225,20 @@
|
||||||
|
next_id: AtomicUsize::new(1),
|
||||||
|
});
|
||||||
|
|
||||||
|
+ let app_state_clone = Arc::clone(&app_state);
|
||||||
|
+ let mut rx = state.activity_tx.subscribe();
|
||||||
|
+ tokio::spawn(async move {
|
||||||
|
+ while let Ok(msg) = rx.recv().await {
|
||||||
|
+ let senders: Vec<_> = app_state_clone
|
||||||
|
+ .clients
|
||||||
|
+ .read()
|
||||||
|
+ .unwrap_or_else(|e| e.into_inner())
|
||||||
|
+ .values()
|
||||||
|
+ .cloned()
|
||||||
|
+ .collect();
|
||||||
|
+ for client_tx in senders {
|
||||||
|
+ let _ = client_tx.try_send(msg.clone());
|
||||||
|
+ }
|
||||||
|
+ }
|
||||||
|
+ });
|
||||||
|
+
|
||||||
|
let app = Router::new()
|
||||||
|
.route(
|
||||||
+2
-6
@@ -15,10 +15,6 @@ fn main() {
|
|||||||
.and_then(|out| String::from_utf8(out.stdout).ok())
|
.and_then(|out| String::from_utf8(out.stdout).ok())
|
||||||
.unwrap_or_else(|| "unknown".to_string());
|
.unwrap_or_else(|| "unknown".to_string());
|
||||||
|
|
||||||
let version = format!(
|
let version = format!("{} ({})", git_date.trim(), git_hash.trim());
|
||||||
"{} ({})",
|
println!("cargo:rustc-env=APP_VERSION={version}");
|
||||||
git_date.trim(),
|
|
||||||
git_hash.trim()
|
|
||||||
);
|
|
||||||
println!("cargo:rustc-env=APP_VERSION={}", version);
|
|
||||||
}
|
}
|
||||||
@@ -1,10 +1,10 @@
|
|||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::io::{BufRead, BufReader, Read, Write};
|
use std::io::{BufRead, BufReader, Write};
|
||||||
use std::process::{Command, Stdio};
|
use std::process::{Command, Stdio};
|
||||||
|
|
||||||
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
|
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
|
||||||
let s = serde_json::to_string(&msg).unwrap();
|
let s = serde_json::to_string(&msg).unwrap();
|
||||||
stdin.write_all(format!("{}\n", s).as_bytes()).unwrap();
|
stdin.write_all(format!("{s}\n").as_bytes()).unwrap();
|
||||||
stdin.flush().unwrap();
|
stdin.flush().unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user