Skip to main content

lean_ctx/tools/registered/
ctx_feedback.rs

1use rmcp::model::Tool;
2use rmcp::ErrorData;
3use serde_json::{json, Map, Value};
4
5use crate::server::tool_trait::{get_int, get_str, McpTool, ToolContext, ToolOutput};
6use crate::tool_defs::tool_def;
7
8pub struct CtxFeedbackTool;
9
10impl McpTool for CtxFeedbackTool {
11    fn name(&self) -> &'static str {
12        "ctx_feedback"
13    }
14
15    fn tool_def(&self) -> Tool {
16        tool_def(
17            "ctx_feedback",
18            "Harness feedback for LLM output tokens/latency (local-first). Actions: record|report|json|reset|status.",
19            json!({
20                "type": "object",
21                "properties": {
22                    "action": {
23                        "type": "string",
24                        "enum": ["record", "report", "json", "reset", "status"],
25                        "description": "Operation to perform (default: report)"
26                    },
27                    "agent_id": { "type": "string", "description": "Agent ID (optional; defaults to current agent)" },
28                    "intent": { "type": "string", "description": "Intent/task string (optional)" },
29                    "model": { "type": "string", "description": "Model identifier (optional)" },
30                    "llm_input_tokens": { "type": "integer", "description": "Required for action=record" },
31                    "llm_output_tokens": { "type": "integer", "description": "Required for action=record" },
32                    "latency_ms": { "type": "integer", "description": "Optional for action=record" },
33                    "note": { "type": "string", "description": "Optional note (no prompts/PII)" },
34                    "limit": { "type": "integer", "description": "For report/json: max recent events (default: 500)" }
35                }
36            }),
37        )
38    }
39
40    fn handle(
41        &self,
42        args: &Map<String, Value>,
43        ctx: &ToolContext,
44    ) -> Result<ToolOutput, ErrorData> {
45        let action = get_str(args, "action").unwrap_or_else(|| "report".to_string());
46        let limit = get_int(args, "limit").map_or(500, |n| n.max(1) as usize);
47
48        let result = match action.as_str() {
49            "record" => {
50                let current_agent_id = ctx
51                    .agent_id
52                    .as_ref()
53                    .and_then(|a| tokio::task::block_in_place(|| a.blocking_read()).clone());
54                let agent_id = get_str(args, "agent_id").or(current_agent_id);
55                let agent_id = agent_id.ok_or_else(|| {
56                    ErrorData::invalid_params(
57                        "agent_id is required (or register an agent via project_root detection first)",
58                        None,
59                    )
60                })?;
61
62                let (ctx_read_last_mode, ctx_read_modes) = if let Some(ref tc) = ctx.tool_calls {
63                    let calls = tokio::task::block_in_place(|| tc.blocking_read());
64                    let mut last: Option<String> = None;
65                    let mut modes: std::collections::BTreeMap<String, u64> =
66                        std::collections::BTreeMap::new();
67                    for rec in calls.iter().rev().take(50) {
68                        if rec.tool != "ctx_read" {
69                            continue;
70                        }
71                        if let Some(m) = rec.mode.as_ref() {
72                            *modes.entry(m.clone()).or_insert(0) += 1;
73                            if last.is_none() {
74                                last = Some(m.clone());
75                            }
76                        }
77                    }
78                    (last, if modes.is_empty() { None } else { Some(modes) })
79                } else {
80                    (None, None)
81                };
82
83                let llm_input_tokens = get_int(args, "llm_input_tokens").ok_or_else(|| {
84                    ErrorData::invalid_params("llm_input_tokens is required", None)
85                })?;
86                let llm_output_tokens = get_int(args, "llm_output_tokens").ok_or_else(|| {
87                    ErrorData::invalid_params("llm_output_tokens is required", None)
88                })?;
89                if llm_input_tokens <= 0 || llm_output_tokens <= 0 {
90                    return Err(ErrorData::invalid_params(
91                        "llm_input_tokens and llm_output_tokens must be > 0",
92                        None,
93                    ));
94                }
95
96                let ev = crate::core::llm_feedback::LlmFeedbackEvent {
97                    agent_id,
98                    intent: get_str(args, "intent"),
99                    model: get_str(args, "model"),
100                    llm_input_tokens: llm_input_tokens as u64,
101                    llm_output_tokens: llm_output_tokens as u64,
102                    latency_ms: get_int(args, "latency_ms").map(|n| n.max(0) as u64),
103                    note: get_str(args, "note"),
104                    ctx_read_last_mode,
105                    ctx_read_modes,
106                    timestamp: chrono::Local::now().to_rfc3339(),
107                };
108                crate::tools::ctx_feedback::record(&ev)
109                    .unwrap_or_else(|e| format!("Error recording feedback: {e}"))
110            }
111            "status" => crate::tools::ctx_feedback::status(),
112            "json" => crate::tools::ctx_feedback::json(limit),
113            "reset" => crate::tools::ctx_feedback::reset(),
114            _ => crate::tools::ctx_feedback::report(limit),
115        };
116
117        Ok(ToolOutput {
118            text: result,
119            original_tokens: 0,
120            saved_tokens: 0,
121            mode: Some(action),
122            path: None,
123        })
124    }
125}