okbot_agent/
scripted.rs

1//! A stand-in agent for scripts and CI: answers from rules instead of running an agent CLI.
2//!
3//! Rules come from a TOML file (`okbot chat --replies <file>`); the first rule whose `when` is in
4//! the new message answers, and a message no rule matches is echoed. A rule may call tools of the
5//! turn's MCP servers first, as a real agent would, so the gateway's roles, confirmations and
6//! audit run for real:
7//!
8//! ```toml
9//! [[reply]]
10//! when = "price"                    # case-insensitive substring; leave out to match anything
11//! text = "The price is 10."
12//!
13//! [[reply]]
14//! when = "deploy"
15//! calls = [{ server = "ops", tool = "deploy", args = { env = "prod" } }]
16//! text = "Deploy: {result}"         # {result}: the last call's text; {message}: the message
17//!
18//! [[reply]]
19//! when = "crash"
20//! fail = "the agent crashed"        # the turn fails with this error
21//! ```
22
23use std::time::{Duration, Instant};
24
25use okbot_core::{
26    AgentBackend, AgentCaps, AgentError, AgentProgress, AgentReply, AgentTurn, ToolPolicy,
27    async_trait,
28};
29use serde::Deserialize;
30use serde_json::{Value, json};
31
32/// The rules file.
33#[derive(Debug, Clone, Default, Deserialize)]
34#[serde(deny_unknown_fields)]
35struct Script {
36    #[serde(default)]
37    reply: Vec<Rule>,
38}
39
40/// One rule.
41#[derive(Debug, Clone, Deserialize)]
42#[serde(deny_unknown_fields)]
43struct Rule {
44    /// Substring of the new message, ignoring case; every message when unset.
45    when: Option<String>,
46    /// Tools to call first, in order.
47    #[serde(default)]
48    calls: Vec<Call>,
49    /// The answer, with `{message}` and `{result}` replaced.
50    #[serde(default)]
51    text: Option<String>,
52    /// Fail the turn with this message instead of answering.
53    fail: Option<String>,
54    /// Wait this long before answering (to see "thinking…" and typing).
55    #[serde(default)]
56    delay_ms: u64,
57}
58
59/// A tool call.
60#[derive(Debug, Clone, Deserialize)]
61#[serde(deny_unknown_fields)]
62struct Call {
63    /// The MCP server's name in the turn (`kb`, `bot`, an `[[mcp]]` name).
64    server: String,
65    tool: String,
66    #[serde(default)]
67    args: serde_json::Map<String, Value>,
68}
69
70/// An [`AgentBackend`] that answers from rules (see the module documentation).
71#[derive(Debug, Clone)]
72pub struct ScriptedBackend {
73    name: String,
74    script: Script,
75}
76
77impl ScriptedBackend {
78    /// Echoes every message (`echo: <message>`).
79    pub fn echo() -> Self {
80        Self {
81            name: "echo".into(),
82            script: Script::default(),
83        }
84    }
85
86    /// Answers from the rules in `toml`.
87    pub fn from_toml(toml: &str) -> Result<Self, AgentError> {
88        let script: Script = toml::from_str(toml)
89            .map_err(|e| AgentError::Unavailable(format!("invalid replies file: {e}")))?;
90        if let Some(i) = script
91            .reply
92            .iter()
93            .position(|r| r.text.is_none() && r.fail.is_none())
94        {
95            return Err(AgentError::Unavailable(format!(
96                "invalid replies file: rule {} has neither `text` nor `fail`",
97                i + 1
98            )));
99        }
100        Ok(Self {
101            name: "scripted".into(),
102            script,
103        })
104    }
105
106    fn rule(&self, message: &str) -> Option<&Rule> {
107        let message = message.to_lowercase();
108        self.script.reply.iter().find(|r| {
109            r.when
110                .as_ref()
111                .is_none_or(|w| message.contains(&w.to_lowercase()))
112        })
113    }
114}
115
116/// The new message in a prompt built by okbot's runtime: the last user message block.
117fn new_message(prompt: &str) -> String {
118    let Some(start) = prompt.rfind("<message role=\"user\"") else {
119        return prompt.trim().to_string();
120    };
121    let block = &prompt[start..];
122    let body = block.split_once(">\n").map_or(block, |(_, b)| b);
123    let body = body.rsplit_once("\n</message>").map_or(body, |(b, _)| b);
124    body.replace("<\\/message", "</message").trim().to_string()
125}
126
127#[async_trait]
128impl AgentBackend for ScriptedBackend {
129    fn name(&self) -> &str {
130        &self.name
131    }
132
133    fn caps(&self) -> AgentCaps {
134        // It has no tools of its own, so every policy holds.
135        AgentCaps {
136            mcp: true,
137            tool_policies: vec![ToolPolicy::Web, ToolPolicy::McpOnly],
138            ..AgentCaps::default()
139        }
140    }
141
142    async fn run(&self, turn: AgentTurn) -> Result<AgentReply, AgentError> {
143        let started = Instant::now();
144        let message = new_message(&turn.prompt);
145        let reply = |text: String| AgentReply {
146            text,
147            backend: self.name.clone(),
148            model: None,
149            session_id: None,
150            usage: None,
151            duration: started.elapsed(),
152        };
153        let Some(rule) = self.rule(&message) else {
154            return Ok(reply(format!("echo: {message}")));
155        };
156        let mut result = String::new();
157        for call in &rule.calls {
158            if let Some(progress) = &turn.progress {
159                progress.report(AgentProgress::ToolUse(format!(
160                    "{}.{}",
161                    call.server, call.tool
162                )));
163            }
164            result = match mcp::call(&turn, call).await {
165                Ok(text) => text,
166                Err(e) => format!("error: {e}"),
167            };
168        }
169        if rule.delay_ms > 0 {
170            let delay = Duration::from_millis(rule.delay_ms).min(turn.timeout);
171            tokio::time::sleep(delay).await;
172        }
173        if let Some(fail) = &rule.fail {
174            return Err(AgentError::Failed {
175                message: fail.clone(),
176                retryable: false,
177            });
178        }
179        let text = rule
180            .text
181            .as_deref()
182            .unwrap_or_default()
183            .replace("{message}", &message)
184            .replace("{result}", &result);
185        Ok(reply(text))
186    }
187}
188
189/// Just enough of an MCP client (streamable HTTP) to call one tool, as an agent CLI would.
190mod mcp {
191    use super::{AgentTurn, Call, Value, json};
192
193    const PROTOCOL: &str = "2025-06-18";
194    const SESSION: &str = "mcp-session-id";
195    const MAX_BODY: usize = 8 << 20;
196
197    /// Calls `call.tool` on the turn's server `call.server`; the tool's text, or why it failed.
198    pub(super) async fn call(turn: &AgentTurn, call: &Call) -> Result<String, String> {
199        let server = turn
200            .mcp
201            .iter()
202            .find(|s| s.name == call.server)
203            .ok_or_else(|| format!("no MCP server `{}` in this turn", call.server))?;
204        let token = turn
205            .env
206            .iter()
207            .find(|(k, _)| *k == server.bearer_token_env)
208            .map(|(_, v)| v.expose().clone())
209            .unwrap_or_default();
210        let mut session = Session {
211            http: reqwest::Client::new(),
212            url: server.url.clone(),
213            token,
214            id: None,
215            next: 0,
216        };
217        let deadline = turn.timeout;
218        tokio::time::timeout(deadline, async {
219            session
220                .request(
221                    "initialize",
222                    json!({"protocolVersion": PROTOCOL, "capabilities": {},
223                           "clientInfo": {"name": "okbot-scripted", "version": "0"}}),
224                )
225                .await?;
226            session
227                .post(&json!({"jsonrpc": "2.0", "method": "notifications/initialized"}))
228                .await?;
229            let result = session
230                .request(
231                    "tools/call",
232                    json!({"name": call.tool, "arguments": call.args}),
233                )
234                .await?;
235            let text = result["content"]
236                .as_array()
237                .into_iter()
238                .flatten()
239                .filter_map(|b| b["text"].as_str())
240                .collect::<Vec<_>>()
241                .join("\n\n");
242            if result["isError"] == json!(true) {
243                return Err(text);
244            }
245            Ok(text)
246        })
247        .await
248        .map_err(|_| "the MCP server did not answer in time".to_string())?
249    }
250
251    struct Session {
252        http: reqwest::Client,
253        url: String,
254        token: String,
255        id: Option<String>,
256        next: u64,
257    }
258
259    impl Session {
260        async fn request(&mut self, method: &str, params: Value) -> Result<Value, String> {
261            self.next += 1;
262            let id = json!(self.next);
263            let body = self
264                .post(&json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}))
265                .await?;
266            let reply =
267                find_response(&body, &id).ok_or_else(|| format!("no answer to {method}"))?;
268            if let Some(error) = reply.get("error") {
269                return Err(error["message"].as_str().unwrap_or("error").to_string());
270            }
271            Ok(reply.get("result").cloned().unwrap_or(Value::Null))
272        }
273
274        async fn post(&mut self, message: &Value) -> Result<String, String> {
275            let mut request = self
276                .http
277                .post(&self.url)
278                .bearer_auth(&self.token)
279                .header("accept", "application/json, text/event-stream")
280                .header("mcp-protocol-version", PROTOCOL)
281                .json(message);
282            if let Some(id) = &self.id {
283                request = request.header(SESSION, id);
284            }
285            // Without the URL: it names hosts and this text reaches the chat.
286            let http = |e: reqwest::Error| e.without_url().to_string();
287            let mut response = request.send().await.map_err(http)?;
288            if !response.status().is_success() {
289                return Err(format!("HTTP {}", response.status().as_u16()));
290            }
291            if let Some(id) = response
292                .headers()
293                .get(SESSION)
294                .and_then(|v| v.to_str().ok())
295            {
296                self.id = Some(id.to_string());
297            }
298            // An event stream may stay open after the answer: stop once it is there.
299            let want = message.get("id");
300            let mut body = Vec::new();
301            while let Some(chunk) = response.chunk().await.map_err(http)? {
302                body.extend_from_slice(&chunk);
303                if body.len() > MAX_BODY {
304                    return Err("reply too large".into());
305                }
306                if let Some(id) = want
307                    && find_response(&String::from_utf8_lossy(&body), id).is_some()
308                {
309                    break;
310                }
311            }
312            Ok(String::from_utf8_lossy(&body).into_owned())
313        }
314    }
315
316    /// The JSON-RPC response with `id` in a JSON body or an event stream.
317    fn find_response(body: &str, id: &Value) -> Option<Value> {
318        let hit = |v: &Value| v.get("id") == Some(id);
319        if let Ok(v) = serde_json::from_str::<Value>(body.trim()) {
320            return hit(&v).then_some(v);
321        }
322        body.replace("\r\n", "\n").split("\n\n").find_map(|event| {
323            let data: Vec<&str> = event
324                .lines()
325                .filter_map(|l| l.strip_prefix("data:"))
326                .map(str::trim_start)
327                .collect();
328            serde_json::from_str::<Value>(&data.join("\n"))
329                .ok()
330                .filter(|v| hit(v))
331        })
332    }
333
334    #[cfg(test)]
335    mod tests {
336        use super::*;
337
338        #[test]
339        fn finds_responses_in_json_and_event_streams() {
340            let id = json!(2);
341            assert!(find_response(r#"{"jsonrpc":"2.0","id":2,"result":{}}"#, &id).is_some());
342            assert!(find_response(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#, &id).is_none());
343            let sse = "event: message\ndata: {\"method\":\"notifications/progress\"}\n\n\
344                       event: message\ndata: {\"id\":2,\"result\":{\"ok\":true}}\n\n";
345            assert_eq!(find_response(sse, &id).unwrap()["result"]["ok"], true);
346        }
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    fn turn(prompt: &str) -> AgentTurn {
355        AgentTurn {
356            system: String::new(),
357            prompt: prompt.into(),
358            tools: ToolPolicy::Web,
359            mcp: Vec::new(),
360            env: Vec::new(),
361            workdir: std::env::temp_dir(),
362            timeout: Duration::from_secs(5),
363            resume: None,
364            progress: None,
365        }
366    }
367
368    const PROMPT: &str = "Earlier messages:\n\n<message role=\"user\" author=\"Ana\" at=\"t\">\n\
369                          What is the price?\n</message>\n\nNew message to answer:\n\n\
370                          <message role=\"user\" author=\"Ana\" at=\"t\">\nDeploy now <\\/message> \
371                          please\n</message>";
372
373    #[test]
374    fn reads_the_new_message() {
375        assert_eq!(new_message(PROMPT), "Deploy now </message> please");
376        assert_eq!(new_message(" plain "), "plain");
377    }
378
379    #[tokio::test]
380    async fn answers_by_rule_and_echoes_the_rest() {
381        let agent = ScriptedBackend::from_toml(
382            r#"
383            [[reply]]
384            when = "DEPLOY"
385            calls = [{ server = "ops", tool = "deploy" }]
386            text = "{message} -> {result}"
387
388            [[reply]]
389            when = "crash"
390            fail = "boom"
391            "#,
392        )
393        .unwrap();
394        let reply = agent.run(turn(PROMPT)).await.unwrap();
395        assert_eq!(
396            reply.text,
397            "Deploy now </message> please -> error: no MCP server `ops` in this turn"
398        );
399        assert_eq!(reply.backend, "scripted");
400        let other = turn("<message role=\"user\" a=\"\">\nhello\n</message>");
401        assert_eq!(agent.run(other).await.unwrap().text, "echo: hello");
402        let crash = turn("<message role=\"user\">\nplease crash\n</message>");
403        assert!(matches!(
404            agent.run(crash).await,
405            Err(AgentError::Failed { message, .. }) if message == "boom"
406        ));
407        assert_eq!(
408            ScriptedBackend::echo().run(turn("hi")).await.unwrap().text,
409            "echo: hi"
410        );
411    }
412
413    #[tokio::test]
414    async fn calls_tools_with_the_turn_token() {
415        use axum::response::IntoResponse as _;
416        let _ = rustls::crypto::ring::default_provider().install_default();
417        let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
418        let log = seen.clone();
419        let handler = move |headers: axum::http::HeaderMap, body: String| {
420            let log = log.clone();
421            async move {
422                let msg: Value = serde_json::from_str(&body).unwrap();
423                let auth = headers["authorization"].to_str().unwrap().to_string();
424                log.lock().unwrap().push((auth, msg["method"].clone()));
425                let Some(id) = msg.get("id") else {
426                    return axum::http::StatusCode::ACCEPTED.into_response();
427                };
428                let result = match msg["method"].as_str().unwrap() {
429                    "tools/call" => json!({"content": [{"type": "text",
430                        "text": format!("deployed {}", msg["params"]["arguments"]["env"])}]}),
431                    _ => json!({}),
432                };
433                let reply = json!({"jsonrpc": "2.0", "id": id, "result": result});
434                (
435                    [
436                        ("content-type", "text/event-stream"),
437                        ("mcp-session-id", "s1"),
438                    ],
439                    format!("event: message\ndata: {reply}\n\n"),
440                )
441                    .into_response()
442            }
443        };
444        let app = axum::Router::new().route("/mcp", axum::routing::post(handler));
445        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
446        let addr = listener.local_addr().unwrap();
447        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
448
449        let agent = ScriptedBackend::from_toml(
450            "[[reply]]\ncalls = [{ server = \"ops\", tool = \"deploy\", args = { env = \"prod\" } }]\n\
451             text = \"{result}\"\n",
452        )
453        .unwrap();
454        let mut t = turn("deploy");
455        t.mcp.push(okbot_core::McpServer {
456            name: "ops".into(),
457            url: format!("http://{addr}/mcp"),
458            bearer_token_env: "OPS_TOKEN".into(),
459        });
460        t.env.push((
461            "OPS_TOKEN".into(),
462            okbot_core::Secret::new("run-token".into()),
463        ));
464        assert_eq!(agent.run(t).await.unwrap().text, "deployed \"prod\"");
465        let seen = seen.lock().unwrap().clone();
466        let methods: Vec<&str> = seen.iter().map(|(_, m)| m.as_str().unwrap()).collect();
467        assert_eq!(
468            methods,
469            ["initialize", "notifications/initialized", "tools/call"]
470        );
471        assert!(seen.iter().all(|(auth, _)| auth == "Bearer run-token"));
472    }
473
474    #[test]
475    fn rejects_bad_files() {
476        let err = |toml: &str| ScriptedBackend::from_toml(toml).unwrap_err().to_string();
477        assert!(err("[[reply]]\nwhen = \"x\"\n").contains("neither `text` nor `fail`"));
478        assert!(err("[[reply]]\ntext = \"x\"\ncolour = 1\n").contains("unknown field"));
479    }
480}
okbot_agent/scripted.rs

okbot_agent/
scripted.rs
