Skip to main content

lean_ctx/tools/registered/
ctx_feedback.rs

1use rmcp::ErrorData;
2use rmcp::model::Tool;
3use serde_json::{Map, Value, json};
4
5use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_int, get_str};
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            "Record and report LLM token/latency metrics — use to track efficiency and optimize context usage.\n\
19             WORKFLOW: action=record during each LLM call, then action=report for readable summary.\n\
20             Actions: record (log event), report (readable summary), json (machine-readable),\n\
21             reset (clear data), status (storage info).\n\
22             ANTIPATTERN: not for debugging code behavior — this tracks token/latency stats only.\n\
23             record requires llm_input_tokens + llm_output_tokens.",
24            json!({
25                "type": "object",
26                "properties": {
27                    "action": {
28                        "type": "string",
29                        "enum": ["record", "report", "json", "reset", "status"],
30                        "description": "record (log event) | report (summary) | json (data) | reset (clear) | status (storage)"
31                    },
32                    "agent_id": { "type": "string", "description": "Agent ID (default: current agent)" },
33                    "intent": { "type": "string", "description": "Intent/task string" },
34                    "model": { "type": "string", "description": "Model identifier" },
35                    "llm_input_tokens": { "type": "integer", "description": "Required for action=record" },
36                    "llm_output_tokens": { "type": "integer", "description": "Required for action=record" },
37                    "latency_ms": { "type": "integer", "description": "Latency in ms (for record)" },
38                    "note": { "type": "string", "description": "Note (no prompts/PII)" },
39                    "limit": { "type": "integer", "description": "Max recent events (default: 500)" }
40                },
41                "allOf": [
42                    {
43                        "if": { "properties": { "action": { "const": "record" } }, "required": ["action"] },
44                        "then": { "required": ["action", "llm_input_tokens", "llm_output_tokens"] }
45                    }
46                ]
47            }),
48        )
49    }
50
51    fn handle(
52        &self,
53        args: &Map<String, Value>,
54        ctx: &ToolContext,
55    ) -> Result<ToolOutput, ErrorData> {
56        let action = get_str(args, "action").unwrap_or_else(|| "report".to_string());
57        let limit = get_int(args, "limit").map_or(500, |n| n.max(1) as usize);
58
59        let result = match action.as_str() {
60            "record" => {
61                let current_agent_id = ctx
62                    .agent_id
63                    .as_ref()
64                    .and_then(|a| tokio::task::block_in_place(|| a.blocking_read()).clone());
65                let agent_id = get_str(args, "agent_id").or(current_agent_id);
66                let agent_id = agent_id.ok_or_else(|| {
67                    ErrorData::invalid_params(
68                        "agent_id is required (or register an agent via project_root detection first)",
69                        None,
70                    )
71                })?;
72
73                let (ctx_read_last_mode, ctx_read_modes) = if let Some(ref tc) = ctx.tool_calls {
74                    let calls = tokio::task::block_in_place(|| tc.blocking_read());
75                    let mut last: Option<String> = None;
76                    let mut modes: std::collections::BTreeMap<String, u64> =
77                        std::collections::BTreeMap::new();
78                    for rec in calls.iter().rev().take(50) {
79                        if rec.tool != "ctx_read" {
80                            continue;
81                        }
82                        if let Some(m) = rec.mode.as_ref() {
83                            *modes.entry(m.clone()).or_insert(0) += 1;
84                            if last.is_none() {
85                                last = Some(m.clone());
86                            }
87                        }
88                    }
89                    (last, if modes.is_empty() { None } else { Some(modes) })
90                } else {
91                    (None, None)
92                };
93
94                let llm_input_tokens = get_int(args, "llm_input_tokens").ok_or_else(|| {
95                    ErrorData::invalid_params("llm_input_tokens is required", None)
96                })?;
97                let llm_output_tokens = get_int(args, "llm_output_tokens").ok_or_else(|| {
98                    ErrorData::invalid_params("llm_output_tokens is required", None)
99                })?;
100                if llm_input_tokens <= 0 || llm_output_tokens <= 0 {
101                    return Err(ErrorData::invalid_params(
102                        "llm_input_tokens and llm_output_tokens must be > 0",
103                        None,
104                    ));
105                }
106
107                let ev = crate::core::llm_feedback::LlmFeedbackEvent {
108                    agent_id,
109                    intent: get_str(args, "intent"),
110                    model: get_str(args, "model"),
111                    llm_input_tokens: llm_input_tokens as u64,
112                    llm_output_tokens: llm_output_tokens as u64,
113                    latency_ms: get_int(args, "latency_ms").map(|n| n.max(0) as u64),
114                    note: get_str(args, "note"),
115                    ctx_read_last_mode,
116                    ctx_read_modes,
117                    timestamp: chrono::Local::now().to_rfc3339(),
118                };
119                crate::tools::ctx_feedback::record(&ev)
120                    .unwrap_or_else(|e| format!("Error recording feedback: {e}"))
121            }
122            "status" => crate::tools::ctx_feedback::status(),
123            "json" => crate::tools::ctx_feedback::json(limit),
124            "reset" => crate::tools::ctx_feedback::reset(),
125            _ => crate::tools::ctx_feedback::report(limit),
126        };
127
128        Ok(ToolOutput {
129            text: result,
130            original_tokens: 0,
131            saved_tokens: 0,
132            mode: Some(action),
133            path: None,
134            changed: false,
135            shell_outcome: None,
136            content_blocks: None,
137        })
138    }
139}