aegis_tools/
control_tool.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use serde_json::{json, Value};
4
5use crate::{Tool, ToolContext};
6
7/// Agent-callable system control tool. Allows the model to manage session
8/// behavior (steering, style, undo) through natural language rather than
9/// requiring users to know slash-command syntax.
10///
11/// Actions that mutate agent state return a `CMD:` prefixed string that the
12/// chat loop intercepts and executes. Read-only actions return results directly.
13pub struct ControlTool;
14
15/// Command prefix that the chat loop recognizes as an agent-issued directive.
16pub const CMD_PREFIX: &str = "CMD:";
17
18#[async_trait]
19impl Tool for ControlTool {
20    fn name(&self) -> &str {
21        "control"
22    }
23
24    fn description(&self) -> &str {
25        "Control aegis session behavior: change output style (normal/concise/minimal), \
26         manage steering instructions (add/remove/list/clear), undo the last turn, \
27         or start a new session. Use this when the user asks to adjust how you respond, \
28         add a persistent instruction, undo something, or start fresh."
29    }
30
31    fn parameters(&self) -> Value {
32        json!({
33            "type": "object",
34            "properties": {
35                "action": {
36                    "type": "string",
37                    "enum": ["style", "steer_add", "steer_remove", "steer_list", "steer_clear", "undo", "new_session"],
38                    "description": "style = set output verbosity; steer_* = manage steering instructions; undo = undo last turn; new_session = start fresh"
39                },
40                "value": {
41                    "type": "string",
42                    "description": "For style: 'normal'|'concise'|'minimal'. For steer_add: the instruction text. For steer_remove: the id prefix."
43                },
44                "turns": {
45                    "type": "integer",
46                    "description": "For steer_add: number of turns before expiry (omit for permanent)"
47                }
48            },
49            "required": ["action"]
50        })
51    }
52
53    async fn execute(&self, args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
54        let action = args["action"].as_str().unwrap_or("");
55        match action {
56            "style" => {
57                let value = args["value"].as_str().unwrap_or("normal");
58                let valid = matches!(value, "normal" | "concise" | "minimal");
59                if !valid {
60                    return Ok("Invalid style. Use: normal, concise, minimal.".to_string());
61                }
62                Ok(format!("{CMD_PREFIX}style:{value}"))
63            }
64            "steer_add" => {
65                let text = args["value"].as_str().unwrap_or("").trim();
66                if text.is_empty() {
67                    return Ok("Error: 'value' (the instruction text) is required.".to_string());
68                }
69                let turns = args["turns"].as_u64().map(|n| n as u32);
70                match turns {
71                    Some(n) => Ok(format!("{CMD_PREFIX}steer_add:{n}:{text}")),
72                    None => Ok(format!("{CMD_PREFIX}steer_add:permanent:{text}")),
73                }
74            }
75            "steer_remove" => {
76                let id = args["value"].as_str().unwrap_or("").trim();
77                if id.is_empty() {
78                    return Ok("Error: 'value' (the steer id or prefix) is required.".to_string());
79                }
80                Ok(format!("{CMD_PREFIX}steer_remove:{id}"))
81            }
82            "steer_list" => Ok(format!("{CMD_PREFIX}steer_list")),
83            "steer_clear" => Ok(format!("{CMD_PREFIX}steer_clear")),
84            "undo" => Ok(format!("{CMD_PREFIX}undo")),
85            "new_session" => Ok(format!("{CMD_PREFIX}new_session")),
86            _ => Ok(format!(
87                "Unknown action '{action}'. Use: style, steer_add, steer_remove, steer_list, steer_clear, undo, new_session."
88            )),
89        }
90    }
91}
92
93/// Parse a CMD: prefixed tool output and execute the corresponding agent mutation.
94/// Returns a human-readable result string, or None if the output is not a command.
95pub fn execute_agent_command(output: &str, agent: &mut dyn AgentControl) -> Option<String> {
96    let cmd = output.strip_prefix(CMD_PREFIX)?;
97
98    if let Some(style) = cmd.strip_prefix("style:") {
99        agent.set_style(style);
100        Some(format!("Output style set to '{style}'."))
101    } else if cmd.starts_with("steer_add:") {
102        let rest = cmd.strip_prefix("steer_add:").unwrap();
103        if let Some((dur, text)) = rest.split_once(':') {
104            let turns = if dur == "permanent" {
105                None
106            } else {
107                dur.parse::<u32>().ok()
108            };
109            let id = agent.steer_add(text, turns);
110            let dur_desc = match turns {
111                None => "permanent".to_string(),
112                Some(n) => format!("{n} turns"),
113            };
114            Some(format!(
115                "Steering instruction added [{:.8}] ({dur_desc}): {text}",
116                id
117            ))
118        } else {
119            Some("Failed to parse steer_add command.".to_string())
120        }
121    } else if let Some(id) = cmd.strip_prefix("steer_remove:") {
122        if agent.steer_remove(id) {
123            Some(format!("Steering instruction '{id}' removed."))
124        } else {
125            Some(format!("No steering instruction found with prefix '{id}'."))
126        }
127    } else if cmd == "steer_list" {
128        let list = agent.steer_list();
129        if list.is_empty() {
130            Some("No steering instructions active.".to_string())
131        } else {
132            Some(list)
133        }
134    } else if cmd == "steer_clear" {
135        agent.steer_clear();
136        Some("All steering instructions cleared.".to_string())
137    } else if cmd == "undo" {
138        if agent.undo_last_turn() {
139            Some("Last turn undone.".to_string())
140        } else {
141            Some("Nothing to undo.".to_string())
142        }
143    } else if cmd == "new_session" {
144        agent.new_session();
145        Some("New session started.".to_string())
146    } else {
147        None
148    }
149}
150
151/// Trait abstracting the Agent mutations that ControlTool needs.
152/// Implemented by Agent in the binary crate to avoid circular deps.
153pub trait AgentControl {
154    fn set_style(&mut self, style: &str);
155    fn steer_add(&mut self, text: &str, turns: Option<u32>) -> String;
156    fn steer_remove(&mut self, id: &str) -> bool;
157    fn steer_list(&self) -> String;
158    fn steer_clear(&mut self);
159    fn undo_last_turn(&mut self) -> bool;
160    fn new_session(&mut self);
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    struct MockAgent {
168        style: String,
169        steers: Vec<(String, String)>,
170    }
171
172    impl AgentControl for MockAgent {
173        fn set_style(&mut self, style: &str) {
174            self.style = style.to_string();
175        }
176        fn steer_add(&mut self, text: &str, _turns: Option<u32>) -> String {
177            let id = format!("mock{:04}", self.steers.len());
178            self.steers.push((id.clone(), text.to_string()));
179            id
180        }
181        fn steer_remove(&mut self, id: &str) -> bool {
182            let len = self.steers.len();
183            self.steers.retain(|(i, _)| !i.starts_with(id));
184            self.steers.len() < len
185        }
186        fn steer_list(&self) -> String {
187            self.steers
188                .iter()
189                .map(|(id, text)| format!("[{id}] {text}"))
190                .collect::<Vec<_>>()
191                .join("\n")
192        }
193        fn steer_clear(&mut self) {
194            self.steers.clear();
195        }
196        fn undo_last_turn(&mut self) -> bool {
197            true
198        }
199        fn new_session(&mut self) {}
200    }
201
202    #[test]
203    fn test_style_command() {
204        let mut agent = MockAgent {
205            style: "normal".into(),
206            steers: vec![],
207        };
208        let result = execute_agent_command("CMD:style:concise", &mut agent);
209        assert_eq!(result, Some("Output style set to 'concise'.".to_string()));
210        assert_eq!(agent.style, "concise");
211    }
212
213    #[test]
214    fn test_steer_add_permanent() {
215        let mut agent = MockAgent {
216            style: "normal".into(),
217            steers: vec![],
218        };
219        let result = execute_agent_command("CMD:steer_add:permanent:be brief", &mut agent);
220        assert!(result.unwrap().contains("be brief"));
221        assert_eq!(agent.steers.len(), 1);
222    }
223
224    #[test]
225    fn test_steer_clear() {
226        let mut agent = MockAgent {
227            style: "normal".into(),
228            steers: vec![("a".into(), "x".into())],
229        };
230        let result = execute_agent_command("CMD:steer_clear", &mut agent);
231        assert!(result.unwrap().contains("cleared"));
232        assert!(agent.steers.is_empty());
233    }
234
235    #[test]
236    fn test_non_command_returns_none() {
237        let mut agent = MockAgent {
238            style: "normal".into(),
239            steers: vec![],
240        };
241        assert!(execute_agent_command("just normal output", &mut agent).is_none());
242    }
243}
aegis_tools/control_tool.rs

aegis_tools/
control_tool.rs
