riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor remaining tools to tagged enums

Browse Source 
Refactored DecisionsTool, TechDebtTool, ManageCheckpointTool, HypothesesTool, AgentSignalsTool to use tagged enums for precise JSON schemas, and updated their handlers.
This commit is contained in:
Riz Ashraf committed 2026-10-10 10:23:58 +01:00
1 parent 8ff43487bc
commit 2cfaeda0cb
3 files changed
+438 -497
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/meta.rs
+95 -140
View File
@@ -383,18 +383,14 @@ impl McpTool for DecisionsHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: DecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req
.namespace
.unwrap_or_else(crate::models::default_namespace);
match req.action {
DecisionAction::Log => {
let title = req.title.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'title' for action 'log'. Next step: Provide ADR 'title' string in request and retry.".to_string())
})?;
let status = req.status.unwrap_or_else(|| "accepted".to_string());
let context = req.context.unwrap_or_default();
let decision = req.decision.unwrap_or_default();
let consequence = req.consequences.unwrap_or_default();
match req {
DecisionsTool::Log { title, status, context, decision, consequences, namespace, repo_name, author, affected_components, alternatives_considered, supersedes, git_commit, git_branch, task_id } => {
let ns = namespace.clone().unwrap_or_else(crate::models::default_namespace);
let title = title;
let status = status.unwrap_or_else(|| "accepted".to_string());
let context = context.unwrap_or_default();
let decision = decision.unwrap_or_default();
let consequence = consequences.unwrap_or_default();
let status_lower = status.to_ascii_lowercase();
let resolved_at = if status_lower == "implemented" || status_lower == "resolved" {
Some(crate::handlers::utils::now_secs())
@@ -407,7 +403,7 @@ impl McpTool for DecisionsHandler {
let mut adrs_to_index = Vec::new();
state.code.adrs.modify(|adrs| {
if let Some(superseded_id) = &req.supersedes {
if let Some(superseded_id) = &supersedes {
for old_adr in adrs.iter_mut() {
if old_adr.id.eq_ignore_ascii_case(superseded_id) {
old_adr.status = "superseded".to_string();
@@ -425,17 +421,17 @@ impl McpTool for DecisionsHandler {
decision: decision.clone(),
consequence,
status,
supersedes: req.supersedes,
supersedes: supersedes,
timestamp: crate::handlers::utils::now_secs(),
namespace: ns,
repo_name: req.repo_name,
alternatives_considered: req.alternatives_considered.unwrap_or_default(),
affected_components: req.affected_components.unwrap_or_default(),
author: req.author,
git_commit: req.git_commit,
git_branch: req.git_branch,
repo_name: repo_name,
alternatives_considered: alternatives_considered.unwrap_or_default(),
affected_components: affected_components.unwrap_or_default(),
author: author,
git_commit: git_commit,
git_branch: git_branch,
resolved_at,
task_id: req.task_id,
task_id: task_id,
};
adrs_to_index.push(a.clone());
@@ -453,17 +449,15 @@ impl McpTool for DecisionsHandler {
);
Ok(format!("Logged decision {}: {}", final_id, title))
}
DecisionAction::Update => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'update'. Next step: Provide ADR 'id' string in request and retry.".to_string())
})?;
DecisionsTool::Update { id, title, status, context, decision, consequences, git_commit, git_branch, task_id, namespace: _, repo_name, author, affected_components, alternatives_considered, supersedes } => {
let id = id;
let mut updated_adr = None;
let mut adrs_to_index = Vec::new();
state.code.adrs.modify(|adrs| {
let target_pos = adrs.iter().position(|a| a.id.eq_ignore_ascii_case(&id));
if let Some(pos) = target_pos {
if let Some(superseded_id) = &req.supersedes
if let Some(superseded_id) = &supersedes
&& let Some(s_pos) = adrs
.iter()
.position(|a| a.id.eq_ignore_ascii_case(superseded_id))
@@ -474,19 +468,19 @@ impl McpTool for DecisionsHandler {
}
let a = &mut adrs[pos];
if let Some(t) = req.title {
if let Some(t) = title {
a.title = t;
}
if let Some(c) = req.context {
if let Some(c) = context {
a.context = c;
}
if let Some(d) = req.decision {
if let Some(d) = decision {
a.decision = d;
}
if let Some(cons) = req.consequences {
if let Some(cons) = consequences {
a.consequence = cons;
}
if let Some(s) = req.status {
if let Some(s) = status {
let s_lower = s.to_ascii_lowercase();
if (s_lower == "implemented" || s_lower == "resolved")
&& a.resolved_at.is_none()
@@ -497,29 +491,29 @@ impl McpTool for DecisionsHandler {
}
a.status = s;
}
if req.supersedes.is_some() {
a.supersedes = req.supersedes;
if supersedes.is_some() {
a.supersedes = supersedes;
}
if req.repo_name.is_some() {
a.repo_name = req.repo_name;
if repo_name.is_some() {
a.repo_name = repo_name;
}
if let Some(alts) = req.alternatives_considered {
if let Some(alts) = alternatives_considered {
a.alternatives_considered = alts;
}
if let Some(aff) = req.affected_components {
if let Some(aff) = affected_components {
a.affected_components = aff;
}
if req.author.is_some() {
a.author = req.author;
if author.is_some() {
a.author = author;
}
if req.git_commit.is_some() {
a.git_commit = req.git_commit;
if git_commit.is_some() {
a.git_commit = git_commit;
}
if req.git_branch.is_some() {
a.git_branch = req.git_branch;
if git_branch.is_some() {
a.git_branch = git_branch;
}
if req.task_id.is_some() {
a.task_id = req.task_id;
if task_id.is_some() {
a.task_id = task_id;
}
adrs_to_index.push(a.clone());
@@ -548,9 +542,10 @@ impl McpTool for DecisionsHandler {
)))
}
}
DecisionAction::Query => {
let limit = req.limit.unwrap_or(20);
let include_body = req.include_body.unwrap_or(true);
DecisionsTool::Query { query, namespace, limit, include_body } => {
let ns = namespace.clone().unwrap_or_else(crate::models::default_namespace);
let limit = limit.unwrap_or(20);
let include_body = include_body.unwrap_or(true);
let data = state.code.adrs.read_with(|adrs| {
let filtered: Vec<_> = adrs
.iter()
@@ -558,7 +553,7 @@ impl McpTool for DecisionsHandler {
if !a.namespace.eq_ignore_ascii_case(&ns) && ns != "global" {
return false;
}
if let Some(q) = &req.query {
if let Some(q) = &query {
crate::handlers::utils::contains_ignore_ascii_case(&a.title, q)
|| crate::handlers::utils::contains_ignore_ascii_case(
&a.context, q,
@@ -601,10 +596,8 @@ impl McpTool for DecisionsHandler {
})?;
Ok(data)
}
DecisionAction::Delete => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'delete'. Next step: Provide ADR 'id' string in request and retry.".to_string())
})?;
DecisionsTool::Delete { id } => {
let id = id;
let mut found = false;
state.code.adrs.modify(|adrs| {
if let Some(pos) = adrs.iter().position(|a| a.id == id) {
@@ -644,15 +637,11 @@ impl McpTool for TechDebtHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: TechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req
.namespace
.unwrap_or_else(crate::models::default_namespace);
match req.action {
TechDebtAction::Log => {
let desc = req.description.or(req.title).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'description' for action 'log'. Next step: Provide tech debt 'description' in request and retry.".to_string())
})?;
let ideal = req.ideal_solution.unwrap_or_default();
match req {
TechDebtTool::Log { description, title, ideal_solution, git_commit, git_branch, symbol_references, line_range, namespace, severity, repo_name, file_path, workaround, effort_estimate } => {
let ns = namespace.clone().unwrap_or_else(crate::models::default_namespace);
let desc = description;
let ideal = ideal_solution.unwrap_or_default();
let text_to_embed = format!("Description: {}\nIdeal Solution: {}", desc, ideal);
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
@@ -665,15 +654,15 @@ impl McpTool for TechDebtHandler {
ideal_solution: ideal,
is_resolved: false,
created_at: crate::handlers::utils::now_secs(),
git_commit: req.git_commit,
git_branch: req.git_branch,
git_commit: git_commit,
git_branch: git_branch,
embedding,
repo_name: req.repo_name,
severity: req.severity,
file_path: req.file_path,
line_range: req.line_range,
workaround: req.workaround,
effort_estimate: req.effort_estimate,
repo_name: repo_name,
severity: severity,
file_path: file_path,
line_range: line_range,
workaround: workaround,
effort_estimate: effort_estimate,
});
if debts.len() > 300 {
let severity_rank =
@@ -697,10 +686,8 @@ impl McpTool for TechDebtHandler {
});
Ok("Tech debt logged".to_string())
}
TechDebtAction::Resolve => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'resolve'. Next step: Provide tech debt 'id' string in request and retry.".to_string())
})?;
TechDebtTool::Resolve { id } => {
let id = id;
let mut found = false;
state.code.tech_debts.modify(|debts| {
for d in debts.iter_mut() {
@@ -720,9 +707,10 @@ impl McpTool for TechDebtHandler {
))
}
}
TechDebtAction::List => {
let inc = req.include_resolved.unwrap_or(false);
let level = req.summary_level.as_deref().unwrap_or("detailed");
TechDebtTool::List { include_resolved, namespace, summary_level, max_tokens } => {
let ns = namespace.clone().unwrap_or_else(crate::models::default_namespace);
let inc = include_resolved.unwrap_or(false);
let level = summary_level.as_deref().unwrap_or("detailed");
let data = state.code.tech_debts.read_with(|debts| {
let filtered: Vec<_> = debts
.iter()
@@ -743,7 +731,7 @@ impl McpTool for TechDebtHandler {
})
.collect();
let mut json_str = serde_yaml::to_string(&filtered)?;
if let Some(max_t) = req.max_tokens {
if let Some(max_t) = max_tokens {
let char_limit = max_t * 4;
if json_str.len() > char_limit {
json_str.truncate(char_limit);
@@ -1199,13 +1187,9 @@ impl McpTool for ManageCheckpointHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ManageCheckpointTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
CheckpointAction::Create => {
let name = req.name_or_id.ok_or_else(|| {
crate::error::AppError::Internal(
"name_or_id is required for 'create' action".to_string(),
)
})?;
match req {
ManageCheckpointTool::Create { name_or_id, description, namespace } => {
let name = name_or_id.ok_or_else(|| crate::error::AppError::Internal("name_or_id is required".to_string()))?;
let target_dir = state.base_dir.join("checkpoints").join(&name);
if let Err(e) = std::fs::create_dir_all(&target_dir) {
@@ -1230,15 +1214,12 @@ impl McpTool for ManageCheckpointHandler {
.read_with(|d| serde_json::to_string(d).unwrap_or_default());
let _ = std::fs::write(target_dir.join("tech_debts.json"), debts_json);
if let Some(desc) = &req.description {
if let Some(desc) = &description {
let snapshot_id = format!(
"SNAP-{}",
uuid::Uuid::new_v4().to_string()[..8].to_uppercase()
);
let ns = req
.namespace
.clone()
.unwrap_or_else(|| "global".to_string());
let ns = namespace.clone().unwrap_or_else(|| "global".to_string());
let snapshot = crate::models::StateSnapshot {
id: snapshot_id,
timestamp: crate::handlers::utils::now_secs(),
@@ -1251,12 +1232,8 @@ impl McpTool for ManageCheckpointHandler {
Ok(format!("Checkpoint '{}' created successfully.", name))
}
CheckpointAction::Restore => {
let name = req.name_or_id.ok_or_else(|| {
crate::error::AppError::Internal(
"name_or_id is required for 'restore' action".to_string(),
)
})?;
ManageCheckpointTool::Restore { name_or_id, namespace } => {
let name = name_or_id.ok_or_else(|| crate::error::AppError::Internal("name_or_id is required".to_string()))?;
let target_dir = state.base_dir.join("checkpoints").join(&name);
if !target_dir.exists() {
@@ -1297,7 +1274,7 @@ impl McpTool for ManageCheckpointHandler {
Ok(format!("Checkpoint '{}' restored successfully.", name))
}
CheckpointAction::List => {
ManageCheckpointTool::List { namespace } => {
let mut list = Vec::new();
let checkpoints_dir = state.base_dir.join("checkpoints");
if let Ok(entries) = std::fs::read_dir(&checkpoints_dir) {
@@ -1315,12 +1292,8 @@ impl McpTool for ManageCheckpointHandler {
}
Ok(serde_yaml::to_string(&list)?)
}
CheckpointAction::Delete => {
let name = req.name_or_id.ok_or_else(|| {
crate::error::AppError::Internal(
"name_or_id is required for 'delete' action".to_string(),
)
})?;
ManageCheckpointTool::Delete { name_or_id, namespace } => {
let name = name_or_id.ok_or_else(|| crate::error::AppError::Internal("name_or_id is required".to_string()))?;
let target_dir = state.base_dir.join("checkpoints").join(&name);
if target_dir.exists() {
let _ = std::fs::remove_dir_all(&target_dir);
@@ -1516,14 +1489,9 @@ impl McpTool for HypothesesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: HypothesesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
HypothesisAction::Log => {
let hyp_text = req.hypothesis.ok_or_else(|| {
crate::error::AppError::Internal(
"Missing required 'hypothesis' for action 'log'".to_string(),
)
})?;
match req {
HypothesesTool::Log { hypothesis, task_id, status, evidence } => {
let hyp_text = hypothesis;
let hyp_id = format!(
"HYP-{}",
uuid::Uuid::new_v4().to_string()[..8].to_uppercase()
@@ -1532,10 +1500,10 @@ impl McpTool for HypothesesHandler {
let record = crate::models::Hypothesis {
id: hyp_id.clone(),
task_id: req.task_id,
task_id: task_id,
hypothesis: hyp_text,
status: req.status.unwrap_or_else(|| "unverified".to_string()),
evidence: req.evidence,
status: status.unwrap_or_else(|| "unverified".to_string()),
evidence: evidence,
timestamp,
..Default::default()
};
@@ -1544,17 +1512,17 @@ impl McpTool for HypothesesHandler {
Ok(format!("Hypothesis '{}' logged successfully.", hyp_id))
}
HypothesisAction::Query => {
HypothesesTool::Query { query, task_id, status } => {
let hypotheses = state.code.hypotheses.read_with(|h| h.clone());
let filtered: Vec<_> = hypotheses
.into_iter()
.filter(|h| {
if let Some(tid) = &req.task_id
if let Some(tid) = &task_id
&& h.task_id.as_ref() != Some(tid)
{
return false;
}
if let Some(q) = &req.query {
if let Some(q) = &query {
let lq = q.to_lowercase();
return h.hypothesis.to_lowercase().contains(&lq)
|| h.evidence
@@ -1665,24 +1633,11 @@ impl McpTool for AgentSignalsHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: AgentSignalsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
AgentSignalAction::Broadcast => {
let sender = req.sender.ok_or_else(|| {
crate::error::AppError::Internal(
"Missing required 'sender' for action 'broadcast'".to_string(),
)
})?;
let signal_type = req.signal_type.ok_or_else(|| {
crate::error::AppError::Internal(
"Missing required 'signal_type' for action 'broadcast'".to_string(),
)
})?;
let payload = req.payload.ok_or_else(|| {
crate::error::AppError::Internal(
"Missing required 'payload' for action 'broadcast'".to_string(),
)
})?;
match req {
AgentSignalsTool::Broadcast { sender, signal_type, payload, ttl_seconds } => {
let sender = sender;
let signal_type = signal_type;
let payload = payload;
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
@@ -1696,7 +1651,7 @@ impl McpTool for AgentSignalsHandler {
signal_type: signal_type.clone(),
payload,
timestamp,
ttl_seconds: req.ttl_seconds,
ttl_seconds: ttl_seconds,
..Default::default()
};
@@ -1724,7 +1679,7 @@ impl McpTool for AgentSignalsHandler {
sig_id, sender
))
}
AgentSignalAction::Query => {
AgentSignalsTool::Query { sender, signal_type, limit } => {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
@@ -1739,19 +1694,19 @@ impl McpTool for AgentSignalsHandler {
{
return false;
}
if let Some(sender) = &req.sender
if let Some(sender) = &sender
&& s.sender.to_lowercase() != sender.to_lowercase()
{
return false;
}
if let Some(st) = &req.signal_type
if let Some(st) = &signal_type
&& s.signal_type.to_lowercase() != st.to_lowercase()
{
return false;
}
true
})
.take(req.limit.unwrap_or(20))
.take(limit.unwrap_or(20))
.cloned()
.collect::<Vec<_>>()
});