Skip to main content

rtk_mcp_qwen/
lib.rs

1use serde::{Deserialize, Serialize};
2
3// ── Qwen PreToolUse input ──────────────────────────────────────────────
4#[derive(Debug, Deserialize, PartialEq)]
5pub struct QwenInput {
6    pub tool_name: String,
7    pub tool_input: ToolInput,
8}
9
10#[derive(Debug, Deserialize, PartialEq)]
11pub struct ToolInput {
12    pub command: Option<String>,
13}
14
15// ── Qwen PreToolUse output ─────────────────────────────────────────────
16#[derive(Debug, Serialize, Deserialize, PartialEq)]
17pub struct QwenOutput {
18    pub hook_specific_output: HookSpecificOutput,
19}
20
21#[derive(Debug, Serialize, Deserialize, PartialEq)]
22pub struct HookSpecificOutput {
23    pub hook_event_name: String,
24    pub permission_decision: String,
25    pub permission_decision_reason: String,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub updated_input: Option<UpdatedInput>,
28}
29
30#[derive(Debug, Serialize, Deserialize, PartialEq)]
31pub struct UpdatedInput {
32    pub command: String,
33}
34
35impl QwenOutput {
36    pub fn pass_through(reason: &str) -> Self {
37        Self {
38            hook_specific_output: HookSpecificOutput {
39                hook_event_name: "PreToolUse".into(),
40                permission_decision: "allow".into(),
41                permission_decision_reason: reason.into(),
42                updated_input: None,
43            },
44        }
45    }
46
47    pub fn rewritten(original: &str, rewritten: &str) -> Self {
48        Self {
49            hook_specific_output: HookSpecificOutput {
50                hook_event_name: "PreToolUse".into(),
51                permission_decision: "allow".into(),
52                permission_decision_reason: format!(
53                    "RTK rewrite applied: {} -> {}",
54                    original, rewritten
55                ),
56                updated_input: Some(UpdatedInput {
57                    command: rewritten.into(),
58                }),
59            },
60        }
61    }
62
63    /// Serialize to JSON string, returning None on failure.
64    pub fn to_json(&self) -> Option<String> {
65        serde_json::to_string(self).ok()
66    }
67
68    /// Parse from JSON string.
69    pub fn from_json(s: &str) -> Option<Self> {
70        serde_json::from_str(s).ok()
71    }
72}
73
74/// The rtk executable used when `RTK_BIN` says nothing.
75pub const DEFAULT_RTK_BIN: &str = "rtk";
76
77/// Decide which rtk executable to invoke.
78///
79/// A blank value counts as unset. `RTK_BIN=` is how a shell commonly spells
80/// "clear this", and `env::var` reports it as `Ok("")` rather than an error —
81/// so taking it literally means trying to spawn the empty string, which fails
82/// with a confusing OS error instead of quietly using rtk from `PATH`.
83///
84/// The variable is deliberately un-prefixed: the sibling `rtk-mcp-agy` and
85/// `rtk-mcp-cc` plugins honour the same one, and all three resolve it the same
86/// way, so relocating rtk takes one variable rather than three.
87pub fn resolve_rtk_bin(override_value: Option<String>) -> String {
88    override_value
89        .map(|v| v.trim().to_string())
90        .filter(|v| !v.is_empty())
91        .unwrap_or_else(|| DEFAULT_RTK_BIN.to_string())
92}
93
94/// Process a single hook input and return the output.
95///
96/// If `rtk_rewrite_fn` returns Some, the command is rewritten.
97/// Otherwise the original command passes through.
98pub fn process_hook(
99    input_json: &str,
100    rtk_rewrite_fn: impl FnOnce(&str) -> Option<String>,
101) -> QwenOutput {
102    // Parse input
103    let data: QwenInput = match serde_json::from_str(input_json) {
104        Ok(v) => v,
105        Err(_) => return QwenOutput::pass_through("JSON parse error, passing through"),
106    };
107
108    // Only process run_shell_command
109    if data.tool_name != "run_shell_command" {
110        return QwenOutput::pass_through("Non-shell command, passing through");
111    }
112
113    let command = match data.tool_input.command {
114        Some(cmd) if !cmd.trim().is_empty() => cmd,
115        _ => return QwenOutput::pass_through("No command in input"),
116    };
117
118    // Ask RTK to rewrite the command
119    match rtk_rewrite_fn(&command) {
120        Some(rewritten) => QwenOutput::rewritten(&command, &rewritten),
121        None => QwenOutput::pass_through(&format!("No RTK rewrite available for: {}", command)),
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn test_rtk_bin_defaults_to_the_path_lookup() {
131        assert_eq!(resolve_rtk_bin(None), DEFAULT_RTK_BIN);
132    }
133
134    #[test]
135    fn test_rtk_bin_honours_an_override() {
136        assert_eq!(
137            resolve_rtk_bin(Some("/opt/rtk/bin/rtk".to_string())),
138            "/opt/rtk/bin/rtk"
139        );
140        assert_eq!(resolve_rtk_bin(Some("  rtk-dev  ".to_string())), "rtk-dev");
141    }
142
143    #[test]
144    fn test_blank_rtk_bin_does_not_erase_the_default() {
145        // `env::var` returns Ok("") for `RTK_BIN=`, so without this the hook
146        // would try to spawn the empty string rather than fall back to PATH.
147        for blank in ["", "   ", "\t", "\n"] {
148            assert_eq!(
149                resolve_rtk_bin(Some(blank.to_string())),
150                DEFAULT_RTK_BIN,
151                "blank value {blank:?} must count as unset"
152            );
153        }
154    }
155
156    #[test]
157    fn test_pass_through_output() {
158        let output = QwenOutput::pass_through("test reason");
159        assert_eq!(output.hook_specific_output.hook_event_name, "PreToolUse");
160        assert_eq!(output.hook_specific_output.permission_decision, "allow");
161        assert_eq!(
162            output.hook_specific_output.permission_decision_reason,
163            "test reason"
164        );
165        assert!(output.hook_specific_output.updated_input.is_none());
166    }
167
168    #[test]
169    fn test_rewritten_output() {
170        let output = QwenOutput::rewritten("ls", "ls -la");
171        assert_eq!(output.hook_specific_output.permission_decision, "allow");
172        assert!(output.hook_specific_output.updated_input.is_some());
173        assert_eq!(
174            output
175                .hook_specific_output
176                .updated_input
177                .as_ref()
178                .unwrap()
179                .command,
180            "ls -la"
181        );
182    }
183
184    #[test]
185    fn test_process_hook_json_parse_error() {
186        let result = process_hook("not json", |_| None);
187        assert_eq!(
188            result.hook_specific_output.permission_decision_reason,
189            "JSON parse error, passing through"
190        );
191    }
192
193    #[test]
194    fn test_process_hook_non_shell_command() {
195        let input = r#"{"tool_name": "read_file", "tool_input": {}}"#;
196        let result = process_hook(input, |_| None);
197        assert_eq!(
198            result.hook_specific_output.permission_decision_reason,
199            "Non-shell command, passing through"
200        );
201    }
202
203    #[test]
204    fn test_process_hook_no_command() {
205        let input = r#"{"tool_name": "run_shell_command", "tool_input": {}}"#;
206        let result = process_hook(input, |_| None);
207        assert_eq!(
208            result.hook_specific_output.permission_decision_reason,
209            "No command in input"
210        );
211    }
212
213    #[test]
214    fn test_process_hook_empty_command() {
215        let input = r#"{"tool_name": "run_shell_command", "tool_input": {"command": "  "}}"#;
216        let result = process_hook(input, |_| None);
217        assert_eq!(
218            result.hook_specific_output.permission_decision_reason,
219            "No command in input"
220        );
221    }
222
223    #[test]
224    fn test_process_hook_with_rewrite() {
225        let input = r#"{"tool_name": "run_shell_command", "tool_input": {"command": "ls"}}"#;
226        let result = process_hook(input, |_| Some("ls -la --color".to_string()));
227        assert!(result.hook_specific_output.updated_input.is_some());
228        assert_eq!(
229            result
230                .hook_specific_output
231                .updated_input
232                .as_ref()
233                .unwrap()
234                .command,
235            "ls -la --color"
236        );
237    }
238
239    #[test]
240    fn test_process_hook_no_rewrite() {
241        let input = r#"{"tool_name": "run_shell_command", "tool_input": {"command": "ls"}}"#;
242        let result = process_hook(input, |_| None);
243        assert!(result.hook_specific_output.updated_input.is_none());
244        assert_eq!(
245            result.hook_specific_output.permission_decision_reason,
246            "No RTK rewrite available for: ls"
247        );
248    }
249
250    /// The keys actually emitted under `hook_specific_output`.
251    fn emitted_fields(output: &QwenOutput) -> serde_json::Map<String, serde_json::Value> {
252        let json = output.to_json().expect("serialize");
253        let value: serde_json::Value = serde_json::from_str(&json).expect("valid JSON");
254        value["hook_specific_output"]
255            .as_object()
256            .expect("hook_specific_output object")
257            .clone()
258    }
259
260    #[test]
261    fn test_pass_through_omits_updated_input_from_the_wire() {
262        // `skip_serializing_if = "Option::is_none"` is a wire contract: a
263        // pass-through must omit the key, not send `"updated_input": null`.
264        //
265        // The roundtrip tests below cannot check this. `Option` deserializes a
266        // missing key and an explicit null identically, so serialize →
267        // deserialize → compare stays equal whether or not the attribute is
268        // there. Only the emitted keys reveal it.
269        let fields = emitted_fields(&QwenOutput::pass_through("reason"));
270        assert!(!fields.contains_key("updated_input"), "got: {fields:?}");
271        // The three fields that must always be present.
272        for key in [
273            "hook_event_name",
274            "permission_decision",
275            "permission_decision_reason",
276        ] {
277            assert!(fields.contains_key(key), "missing {key}: {fields:?}");
278        }
279    }
280
281    #[test]
282    fn test_rewrite_includes_updated_input_on_the_wire() {
283        let fields = emitted_fields(&QwenOutput::rewritten("ls", "ls -la"));
284        assert_eq!(
285            fields["updated_input"]["command"], "ls -la",
286            "got: {fields:?}"
287        );
288    }
289
290    #[test]
291    fn test_output_json_roundtrip() {
292        let output = QwenOutput::rewritten("ls", "ls -la");
293        let json = output.to_json().expect("serialize");
294        let parsed = QwenOutput::from_json(&json).expect("deserialize");
295        assert_eq!(output, parsed);
296    }
297
298    #[test]
299    fn test_output_json_pass_through() {
300        let output = QwenOutput::pass_through("reason");
301        let json = output.to_json().expect("serialize");
302        let parsed = QwenOutput::from_json(&json).expect("deserialize");
303        assert_eq!(output, parsed);
304    }
305}