254 lines
9.6 KiB
Rust
254 lines
9.6 KiB
Rust
use crate::router::McpTool;
|
|
use crate::state::MemoryState;
|
|
use crate::tools::GetActiveWorktreeContextTool;
|
|
use async_trait::async_trait;
|
|
use serde_json::{Value, json};
|
|
use std::env;
|
|
use std::sync::Arc;
|
|
|
|
pub struct GetActiveWorktreeContextHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for GetActiveWorktreeContextHandler {
|
|
fn name(&self) -> &'static str {
|
|
"get_active_worktree_context"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<GetActiveWorktreeContextTool>(
|
|
"get_active_worktree_context",
|
|
"Get the active worktree context, including branch name, modified files, and a truncated git diff.",
|
|
)
|
|
}
|
|
|
|
async fn execute(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<String> {
|
|
let result =
|
|
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
|
let cwd = env::current_dir().map_err(|e| crate::error::AppError::Internal(format!("Failed to get current directory: {}", e)))?;
|
|
let repo = git2::Repository::discover(&cwd).map_err(|e| {
|
|
crate::error::AppError::Internal(format!("Not in a git repository: {}", e))
|
|
})?;
|
|
|
|
let mut branch_name = String::new();
|
|
if let Ok(head) = repo.head()
|
|
&& let Some(name) = head.shorthand()
|
|
{
|
|
branch_name = name.to_string();
|
|
}
|
|
|
|
let mut opts = git2::DiffOptions::new();
|
|
let mut diff = None;
|
|
|
|
// Try to diff against HEAD
|
|
if let Ok(tree) = repo.head().and_then(|h| h.peel_to_tree()) {
|
|
diff = repo
|
|
.diff_tree_to_workdir_with_index(Some(&tree), Some(&mut opts))
|
|
.ok();
|
|
}
|
|
|
|
let mut files = Vec::new();
|
|
let mut diff_output = String::new();
|
|
|
|
if let Some(diff) = diff {
|
|
let _ = diff.print(git2::DiffFormat::Patch, |_delta, _hunk, line| {
|
|
match line.origin() {
|
|
'+' | '-' | ' ' => diff_output.push(line.origin()),
|
|
_ => {}
|
|
}
|
|
let content = std::str::from_utf8(line.content()).unwrap_or("");
|
|
diff_output.push_str(content);
|
|
true
|
|
});
|
|
|
|
for delta in diff.deltas() {
|
|
if let Some(path) = delta.new_file().path() {
|
|
let path_str = path.to_string_lossy().into_owned();
|
|
files.push(crate::handlers::utils::sanitize_path(&path_str));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Truncate diff output if it's too large to save tokens
|
|
if diff_output.len() > 10000 {
|
|
let valid_boundary = diff_output.floor_char_boundary(10000);
|
|
diff_output.truncate(valid_boundary);
|
|
diff_output.push_str("\n... [Diff truncated due to size]");
|
|
}
|
|
|
|
let mut out = String::new();
|
|
out.push_str("## Active Worktree Context\n\n");
|
|
out.push_str(&format!("**Branch:** `{}`\n\n", branch_name));
|
|
|
|
out.push_str("### Modified Files\n");
|
|
if files.is_empty() {
|
|
out.push_str("No modified files.\n");
|
|
} else {
|
|
for f in files {
|
|
out.push_str(&format!("- `{}`\n", f));
|
|
}
|
|
}
|
|
|
|
out.push_str("\n### Diff\n");
|
|
if diff_output.is_empty() {
|
|
out.push_str("No changes.\n");
|
|
} else {
|
|
out.push_str("```diff\n");
|
|
out.push_str(&diff_output);
|
|
out.push_str("\n```\n");
|
|
}
|
|
|
|
Ok(out)
|
|
})
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
|
|
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
pub struct QueryGitDiffsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for QueryGitDiffsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"query_git_diffs"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<crate::tools::QueryGitDiffsTool>(
|
|
"query_git_diffs",
|
|
"Query recent git commit history, diffs, and change ledger entries.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: crate::tools::QueryGitDiffsTool =
|
|
serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let limit = req.limit.unwrap_or(5);
|
|
let q = req.query.to_lowercase();
|
|
|
|
let mut matches = Vec::new();
|
|
|
|
state.code.ledger.read_with(|ledger| {
|
|
for change in ledger {
|
|
if change.file_path.to_lowercase().contains(&q)
|
|
|| change.description.to_lowercase().contains(&q)
|
|
|| change.git_commit.as_ref().is_some_and(|c| c.contains(&q))
|
|
{
|
|
matches.push(json!({
|
|
"file_path": crate::handlers::utils::sanitize_path(&change.file_path),
|
|
"description": change.description,
|
|
"commit": change.git_commit,
|
|
"branch": change.git_branch,
|
|
"timestamp": change.timestamp,
|
|
}));
|
|
if matches.len() >= limit {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
if matches.len() < limit {
|
|
let remaining = limit - matches.len();
|
|
let git_matches = tokio::task::spawn_blocking(move || {
|
|
let mut results = Vec::new();
|
|
let cwd = env::current_dir().unwrap_or_default();
|
|
if let Ok(repo) = git2::Repository::discover(&cwd)
|
|
&& let Ok(mut revwalk) = repo.revwalk()
|
|
{
|
|
let _ = revwalk.push_head();
|
|
let mut count = 0;
|
|
for oid in revwalk.flatten() {
|
|
if count >= remaining {
|
|
break;
|
|
}
|
|
if let Ok(commit) = repo.find_commit(oid) {
|
|
let summary = commit.summary().unwrap_or("");
|
|
if summary.to_lowercase().contains(&q) {
|
|
count += 1;
|
|
results.push(json!({
|
|
"commit_id": oid.to_string(),
|
|
"author": commit.author().name().unwrap_or("unknown"),
|
|
"message": summary,
|
|
"timestamp": commit.time().seconds(),
|
|
}));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
results
|
|
})
|
|
.await
|
|
.unwrap_or_default();
|
|
matches.extend(git_matches);
|
|
}
|
|
|
|
let mut out = String::new();
|
|
out.push_str(&format!("## Git Diffs for `{}`\n\n", req.query));
|
|
for m in &matches {
|
|
if let Some(file) = m.get("file_path").and_then(|f| f.as_str()) {
|
|
let desc = m.get("description").and_then(|d| d.as_str()).unwrap_or("");
|
|
let commit = m.get("commit").and_then(|c| c.as_str()).unwrap_or("");
|
|
out.push_str(&format!("- **File:** `{}` | **Commit:** `{}`\n - {}\n", file, commit, desc));
|
|
} else {
|
|
let commit = m.get("commit_id").and_then(|c| c.as_str()).unwrap_or("");
|
|
let author = m.get("author").and_then(|a| a.as_str()).unwrap_or("");
|
|
let msg = m.get("message").and_then(|m| m.as_str()).unwrap_or("");
|
|
out.push_str(&format!("- **Commit:** `{}` | **Author:** {}\n - {}\n", commit, author, msg));
|
|
}
|
|
}
|
|
if matches.is_empty() {
|
|
out.push_str("No git diffs found.\n");
|
|
}
|
|
|
|
Ok(out)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
use std::sync::Arc;
|
|
use tempfile::tempdir;
|
|
|
|
#[tokio::test]
|
|
async fn test_get_active_worktree_context() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = GetActiveWorktreeContextHandler;
|
|
|
|
let result = handler
|
|
.execute(json!({}), state)
|
|
.await
|
|
.map_err(|e| format!("Failed to get worktree context: {}", e))
|
|
.unwrap();
|
|
|
|
assert!(result.contains("## Active Worktree Context"));
|
|
assert!(result.contains("**Branch:**"));
|
|
assert!(result.contains("### Modified Files"));
|
|
assert!(result.contains("### Diff"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_get_active_worktree_context_empty_git_repo() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
let handler = GetActiveWorktreeContextHandler;
|
|
|
|
let result = handler
|
|
.execute(serde_json::json!({}), state)
|
|
.await
|
|
.map_err(|e| format!("Failed to get worktree context: {}", e))
|
|
.unwrap();
|
|
|
|
assert!(result.contains("## Active Worktree Context"));
|
|
}
|
|
}
|